From 556b654c6ae99570b26f0b4a32481109b829ac91 Mon Sep 17 00:00:00 2001
From: Conrad <conrad.schloer@gmail.com>
Date: Mon, 18 Dec 2017 11:22:15 -0500
Subject: [PATCH] Wrap the cursor and execute functions

---
 app/controllers/data.py      | 13 ++++------
 app/lib/redshift_database.py | 48 +++++++++++++++++++++++++++++++-----
 2 files changed, 47 insertions(+), 14 deletions(-)

diff --git a/app/controllers/data.py b/app/controllers/data.py
index 5b9c6cc..e4404da 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 d5409f2..f78c9a5 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):
-- 
GitLab