feat: Use psycopg instead of psycopg2 for postgres.

Dockerfile

@@ -18,7 +18,7 @@ COPY src src
 
 RUN poetry build
 RUN (export version=$(find dist -name '*.whl'); \
-    pip install "${version}[s3,psycopg2]")
+    pip install "${version}[s3,postgres]")
 
 FROM python:3.9
 

Makefile

@@ -3,7 +3,7 @@
 VERSION=$(shell python -c 'from importlib import metadata; print(metadata.version("databudgie"))')
 
 install:
-	poetry install -E psycopg2-binary -E s3
+	poetry install -E postgres -E s3
 
 
 format:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databudgie"
-version = "2.8.5"
+version = "2.9.0"
 packages = [
     { include = "databudgie", from = "src" },
 ]
@@ -38,13 +38,11 @@ typing-extensions = {version = ">=3.10.0", python = "<3.8"}
 importlib-metadata = {version = "*", python = "<3.8"}
 
 boto3 = { version = "*", optional = true }
-psycopg2 = { version = ">=2.7", optional = true }
-psycopg2-binary = { version = ">=2.7", optional = true }
+psycopg = { version = ">=2.7" }
 
 [tool.poetry.extras]
 s3 = ["boto3"]
-psycopg2 = ["psycopg2"]
-psycopg2-binary = ["psycopg2-binary"]
+postgres = ["psycopg"]
 
 [tool.poetry.dev-dependencies]
 boto3 = "^1.16.10,<1.17.72"

src/databudgie/adapter/postgres.py

@@ -6,7 +6,7 @@
 import subprocess  # nosec
 from typing import cast, Dict, List
 
-from psycopg2._psycopg import cursor
+from psycopg import Cursor, sql
 from sqlalchemy import text
 from sqlalchemy.engine import create_engine, Engine
 from sqlalchemy.engine.url import URL
@@ -46,11 +46,13 @@ def export_query(self, query: str) -> QueryResult:
         result = QueryResult()
         with result.binary_buffer() as buffer:
             with contextlib.closing(engine.raw_connection()) as conn:
-                with cast(cursor, conn.cursor()) as cursor_:
-                    copy = f"COPY ({query}) TO STDOUT CSV HEADER"
+                with cast(Cursor, conn.cursor()) as cursor:
+                    command = sql.SQL(f"COPY ({query}) TO STDOUT CSV HEADER")
+                    with cursor.copy(command) as copy:
+                        for data in copy:
+                            buffer.write(data)
 
-                    cursor_.copy_expert(copy, buffer)
-                    result.row_count = cursor_.rowcount
+                    result.row_count = cursor.rowcount
 
         return result
 
@@ -59,11 +61,14 @@ def import_csv(self, csv_file: io.TextIOBase, table: str):
 
         # Reading the header line from the buffer removes it for the ingest
         columns: List[str] = [f'"{c}"' for c in csv_file.readline().strip().split(",")]
-        copy = "COPY {table} ({columns}) FROM STDIN CSV".format(table=table, columns=",".join(columns))
+        command = sql.SQL("COPY {table} ({columns}) FROM STDIN CSV".format(table=table, columns=",".join(columns)))
 
         with contextlib.closing(engine.raw_connection()) as conn:
-            with cast(cursor, conn.cursor()) as cursor_:
-                cursor_.copy_expert(copy, csv_file)
+            with cast(Cursor, conn.cursor()) as cursor:
+                with cursor.copy(command) as copy:
+                    while data := csv_file.read():
+                        copy.write(data)
+
                 conn.commit()
 
     def export_schema_ddl(self, name: str, console: Console = default_console) -> bytes:

src/databudgie/config.py

@@ -44,7 +44,7 @@ def push(self, config: dict):
 class Config(metaclass=abc.ABCMeta):
     def to_dict(self) -> dict:
         result = {}
-        for f in fields(self):
+        for f in fields(self):  # type: ignore
             v = getattr(self, f.name)
             if isinstance(v, Config):
                 value: Any = v.to_dict()

tests/cli/test_base.py

@@ -17,7 +17,7 @@ def test_create_postgres_session_str(pg_engine):
     except AttributeError:
         url_str = str(url)
 
-    assert url_str.startswith("postgresql+psycopg2://")
+    assert url_str.startswith("postgresql+psycopg://")
     config = BackupConfig.from_stack(ConfigStack({"url": url_str}))
 
     session = _create_postgres_session(config)
@@ -41,7 +41,7 @@ def test_connection_selection(pg_engine):
     except AttributeError:
         url_str = str(url)
 
-    assert url_str.startswith("postgresql+psycopg2://")
+    assert url_str.startswith("postgresql+psycopg://")
     config = BackupConfig.from_stack(ConfigStack({"connections": {"foo": url_str}, "connection": "foo"}))
 
     session = _create_postgres_session(config)