https://github.com/WikiApiary/WikiApiary/issues/42

Paul Decoursey · Paul Decoursey · commit 86c2ff6db586 · 2014-08-05T19:47:22.000-05:00
diff --git a/apiary/ApiaryAPI/api.py b/apiary/ApiaryAPI/api.py
@@ -16,7 +16,7 @@ def data(site_id):
 
     duration = request.args.get('duration')
 
-    cur = apiary_db.cursor()
+    cur = apiary_db.apiary_db.cursor()
     # SELECT capture_date, articles, pages FROM statistics WHERE website_id = :id AND capture_date > :date_filter
     temp_sql = "SELECT capture_date, articles, pages FROM statistics WHERE website_id = %d" % site_id
     cur.execute(temp_sql)
diff --git a/apiary/connect_mysql.py b/apiary/connect_mysql.py
@@ -30,11 +30,62 @@
 else:
     LOGGER.info("No configuration file detected.")
 
-LOGGER.info("Opening MySQL connection")
-apiary_db = mdb.connect(
-    host=APIARYDB_HOSTNAME,
-    db=APIARYDB_DATABASE,
-    user=APIARYDB_USERNAME,
-    passwd=APIARYDB_PASSWORD,
-    charset='utf8')
+
+
+class ApiaryDB(object):
+
+    def __init__(self, host, database, username, password):
+        self.host = host
+        self.database = database
+        self.username = username
+        self.password = password
+        self.reconnect()
+
+    def reconnect(self):
+        LOGGER.info("Opening MySQL connection")
+        self.apiary_db = mdb.connect(
+            host=self.host,
+            db=self.database,
+            user=self.username,
+            passwd=self.password,
+            charset='utf8')
+
+    def fetch_one(self, sql, retry_count=0):
+        try:
+            cur = self.apiary_db.cursor()
+            cur.execute(sql)
+            return cur.fetchone()
+        except MySQLdb.OperationalError:
+            if retry_count > 10: # seemed like a good count
+                raise Exception("max retries...")
+            self.reconnect()
+            retry_count += 1
+            return self.fetch_one(sql, retry_count)
+        except Exception, e:
+            cur.close()
+            LOGGER.error("SQL Command: %s" % sql)
+            raise Exception(e)
+
+
+    def runSql(self, sql_command, retry_count=0):
+        """Helper to run a SQL command and catch errors"""
+        LOGGER.debug("SQL: %s" % sql_command)
+        try:
+            cur = self.apiary_db.cursor()
+            cur.execute(sql_command)
+            cur.close()
+            self.apiary_db.commit()
+            return True, cur.rowcount
+        except MySQLdb.OperationalError:
+            if retry_count > 10: # seemed like a good count
+                raise Exception("max retries...")
+            self.reconnect()
+            retry_count += 1
+            return self.runSql(sql_command, retry_count)
+        except Exception, e:
+            cur.close()
+            LOGGER.error("SQL Command: %s" % sql_command)
+            raise Exception(e)
+
+apiary_db = ApiaryDB(APIARYDB_HOSTNAME, APIARYDB_DATABASE, APIARYDB_USERNAME, APIARYDB_PASSWORD)
 
diff --git a/apiary/tasks/__init__.py b/apiary/tasks/__init__.py
@@ -78,11 +78,7 @@ def runSql(self, sql_command):
         """Helper to run a SQL command and catch errors"""
         LOGGER.debug("SQL: %s" % sql_command)
         try:
-            cur = self.apiary_db.cursor()
-            cur.execute(sql_command)
-            cur.close()
-            self.apiary_db.commit()
-            return True, cur.rowcount
+            return self.apiary_db.runSql(sql_command)
         except Exception, e:
             cur.close()
             LOGGER.error("SQL Command: %s" % sql_command)
diff --git a/apiary/tasks/bot/updatetotaledits.py b/apiary/tasks/bot/updatetotaledits.py
@@ -33,9 +33,7 @@ def run(self):
 """
 
         # Get the total edit count
-        cur = self.apiary_db.cursor()
-        cur.execute(sql_query)
-        data = cur.fetchone()
+        data = self.apiary_db.fetch_one(sql_query)
         LOGGER.info("Total edits: %d Total active users: %d Total pages: %d", data[0], data[1], data[2])
 
         # Update the wiki with the new value
