diff --git a/app/controllers/data.py b/app/controllers/data.py index 5b9c6cc2f89a817244ada3e931e18a057082ff8a..e4404da9f426169a74f15434749aacf458240666 100644 --- a/app/controllers/data.py +++ b/app/controllers/data.py @@ -30,12 +30,9 @@ class DataController(RestController): where_statement = where_statement[:-4] + ')' sql = "SELECT ts, value, unit, unit_id, name, sn, mod FROM data {} ORDER BY ts".format(where_statement) - cur = self.redshift.db.cursor(cursor_factory=psycopg2.extras.DictCursor) - try: - cur.execute(sql, where_tuple) - except Exception as e: - cur.close() - self.redshift.db.rollback() - raise e + cur = self.redshift.get_cursor() + cur = self.redshift.execute(cur, sql, where_tuple) + results = [dict(record) for record in cur] + cur.close() - return [dict(record) for record in cur] + return results diff --git a/app/lib/redshift_database.py b/app/lib/redshift_database.py index d5409f240ecfd8c0c3cc0576c91c95fb17880a4d..f78c9a578487c89bdd88aa491bca2e064bc13e56 100644 --- a/app/lib/redshift_database.py +++ b/app/lib/redshift_database.py @@ -5,19 +5,55 @@ 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.""" - uri = urlparse(app.config.get('REDSHIFT_DATABASE_URI')) + self.uri = app.config.get('REDSHIFT_DATABASE_URI') - self.db = psycopg2.connect( - user=uri.username, - password=uri.password, - host=uri.hostname, + 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.path[1:] + 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 + """ + try: + cursor.execute(sql, where_tuple) + except psycopg2.ProgrammingError as e: + self.current_app.logger.info( + 'ProgrammingError while executing redshift db call: {}'.format(e.message) + ) + self.db.rollback() + except psycopg2.InterfaceError as e: + self.current_app.logger.info( + 'InterfaceError while executing redshift db call: {}'.format(e.message) + ) + self.db = self.make_connection() + return self.get_cursor() + return cursor + + redshift = RedshiftWrapper() def register(app):