simple retriable execute (for sql)

the idea is that db calls occasionally break due to transient issues and we’d rather have the calls retried before raising exceptions etc

2 files changed, 29 insertions, 1 deletions

diff --git a/python/skytools/scripting.py b/python/skytools/scripting.py
index 840f3cf4..0d53da71 100644
--- a/python/skytools/scripting.py
+++ b/python/skytools/scripting.py
@@ -921,6 +921,34 @@ class DBScript(BaseScript):
             # error is already logged
             sys.exit(1)
 
+    def execute_with_retry (self, cursor, stmt, args, exceptions = None):
+        """ Execute SQL and retry if it fails.
+        """
+        self.sql_retry_max_count = self.cf.getint("sql_retry_max_count", 10)
+        self.sql_retry_max_time = self.cf.getint("sql_retry_max_time", 300)
+        self.sql_retry_formula_a = self.cf.getint("sql_retry_formula_a", 1)
+        self.sql_retry_formula_b = self.cf.getint("sql_retry_formula_b", 5)
+        self.sql_retry_formula_cap = self.cf.getint("sql_retry_formula_cap", 60)
+        import psycopg2
+        elist = exceptions or (psycopg2.OperationalError,)
+        stime = time.time()
+        tried = 0
+        while True:
+            try:
+                cursor.execute (stmt, args)
+                break
+            except elist:
+                if tried >= self.sql_retry_max_count or time.time() - stime >= self.sql_retry_max_time:
+                    raise
+            except:
+                raise
+            # y = a + bx , apply cap
+            y = self.sql_retry_formula_a + self.sql_retry_formula_b * tried
+            if self.sql_retry_formula_cap is not None and y > self.sql_retry_formula_cap:
+                y = self.sql_retry_formula_cap
+            self.sleep(y)
+            tried += 1
+
     def listen(self, dbname, channel):
         """Make connection listen for specific event channel.
 
diff --git a/scripts/simple_local_consumer.py b/scripts/simple_local_consumer.py
index 87c65868..f6951852 100755
--- a/scripts/simple_local_consumer.py
+++ b/scripts/simple_local_consumer.py
@@ -59,7 +59,7 @@ class SimpleLocalConsumer(pgq.LocalConsumer):
         payload['pgq.ev_extra4'] = ev.ev_extra4
 
         self.log.debug(self.dst_query, payload)
-        curs.execute(self.dst_query, payload)
+        self.execute_with_retry(curs, self.dst_query, payload)
         if curs.statusmessage[:6] == 'SELECT':
             res = curs.fetchall()
             self.log.debug(res)