Nicer DB connection loss handling (InvalidDBCursor)

Author: Anze

dbutils.py

@@ -49,25 +49,36 @@ def get_db_connection():
         conn.autocommit = True
         conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
         yield conn
-    except DBConnectionError:
+    except (DBConnectionError, psycopg2.OperationalError):
+        db_pool = None  # make sure that we reconnect next time
         yield None
     finally:
-        if db_pool is not None and conn is not None:
+        if db_pool is not None:
             db_pool.putconn(conn)
 
 
 @contextmanager
 def get_db_cursor():
     with get_db_connection() as connection:
         if connection is None:
-            yield None
+            yield InvalidDBCursor()
             return
 
         cursor = connection.cursor()
         try:
             yield cursor
         finally:
-            cursor.close()
+            if not isinstance(cursor, InvalidDBCursor):
+                cursor.close()
+
+
+# In python it is not possible to throw an exception within the __enter__ phase of a with statement:
+#   https://www.python.org/dev/peps/pep-0377/
+# If we want to handle DB connection failures gracefully we return a cursor which will throw
+# DBConnectionError exception whenever it is accessed.
+class InvalidDBCursor(object):
+    def __getattr__(self, attr):
+        raise DBConnectionError()
 
 
 def db_connect():
@@ -105,10 +116,13 @@ def db_disconnect():
 def initial_wait_for_db():
     while True:
         with get_db_cursor() as c:
-            if c is not None:
+            try:
+                c.execute('SELECT 1;')
+                res = c.fetchone()
                 return
-            log.info("DB connection failed - waiting for DB to become available, sleeping 5s")
-            time.sleep(5)
+            except DBConnectionError:
+                log.info("DB connection failed - waiting for DB to become available, sleeping 5s")
+                time.sleep(5)
 
 
 ###########################

snmpbot.py

@@ -40,8 +40,6 @@ class InvalidOutputPath(Exception):
 
 def _get_previous_counter_value(counter_ident):
     with get_db_cursor() as c:
-        if c is None:
-            raise DBConnectionError()
         try:
             c.execute(f'SELECT value, ts FROM {DB_PREFIX}bot_counters WHERE id = %s;', (counter_ident,))
             rec = c.fetchone()
@@ -56,8 +54,6 @@ def _get_previous_counter_value(counter_ident):
 
 def _save_current_counter_value(new_value, now, counter_ident):
     with get_db_cursor() as c:
-        if c is None:
-            raise DBConnectionError()
         c.execute(f"INSERT INTO {DB_PREFIX}bot_counters (id, value, ts) VALUES (%s, %s, %s) ON CONFLICT (id) DO UPDATE SET value = %s, ts = %s;",
                 (counter_ident, new_value, now, new_value, now))