diff --git a/app/__init__.py b/app/__init__.py index 1e1efe31c1900e5559409fe37fcb31cd0f5adada..d7684603a91825ec99673a3c5ef6a1977825745e 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -32,8 +32,8 @@ def create_app(config): from . import views views.register(app) - from app.lib import database, red, exceptions, service, session, auth0, redshift_database - services = (database, red, exceptions, service, session, auth0, redshift_database) + from app.lib import database, red, exceptions, service, session, auth0 + services = (database, red, exceptions, service, session, auth0) for service in services: service.register(app) diff --git a/app/controllers/base.py b/app/controllers/base.py index 2e60fb7a7a78c713ad6c4d02ed70b605b8e8885e..cf86692dd25dcda3e91e3b7f0defdc7a89b4c12b 100644 --- a/app/controllers/base.py +++ b/app/controllers/base.py @@ -1,6 +1,6 @@ """Flask RestController wrapper""" from bpvalve.flask.controllers import RestController as BaseRestController -from app.lib import db, redis, auth0_, redshift +from app.lib import db, redis, auth0_ class RestController(BaseRestController): @@ -10,4 +10,3 @@ class RestController(BaseRestController): db = db redis = redis auth0 = auth0_ - redshift = redshift diff --git a/app/controllers/data.py b/app/controllers/data.py index 253e18fa287f56ae8f47e7cb7931d4c01d41a952..22f695b812a12df0df0de9d29d3f50cc740ed53e 100644 --- a/app/controllers/data.py +++ b/app/controllers/data.py @@ -1,9 +1,12 @@ -from .base import RestController +from flask import current_app import datetime import json import psycopg2 from werkzeug.exceptions import BadRequest +from app.lib import get_redshift_db +from .base import RestController + class DataController(RestController): """A data controller.""" @@ -35,8 +38,9 @@ class DataController(RestController): where_tuple = where_tuple + (filter_data['unit_id'],) sql = "SELECT ts, value, unit, unit_id, name, sn, mod FROM data {} ORDER BY ts".format(where_statement) - cur = self.redshift.get_cursor() - cur = self.redshift.execute(cur, sql, where_tuple) + db = get_redshift_db(current_app) + cur = db.cursor(cursor_factory=psycopg2.extras.DictCursor) + cur.execute(sql, where_tuple) results = [dict(record) for record in cur] cur.close() diff --git a/app/lib/__init__.py b/app/lib/__init__.py index 90e9f86aab7bab7275fc8aaff0ec80e319422608..9e859d01100e454b76f89aa28f1a05b5489d8f42 100644 --- a/app/lib/__init__.py +++ b/app/lib/__init__.py @@ -1,4 +1,4 @@ from .database import db from .red import redis from .auth0 import auth0_ -from .redshift_database import redshift +from .redshift_database import get_redshift_db diff --git a/app/lib/redshift_database.py b/app/lib/redshift_database.py index 5ad1b78109d285a4b174650c0fc5841f67052581..282d64127d4450a0b06035fca85fdb5078d58060 100644 --- a/app/lib/redshift_database.py +++ b/app/lib/redshift_database.py @@ -1,64 +1,28 @@ +from flask import Flask, g from urllib.parse import urlparse import psycopg2 - -class RedshiftWrapper(object): - """A wrapper for Redshift.""" - db = None - current_app = None - uri = '' - - def init_app(self, app): - """Stores information about the redis database from the URI.""" - self.uri = app.config.get('REDSHIFT_DATABASE_URI') - - self.db = self.make_connection() - self.current_app = app - - def make_connection(self): - uri_data = urlparse(self.uri) - return psycopg2.connect( - user=uri_data.username, - password=uri_data.password, - host=uri_data.hostname, - port=5439, - dbname=uri_data.path[1:] - ) - - def get_cursor(self): - try: - cursor = self.db.cursor(cursor_factory=psycopg2.extras.DictCursor) - except: - self.db.close() - self.db = self.make_connection() - cursor = self.db.cursor(cursor_factory=psycopg2.extras.DictCursor) - return cursor - - def execute(self, cursor, sql, where_tuple): - """ Execute a database call and handle the InterfaceError - Returns the cursor that should be used to make future calls - """ - retry_counter = 0 - while retry_counter < 5: - retry_counter += 1 - try: - cursor.execute(sql, where_tuple) - break - except psycopg2.ProgrammingError as e: - self.current_app.logger.info( - 'ProgrammingError while executing redshift db call: {}'.format(e) - ) - self.db.rollback() - except (psycopg2.InterfaceError, psycopg2.OperationalError) as e: - self.current_app.logger.info( - 'Interface or Operational Error while executing redshift db call: {}'.format(e) - ) - self.db = self.make_connection() - cursor = self.get_cursor() - return cursor - - -redshift = RedshiftWrapper() - -def register(app): - redshift.init_app(app) +app = Flask(__name__) + +def get_redshift_db(app): + db = getattr(g, 'redshift_database', None) + if db is None: + db = g._database = connect_to_database(app) + return db + +@app.teardown_appcontext +def teardown_db(exception): + db = getattr(g, 'redshift_database', None) + if db is not None: + db.close() + +def connect_to_database(app): + uri = app.config.get('REDSHIFT_DATABASE_URI') + uri_data = urlparse(uri) + return psycopg2.connect( + user=uri_data.username, + password=uri_data.password, + host=uri_data.hostname, + port=5439, + dbname=uri_data.path[1:] + )