Add DB utility functions fetchone fetchall execute

plugin/teksi_wastewater/interlis/interlis_importer_exporter.py

@@ -263,11 +263,8 @@ def _import_from_intermediate_schema(self, import_model):
         return interlisImporterToIntermediateSchema.session_tww
 
     def _import_manage_organisations(self):
-        with DatabaseUtils.PsycopgConnection() as connection:
-            cursor = connection.cursor()
-
-            logger.info("Update organisation tww_active")
-            cursor.execute("SELECT tww_app.set_organisations_active();")
+        logger.info("Update organisation tww_active")
+        DatabaseUtils.execute("SELECT tww_app.set_organisations_active();")
 
     def _import_update_main_cover_and_refresh_mat_views(self):
         with DatabaseUtils.PsycopgConnection() as connection:

plugin/teksi_wastewater/utils/database_utils.py

@@ -48,6 +48,36 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             self.connection.commit()
             self.connection.close()
 
+    @staticmethod
+    def fetchone(query: str):
+        with DatabaseUtils.PsycopgConnection() as connection:
+            cursor = connection.cursor()
+
+            cursor.execute(query)
+            return cursor.fetchone()
+
+    @staticmethod
+    def fetchall(query: str):
+        with DatabaseUtils.PsycopgConnection() as connection:
+            cursor = connection.cursor()
+
+            cursor.execute(query)
+            return cursor.fetchall()
+
+    @staticmethod
+    def execute(query: str):
+        with DatabaseUtils.PsycopgConnection() as connection:
+            cursor = connection.cursor()
+            cursor.execute(query)
+
+    @staticmethod
+    def get_psycopg_connection():
+        connection = psycopg.connect(
+            DatabaseUtils.get_pgconf_as_psycopg_dsn(), **DEFAULTS_CONN_ARG
+        )
+        if PSYCOPG_VERSION == 2:
+            connection.set_session(autocommit=True)
+
     @staticmethod
     def read_pgservice(service_name):
         """
@@ -116,30 +146,19 @@ def get_pgconf_as_psycopg_dsn() -> List[str]:
 
     @staticmethod
     def disable_symbology_triggers():
-        with DatabaseUtils.PsycopgConnection() as connection:
-            cursor = connection.cursor()
-
-            logger.info("Disable symbology triggers")
-            cursor.execute("SELECT tww_sys.disable_symbology_triggers();")
+        logger.info("Disable symbology triggers")
+        DatabaseUtils.execute("SELECT tww_sys.disable_symbology_triggers();")
 
     @staticmethod
     def enable_symbology_triggers():
-        with DatabaseUtils.PsycopgConnection() as connection:
-            cursor = connection.cursor()
-
-            logger.info("Enable symbology triggers")
-            cursor.execute("SELECT tww_sys.enable_symbology_triggers();")
+        logger.info("Enable symbology triggers")
+        DatabaseUtils.execute("SELECT tww_sys.enable_symbology_triggers();")
 
     @staticmethod
     def check_symbology_triggers_enabled():
-        with DatabaseUtils.PsycopgConnection() as connection:
-            cursor = connection.cursor()
-
-            logger.info("Enable symbology triggers")
-            cursor.execute("SELECT tww_sys.check_symbology_triggers_enabled();")
-
-            row = cursor.fetchone()
-            return row[0]
+        logger.info("Check symbology triggers enabled")
+        row = DatabaseUtils.fetchone("SELECT tww_sys.check_symbology_triggers_enabled();")
+        return row[0]
 
     @staticmethod
     def update_symbology():
@@ -154,44 +173,42 @@ def update_symbology():
     @staticmethod
     def check_oid_prefix() -> List[str]:
         """Check whether the oid_prefix is set up for production"""
-        with DatabaseUtils.PsycopgConnection() as connection:
-            cursor = connection.cursor()
+        logger.info("Checking setup of oid prefix")
+        prefixes = DatabaseUtils.fetchall("SELECT prefix FROM tww_sys.oid_prefixes WHERE active;")
 
-            logger.info("Checking setup of oid prefix")
-            cursor.execute("SELECT prefix FROM tww_sys.oid_prefixes WHERE active;")
-            prefixes = cursor.fetchall()
-            msg_list = []
-            if len(prefixes) > 1:
-                msg_list.append(
-                    "more than one oid_prefix set to active. Generation of Object-ID will not work. Set the OID prefix for your production database to active."
-                )
+        msg_list = []
+        if len(prefixes) > 1:
+            msg_list.append(
+                "more than one oid_prefix set to active. Generation of Object-ID will not work. Set the OID prefix for your production database to active."
+            )
 
-            if len(prefixes) > 0:
-                active_pref = prefixes[0][0]
-                if active_pref == "ch000000":
-                    msg_list.append(
-                        "OID prefix set to 'ch000000'. Database not safe for production"
-                    )
+        if len(prefixes) > 0:
+            active_pref = prefixes[0][0]
+            if active_pref == "ch000000":
+                msg_list.append("OID prefix set to 'ch000000'. Database not safe for production")
 
-            return msg_list
+        return msg_list
 
     @staticmethod
     def check_fk_defaults() -> List[str]:
         """Check whether the database is set up for production"""
-        with DatabaseUtils.PsycopgConnection() as connection:
-            cursor = connection.cursor()
+        logger.info("Checking setup of default_values")
+
+        defaults = []
+        vals = []
+        for item in DatabaseUtils.fetchall(
+            "SELECT fieldname,value_obj_id from tww_od.default_values;"
+        ):
+            defaults.append(item["fieldname"])
+            vals.append(item["value_obj_id"])
 
-            logger.info("Checking setup of default_values")
-            cursor.execute("SELECT fieldname,value_obj_id from tww_od.default_values;")
-            defaults = [item["fieldname"] for item in cursor.fetchall()]
-            vals = [item["value_obj_id"] for item in cursor.fetchall()]
-            msg_list = []
-            if None in vals:
-                msg_list.append("There is an undefined default value in tww_od.default_values")
-            elif not all(x in defaults for x in ["fk_provider", "fk_dataowner"]):
-                msg_list.append("'fk_provider' or 'fk_dataowner' not set in tww_od.default_values")
+        msg_list = []
+        if None in vals:
+            msg_list.append("There is an undefined default value in tww_od.default_values")
+        elif not all(x in defaults for x in ["fk_provider", "fk_dataowner"]):
+            msg_list.append("'fk_provider' or 'fk_dataowner' not set in tww_od.default_values")
 
-            return msg_list
+        return msg_list
 
     @staticmethod
     def get_validity_check_issues() -> List[str]:
@@ -203,11 +220,3 @@ def get_validity_check_issues() -> List[str]:
             messages.append("Symbology triggers are disabled")
 
         return messages
-
-    @staticmethod
-    def get_psycopg_connection():
-        connection = psycopg.connect(
-            DatabaseUtils.get_pgconf_as_psycopg_dsn(), **DEFAULTS_CONN_ARG
-        )
-        if PSYCOPG_VERSION == 2:
-            connection.set_session(autocommit=True)