move to psycopg3

pirogue/information_schema.py

@@ -1,4 +1,4 @@
-from psycopg2.extensions import cursor
+from psycopg import Cursor
 
 from pirogue.exceptions import (
     InvalidSkipColumns,
@@ -7,7 +7,7 @@
 )
 
 
-def primary_key(pg_cur: cursor, schema_name: str, table_name: str) -> str:
+def primary_key(pg_cur: Cursor, schema_name: str, table_name: str) -> str:
     """
     Returns the primary of a table
 
@@ -38,7 +38,7 @@ def primary_key(pg_cur: cursor, schema_name: str, table_name: str) -> str:
 
 
 def columns(
-    pg_cur: cursor,
+    pg_cur: Cursor,
     table_schema: str,
     table_name: str,
     table_type: str = "table",
@@ -106,7 +106,7 @@ def columns(
 
 
 def reference_columns(
-    pg_cur: cursor,
+    pg_cur: Cursor,
     table_schema: str,
     table_name: str,
     foreign_table_schema: str,
@@ -157,7 +157,7 @@ def reference_columns(
     return cols
 
 
-def default_value(pg_cur: cursor, table_schema: str, table_name: str, column: str) -> str:
+def default_value(pg_cur: Cursor, table_schema: str, table_name: str, column: str) -> str:
     """
     Returns the default value of the column
 
