Fix: close DB connection in main process to avoid multithreading issues caused by creating DB connection pool before the workers were forked

Author: Anze

dbutils.py

@@ -74,7 +74,16 @@ def db_connect():
                               connect_timeout=connect_timeout)
     except:
         db_pool = None
-        log.error("DB connection failed")
+        log.exception("DB connection failed")
+
+
+def db_disconnect():
+    global db_pool
+    if not db_pool:
+        return
+    db_pool.closeall()
+    db_pool = None
+    log.info("DB connection is closed")
 
 
 ###########################

snmpbot.py

@@ -14,7 +14,7 @@
 import psycopg2
 
 from grafoleancollector import Collector
-from dbutils import get_db_cursor, DB_PREFIX, migrate_if_needed
+from dbutils import get_db_cursor, DB_PREFIX, migrate_if_needed, db_disconnect
 
 
 logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
@@ -410,7 +410,6 @@ def jobs(self):
             Each entity (device) is a single job, no matter how many sensors it has. The reason is
             that when the intervals align, we can then issue a single SNMP Bulk GET/WALK.
         """
-        migrate_if_needed()
         for entity_info in self.fetch_job_configs('snmp'):
             intervals = list(set([sensor_info["interval"] for sensor_info in entity_info["sensors"]]))
             job_info = { **entity_info, "backend_url": self.backend_url, "bot_token": self.bot_token }
@@ -444,6 +443,9 @@ def wait_for_grafolean(backend_url):
 if __name__ == "__main__":
     dotenv.load_dotenv()
 
+    migrate_if_needed()
+    db_disconnect()  # each worker should open their own connection pool
+
     backend_url = os.environ.get('BACKEND_URL')
     jobs_refresh_interval = int(os.environ.get('JOBS_REFRESH_INTERVAL', 120))