Cleanup: replace db.cursor() with get_db_cursor()

Author: grafolean

dbutils.py

@@ -44,13 +44,13 @@ def get_db_connection():
 @contextmanager
 def get_db_cursor(commit=False):
     with get_db_connection() as connection:
-      cursor = connection.cursor()
-      try:
-          yield cursor
-          if commit:
-              connection.commit()
-      finally:
-          cursor.close()
+        cursor = connection.cursor()
+        try:
+            yield cursor
+            if commit:
+                connection.commit()
+        finally:
+            cursor.close()
 
 
 def db_connect():
@@ -75,21 +75,14 @@ def db_connect():
         db_pool = None
         log.error("DB connection failed")
 
-# This class is only needed until we replace all db.cursor() calls with get_db_cursor()
-class ThinDBWrapper(object):
-    @staticmethod
-    def cursor():
-        return get_db_cursor()
-db = ThinDBWrapper
-
 
 ###########################
 #   DB schema migration   #
 ###########################
 
 def get_existing_schema_version():
     existing_schema_version = 0
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         try:
             c.execute(f'SELECT schema_version FROM {DB_PREFIX}runtime_data;')
             res = c.fetchone()
@@ -120,7 +113,7 @@ def migrate_if_needed():
         method_to_call = getattr(sys.modules[__name__], method_name)
         method_to_call()
         # automatically upgrade schema version if there is no exception:
-        with db.cursor() as c:
+        with get_db_cursor() as c:
             c.execute(f'UPDATE {DB_PREFIX}runtime_data SET schema_version = %s;', (try_migrating_to,))
         try_migrating_to += 1
     if try_migrating_to == existing_schema_version + 1:
@@ -130,12 +123,12 @@ def migrate_if_needed():
 
 
 def migration_step_1():
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         c.execute(f'CREATE TABLE {DB_PREFIX}runtime_data (schema_version SMALLSERIAL NOT NULL);')
         c.execute(f'INSERT INTO {DB_PREFIX}runtime_data (schema_version) VALUES (1);')
 
 def migration_step_2():
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         # UNLOGGED: Disabling WAL avoids high I/O load. Since NetFlow data is of temporary nature, this still
         # allows us to perform queries, but if the database crashes it is acceptable to lose all of the records.
         c.execute(f'CREATE UNLOGGED TABLE {DB_PREFIX}records (seq BIGSERIAL NOT NULL PRIMARY KEY, ts NUMERIC(16,6) NOT NULL, client_ip TEXT);')

netflowbot.py

@@ -13,7 +13,7 @@
 import requests
 
 from grafoleancollector import Collector, send_results_to_grafolean
-from dbutils import db, DB_PREFIX
+from dbutils import get_db_cursor, DB_PREFIX
 from lookup import PROTOCOLS, DIRECTION_INGRESS, DIRECTION_EGRESS
 
 logging.basicConfig(format='%(asctime)s.%(msecs)03d | %(levelname)s | %(message)s',
@@ -26,7 +26,7 @@
 
 
 def _get_last_used_seq(job_id):
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         c.execute(f'SELECT j.last_used_seq, r.ts FROM {DB_PREFIX}bot_jobs j, {DB_PREFIX}records r WHERE j.job_id = %s AND j.last_used_seq = r.seq;', (job_id,))
         rec = c.fetchone()
         if rec is None:
@@ -35,7 +35,7 @@ def _get_last_used_seq(job_id):
         return last_used_seq, ts
 
 def _get_current_max_seq():
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         c.execute(f"SELECT seq, ts FROM {DB_PREFIX}records WHERE seq = (SELECT MAX(seq) FROM {DB_PREFIX}records);")
         rec = c.fetchone()
         if rec is None:
@@ -44,7 +44,7 @@ def _get_current_max_seq():
         return max_seq, now_ts
 
 def _save_current_max_seq(job_id, seq):
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         c.execute(f"INSERT INTO {DB_PREFIX}bot_jobs (job_id, last_used_seq) VALUES (%s, %s) ON CONFLICT (job_id) DO UPDATE SET last_used_seq = %s;", (job_id, seq, seq))
 
 
@@ -178,7 +178,7 @@ def construct_output_path_prefix(interval_label, direction, entity_id, interface
     @staticmethod
     def get_traffic_for_entity(interval_label, last_seq, max_seq, time_between, direction, entity_id, entity_ip):
         # returns cumulative traffic for the whole entity, and traffic per interface for this entity
-        with db.cursor() as c:
+        with get_db_cursor() as c:
 
             c.execute(f"""
                 SELECT
@@ -218,7 +218,7 @@ def get_traffic_for_entity(interval_label, last_seq, max_seq, time_between, dire
     # @staticmethod
     # def get_traffic_all_entities(interval_label, last_seq, max_seq, direction):
     #     output_path = NetFlowBot.construct_output_path_prefix(interval_label, direction, entity_id=None, interface=None)
-    #     with db.cursor() as c:
+    #     with get_db_cursor() as c:
     #         c.execute(f"""
     #             SELECT
     #                 sum(f.in_bytes)
@@ -242,7 +242,7 @@ def get_traffic_for_entity(interval_label, last_seq, max_seq, time_between, dire
 
     # @staticmethod
     # def get_top_N_IPs(output_path_prefix, from_time, to_time, interface_index, is_direction_in=True):
-    #     with db.cursor() as c:
+    #     with get_db_cursor() as c:
     #         # TODO: missing check for IP: r.client_ip = %s AND
     #         c.execute(f"""
     #             SELECT
@@ -275,7 +275,7 @@ def get_traffic_for_entity(interval_label, last_seq, max_seq, time_between, dire
 
     # @staticmethod
     # def get_top_N_protocols(output_path_prefix, from_time, to_time, interface_index, is_direction_in=True):
-    #     with db.cursor() as c:
+    #     with get_db_cursor() as c:
     #         # TODO: missing check for IP: r.client_ip = %s AND
     #         c.execute(f"""
     #             SELECT

netflowwriter.py

@@ -13,7 +13,7 @@
 from colors import color
 
 from lookup import PROTOCOLS
-from dbutils import migrate_if_needed, db, DB_PREFIX
+from dbutils import migrate_if_needed, get_db_cursor, DB_PREFIX
 
 
 logging.basicConfig(format='%(asctime)s.%(msecs)03d | %(levelname)s | %(message)s',
@@ -69,7 +69,7 @@ def write_record(j):
     # }
     # https://www.cisco.com/en/US/technologies/tk648/tk362/technologies_white_paper09186a00800a3db9.html#wp9001622
 
-    with db.cursor() as c:
+    with get_db_cursor() as c:
         # first save the flow record:
         ts = j['ts']
         log.info(f"Received record [{j['seq']}]: {datetime.utcfromtimestamp(ts)} from {j['client']}")