handle cursors

teksi · Apr 9, 2024 · f303621 · f303621
1 parent 2f298eb
commit f303621
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 6 deletions.
diff --git a/datamodel/test/test_network.py b/datamodel/test/test_network.py
@@ -5,6 +5,7 @@
     import psycopg
 except ImportError:
     import psycopg2 as psycopg
+    import psycopg2.exras as psycopg_extras
 
 from .utils import DEFAULT_PG_SERVICE, DbTestBase
 
@@ -91,7 +92,11 @@ def downstream_nodes_depths(self, node_id):
         SELECT obj_id, depth
         FROM downstream;
         """
-        cur = self.cursor(row_factory=psycopg.rows.dict_row)
+        try:
+            cur = self.cursor(row_factory=psycopg.rows.dict_row)
+        except AttributeError:
+            # remove when dropping psycopg2 support
+            cur = self.cursor(cursor_factory=psycopg_extras.DictCursor)
         cur.execute(query, (node_id,))
         return {row["obj_id"]: row["depth"] for row in cur.fetchall()}
 
@@ -118,7 +123,11 @@ def upstream_nodes_depths(self, node_id):
         SELECT obj_id, depth
         FROM upstream
         """
-        cur = self.cursor(row_factory=psycopg.rows.dict_row)
+        try:
+            cur = self.cursor(row_factory=psycopg.rows.dict_row)
+        except AttributeError:
+            # remove when dropping psycopg2 support
+            cur = self.cursor(cursor_factory=psycopg_extras.DictCursor)
         cur.execute(query, (node_id,))
         return {row["obj_id"]: row["depth"] for row in cur.fetchall()}
 

diff --git a/datamodel/test/test_removed_fields.py b/datamodel/test/test_removed_fields.py
@@ -5,6 +5,7 @@
     import psycopg
 except ImportError:
     import psycopg2 as psycopg
+    import psycopg2.exras as psycopg_extras
 
 from .utils import DEFAULT_PG_SERVICE, DbTestBase
 
@@ -20,7 +21,11 @@ def setUpClass(cls):
         cls.conn = psycopg.connect(f"service={pgservice}")
 
     def test_dataowner(self):
-        cur = self.conn.cursor(row_factory=psycopg.rows.dict_row)
+        try:
+            cur = self.conn.cursor(row_factory=psycopg.rows.dict_row)
+        except AttributeError:
+            # remove when dropping psycopg2 support
+            cur = self.conn.cursor(cursor_factory=psycopg_extras.DictCursor)
 
         cur.execute("SELECT * FROM tww_od.wastewater_structure LIMIT 1")
         colnames = [desc[0] for desc in cur.description]

diff --git a/datamodel/test/test_views.py b/datamodel/test/test_views.py
@@ -7,6 +7,7 @@
     import psycopg
 except ImportError:
     import psycopg2 as psycopg
+    import psycopg2.exras as psycopg_extras
 
 from .utils import DEFAULT_PG_SERVICE, DbTestBase
 
@@ -95,7 +96,11 @@ def test_vw_tww_wastewater_structure(self):
 
         self.update_check("vw_tww_wastewater_structure", row, obj_id)
 
-        cur = self.cursor(row_factory=psycopg.rows.dict_row)
+        try:
+            cur = self.cursor(row_factory=psycopg.rows.dict_row)
+        except AttributeError:
+            # remove when dropping psycopg2 support
+            cur = self.cursor(cursor_factory=psycopg_extras.DictCursor)
 
         cur.execute(
             "SELECT * FROM tww_od.wastewater_networkelement NE LEFT JOIN tww_od.wastewater_node NO ON NO.obj_id = NE.obj_id WHERE fk_wastewater_structure='{obj_id}' ".format(

diff --git a/datamodel/test/utils.py b/datamodel/test/utils.py
@@ -2,6 +2,7 @@
     import psycopg
 except ImportError:
     import psycopg2 as psycopg
+    import psycopg2.exras as psycopg_extras
 
 DEFAULT_PG_SERVICE = "pg_tww"
 
@@ -26,8 +27,11 @@ def geom_as_text(cls, wkb: str) -> str:
 
     @classmethod
     def select(cls, table, obj_id, schema="tww_app"):
-        cur = cls.conn.cursor(row_factory=psycopg.rows.dict_row)
-        # cur = cls.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
+        try:
+            cur = cls.conn.cursor(row_factory=psycopg.rows.dict_row)
+        except AttributeError:
+            # remove when dropping psycopg2 support
+            cur = cls.conn.cursor(cursor_factory=psycopg_extras.DictCursor)
         cur.execute(f"SELECT * FROM {schema}.{table} WHERE obj_id=%(obj_id)s", {"obj_id": obj_id})
         return cur.fetchone()