@@ -188,7 +188,7 @@ def default_value(pg_cur: cursor, table_schema: str, table_name: str, column: st
 
 
 def geometry_type(
-    pg_cur: cursor, table_schema: str, table_name: str, column: str = "geometry"
+    pg_cur: Cursor, table_schema: str, table_name: str, column: str = "geometry"
 ) -> (str, int):
     """
     Returns the geometry type of a column as a tuple (type, srid)

pirogue/multiple_inheritance.py

@@ -1,7 +1,6 @@
 import os
 
-import psycopg2
-import psycopg2.extras
+import psycopg
 
 from pirogue.exceptions import InvalidDefinition, TableHasNoPrimaryKey, VariableError
 from pirogue.information_schema import (
@@ -57,7 +56,7 @@ def __init__(
         self.pg_service = pg_service
         if self.pg_service is None:
             self.pg_service = os.getenv("PGSERVICE")
-        self.conn = psycopg2.connect(f"service={self.pg_service}")
+        self.conn = psycopg.connect(f"service={self.pg_service}")
         self.cursor = self.conn.cursor()
 
         # check definition validity
@@ -193,7 +192,7 @@ def create(self) -> bool:
                 continue
             try:
                 if self.variables:
-                    self.cursor.execute(sql, self.variables)
+                    self.cursor.execute(psycopg.sql.SQL(sql).format(self.variables))
                 else:
                     self.cursor.execute(sql)
             except TypeError:
@@ -207,7 +206,7 @@ def create(self) -> bool:
                         svars=list(self.variables.keys())
                     )
                 )
-            except psycopg2.Error as e:
+            except psycopg.Error as e:
                 print(f"*** Failing:\n{sql}\n***")
                 raise e
         self.conn.commit()

pirogue/simple_joins.py

@@ -1,7 +1,6 @@
 import os
 
-import psycopg2
-import psycopg2.extras
+import psycopg
 
 from pirogue.exceptions import InvalidDefinition, NoReferenceFound, TableHasNoPrimaryKey
 from pirogue.information_schema import primary_key, reference_columns
@@ -44,7 +43,7 @@ def __init__(self, definition: dict, pg_service: str = None):
 
         if pg_service is None:
             pg_service = os.getenv("PGSERVICE")
-        self.conn = psycopg2.connect(f"service={pg_service}")
+        self.conn = psycopg.connect(f"service={pg_service}")
         self.cursor = self.conn.cursor()
 
         (self.parent_schema, self.parent_table) = table_parts(definition["table"])
@@ -98,7 +97,7 @@ def create(self) -> bool:
         success = True
         try:
             self.cursor.execute(sql)
-        except psycopg2.Error as e:
+        except psycopg.Error as e:
             success = False
             print(f"*** Failing:\n{sql}\n***")
             raise e

pirogue/single_inheritance.py

@@ -1,7 +1,6 @@
 import os
 
-import psycopg2
-import psycopg2.extras
+import psycopg
 
 from pirogue.exceptions import TableHasNoPrimaryKey
 from pirogue.information_schema import default_value, primary_key, reference_columns
@@ -46,7 +45,7 @@ def __init__(
 
         if pg_service is None:
             pg_service = os.getenv("PGSERVICE")
-        self.conn = psycopg2.connect(f"service={pg_service}")
+        self.conn = psycopg.connect(f"service={pg_service}")
         self.cursor = self.conn.cursor()
 
         self.pkey_default_value = pkey_default_value
@@ -96,7 +95,7 @@ def create(self) -> bool:
             try:
                 if sql:
                     self.cursor.execute(sql)
-            except psycopg2.Error as e:
+            except psycopg.Error as e:
                 success = False
                 print(f"*** Failing:\n{sql}\n***")
                 raise e

pirogue/utils.py

@@ -1,4 +1,4 @@
-from psycopg2.extensions import cursor
+from psycopg import Cursor
 
 from pirogue.exceptions import InvalidColumn, TableHasNoPrimaryKey
 from pirogue.information_schema import columns, default_value, primary_key
@@ -20,7 +20,7 @@ def table_parts(name: str) -> (str, str):
 
 
 def select_columns(
-    pg_cur: cursor,
+    pg_cur: Cursor,
     table_schema: str,
     table_name: str,
     table_type: str = "table",
@@ -143,7 +143,7 @@ def print_comma(first_column_printed, print: bool) -> str:
 
 
 def insert_command(
-    pg_cur: cursor,
+    pg_cur: Cursor,
     table_schema: str,
     table_name: str,
     table_type: str = "table",
@@ -199,8 +199,8 @@ def insert_command(
         add a prefix to the columns (do not applied to remapped columns)
     returning
         returning command
-    indent add
-         an indent in front
+    indent
+         add an indent in front
     inner_defaults
         dictionary of other columns to default to in case the provided value is null (can be used instead of insert_values to make it easier to reuse other columns definitions)
     """
@@ -306,7 +306,7 @@ def value(col):
 
 
 def update_command(
-    pg_cur: cursor,
+    pg_cur: Cursor,
     table_schema: str,
     table_name: str,
     table_alias: str = None,

requirements.txt

@@ -1,2 +1,2 @@
-psycopg2-binary
-pyyaml
+pyyaml>=6.0.1
+psycopg>=3.1.18

test/test_multiple_inheritance.py

@@ -2,8 +2,7 @@
 
 import unittest
 
-import psycopg2
-import psycopg2.extras
+import psycopg
 import yaml
 
 from pirogue import MultipleInheritance
@@ -15,7 +14,7 @@
 
 class TestMultipleInheritance(unittest.TestCase):
     def setUp(self):
-        self.conn = psycopg2.connect(f"service={pg_service}")
+        self.conn = psycopg.connect(f"service={pg_service}")
         self.cur = self.conn.cursor()
 
         sql = open("test/demo_data.sql").read()
@@ -68,7 +67,7 @@ def test_type_change_not_allowed(self):
             self.cur.execute(
                 "UPDATE pirogue_test.vw_animal_no_type_change SET animal_type = 'cat';"
             )
-        except psycopg2.errors.RaiseException:
+        except psycopg.errors.RaiseException:
             error_caught = True
         self.assertTrue(error_caught)
 

test/test_simple_joins.py

@@ -2,8 +2,7 @@
 
 import unittest
 
-import psycopg2
-import psycopg2.extras
+import psycopg
 import yaml
 
 from pirogue import MultipleInheritance, SimpleJoins
@@ -14,7 +13,7 @@
 
 class TestSimpleJoins(unittest.TestCase):
     def setUp(self):
-        self.conn = psycopg2.connect(f"service={pg_service}")
+        self.conn = psycopg.connect(f"service={pg_service}")
         self.cur = self.conn.cursor()
 
         sql = open("test/demo_data.sql").read()