Gracefully handle DB connection loss and startup time connectivity problems

Author: Anze

dbutils.py

@@ -5,6 +5,7 @@
 import sys
 import copy
 import json
+import time
 
 import psycopg2
 from psycopg2.pool import ThreadedConnectionPool
@@ -21,6 +22,10 @@
 log = logging.getLogger("{}.{}".format(__name__, "dbutils"))
 
 
+class DBConnectionError(Exception):
+    pass
+
+
 db_pool = None
 DB_PREFIX = 'snmp_'
 register_adapter(dict, Json)
@@ -32,24 +37,35 @@ def get_db_connection():
     global db_pool
     if db_pool is None:
         db_connect()
-
     try:
+        if db_pool is None:
+            # connecting to DB failed
+            raise DBConnectionError()
         conn = db_pool.getconn()
+        if conn is None:
+            # pool wasn't able to return a valid connection
+            raise DBConnectionError()
+
         conn.autocommit = True
         conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
         yield conn
+    except DBConnectionError:
+        yield None
     finally:
-        db_pool.putconn(conn)
+        if db_pool is not None and conn is not None:
+            db_pool.putconn(conn)
 
 
 @contextmanager
-def get_db_cursor(commit=False):
+def get_db_cursor():
     with get_db_connection() as connection:
+        if connection is None:
+            yield None
+            return
+
         cursor = connection.cursor()
         try:
             yield cursor
-            if commit:
-                connection.commit()
         finally:
             cursor.close()
 
@@ -74,7 +90,7 @@ def db_connect():
                               connect_timeout=connect_timeout)
     except:
         db_pool = None
-        log.exception("DB connection failed")
+        log.warning("DB connection failed")
 
 
 def db_disconnect():
@@ -86,6 +102,15 @@ def db_disconnect():
     log.info("DB connection is closed")
 
 
+def initial_wait_for_db():
+    while True:
+        with get_db_cursor() as c:
+            if c is not None:
+                return
+            log.info("DB connection failed - waiting for DB to become available, sleeping 5s")
+            time.sleep(5)
+
+
 ###########################
 #   DB schema migration   #
 ###########################

snmpbot.py

@@ -14,7 +14,7 @@
 import psycopg2
 
 from grafoleancollector import Collector
-from dbutils import get_db_cursor, DB_PREFIX, migrate_if_needed, db_disconnect
+from dbutils import get_db_cursor, DB_PREFIX, initial_wait_for_db, migrate_if_needed, db_disconnect, DBConnectionError
 
 
 logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
@@ -40,6 +40,8 @@ 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()
@@ -54,6 +56,8 @@ 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))
 
@@ -67,24 +71,28 @@ def _convert_counters_to_values(results, now, counter_ident_prefix):
         if v.snmp_type not in ['COUNTER', 'COUNTER64']:
             new_results.append(v)
             continue
+
         # counter - deal with it:
-        counter_ident = counter_ident_prefix + f'/{i}/{v.oid}/{v.oid_index}'
-        old_value, t = _get_previous_counter_value(counter_ident)
         new_value = int(float(v.value))
-        _save_current_counter_value(new_value, now, counter_ident)
-        if old_value is None:
-            new_results.append(SNMPVariable(oid=v.oid, oid_index=v.oid_index, value=None, snmp_type='COUNTER_PER_S'))
-            continue
+        counter_ident = counter_ident_prefix + f'/{i}/{v.oid}/{v.oid_index}'
+        try:
+            old_value, t = _get_previous_counter_value(counter_ident)
+            _save_current_counter_value(new_value, now, counter_ident)
+            if old_value is None:
+                new_results.append(SNMPVariable(oid=v.oid, oid_index=v.oid_index, value=None, snmp_type='COUNTER_PER_S'))
+                continue
 
-        # it seems like the counter overflow happened, discard result:
-        if new_value < old_value:
-            new_results.append(SNMPVariable(oid=v.oid, oid_index=v.oid_index, value=None, snmp_type='COUNTER_PER_S'))
-            log.warning(f"Counter overflow detected for oid {v.oid}, oid index {v.oid_index}, discarding value - if this happens often, consider using OIDS with 64bit counters (if available) or decreasing polling interval.")
-            continue
+            # it seems like the counter overflow happened, discard result:
+            if new_value < old_value:
+                new_results.append(SNMPVariable(oid=v.oid, oid_index=v.oid_index, value=None, snmp_type='COUNTER_PER_S'))
+                log.warning(f"Counter overflow detected for oid {v.oid}, oid index {v.oid_index}, discarding value - if this happens often, consider using OIDS with 64bit counters (if available) or decreasing polling interval.")
+                continue
 
-        dt = now - t
-        dv = (new_value - old_value) / dt
-        new_results.append(SNMPVariable(oid=v.oid, oid_index=v.oid_index, value=dv, snmp_type='COUNTER_PER_S'))
+            dt = now - t
+            dv = (new_value - old_value) / dt
+            new_results.append(SNMPVariable(oid=v.oid, oid_index=v.oid_index, value=dv, snmp_type='COUNTER_PER_S'))
+        except DBConnectionError:
+            log.error(f"Could not convert counter due to DB error: {counter_ident} / {new_value}")
     return new_results
 
 
@@ -443,6 +451,7 @@ def wait_for_grafolean(backend_url):
 if __name__ == "__main__":
     dotenv.load_dotenv()
 
+    initial_wait_for_db()
     migrate_if_needed()
     db_disconnect()  # each worker should open their own connection pool