Merge pull request #120 from schireson/dc/psycopg

feat: Use psycopg instead of psycopg2 for postgres.

Author: DanCardin

.github/workflows/build.yml

@@ -37,7 +37,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Install dependencies
-        run: uv sync --extra psycopg2-binary --extra s3
+        run: uv sync --extra postgres --extra s3
 
       - name: Install specific sqlalchemy version
         run: uv pip install 'sqlalchemy~=${{ matrix.sqlalchemy-version }}'

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+### v2.9.0
+
+* Swap to psycopg from psycopg2
+
 ### [v2.8.5](https://github.com/schireson/databudgie/compare/v2.8.4...v2.8.5) (2024-02-20)
 
 #### Features

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,17 +3,16 @@
 VERSION=$(shell python -c 'from importlib import metadata; print(metadata.version("databudgie"))')
 
 install:
-	uv sync --extra psycopg2-binary --extra s3
+	uv sync --extra postgres --extra s3
 
 
 format:
 	uv run ruff check --fix src tests
-	uv run black src tests
+	uv run ruff format src tests
 
 lint:
 	uv run ruff check src tests || exit 1
 	uv run mypy --namespace-packages src tests || exit 1
-	uv run black --check --diff src tests || exit 1
 
 test:
 	uv run coverage run -a -m pytest src tests

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databudgie"
-version = "2.8.6"
+version = "2.9.0"
 description = "Ergonomic and flexible tool for database backup and restore"
 readme = "README.md"
 license = { text = "MIT" }
@@ -20,26 +20,7 @@ dependencies = [
 
 [project.optional-dependencies]
 s3 = ["boto3"]
-psycopg2 = ["psycopg2>=2.7"]
-psycopg2-binary = ["psycopg2-binary>=2.7"]
-dev = [
-    "boto3==1.34.100",
-    "black==22.3.0",
-    "boto3-stubs[s3]>=1.18.38",
-    "coverage>=5",
-    "freezegun",
-    "mypy>=1.14",
-    "pytest>=6.2.4",
-    "pytest-mock-resources[docker]>=2.1.10",
-    "responses>=0.10.9",
-    "ruff==0.0.254",
-    "types-freezegun>=0.1.3",
-    "types-requests>=0.1.11",
-    "faker>=8.12.1",
-    "moto[s3]>=5.0.0",
-    "sqlalchemy-model-factory>=0.4.5",
-    "types-click>=7.1.5",
-]
+postgres = ["psycopg[binary]>=3.0"]
 
 [project.scripts]
 databudgie = "databudgie.__main__:run"
@@ -52,16 +33,15 @@ package = true
 
 [dependency-groups]
 dev = [
-    "boto3==1.34.100",
-    "black==22.3.0",
-    "boto3-stubs[s3]>=1.18.38",
+    "boto3",
+    "boto3-stubs[s3]",
     "coverage>=5",
     "freezegun",
     "mypy>=0.991",
     "pytest>=6.2.4",
     "pytest-mock-resources[docker]>=2.1.10",
     "responses>=0.10.9",
-    "ruff==0.0.254",
+    "ruff>=0.9",
     "types-freezegun>=0.1.3",
     "types-requests>=0.1.11",
     "faker>=8.12.1",
@@ -70,16 +50,40 @@ dev = [
     "types-click>=7.1.5",
 ]
 
-[tool.black]
-line_length = 120
+[tool.ruff]
+line-length = 120
+src = ["src", "tests"]
+
+[tool.ruff.lint]
+select = ["C", "D", "E", "F", "I", "N", "Q", "RET", "RUF", "S", "T", "UP", "YTT"]
+ignore = ["C901", "E501", "S101", "D1", "D203", "D213", "D406", "D407", "D408", "D409", "D413"]
+extend-ignore = [
+    "D1",
+
+    "D203",
+    "D204",
+    "D213",
+    "D215",
+    "D400",
+    "D404",
+    "D406",
+    "D407",
+    "D408",
+    "D409",
+    "D413",
+]
+
+[tool.ruff.lint.isort]
+order-by-type = false
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["T201", "D", "S", "N801", "N802", 'N806']
 
-[tool.isort]
-profile = 'black'
-known_first_party = 'databudgie,tests'
-line_length = 120
-float_to_top = true
-order_by_type = false
-use_parentheses = true
+[tool.ruff.lint.pyupgrade]
+keep-runtime-typing = true
+
+[tool.ruff.format]
+docstring-code-format = true
 
 [tool.mypy]
 strict_optional = true
@@ -100,10 +104,6 @@ exclude_lines = ["pragma: no cover", "if TYPE_CHECKING:", "if __name__ == .__mai
 source = ["src"]
 branch = true
 
-[tool.pydocstyle]
-ignore = 'D1,D200,D202,D203,D204,D213,D406,D407,D413'
-match_dir = '^[^\.{]((?!igrations).)*'
-
 [tool.pytest.ini_options]
 doctest_optionflags = "NORMALIZE_WHITESPACE IGNORE_EXCEPTION_DETAIL ELLIPSIS"
 addopts = "--doctest-modules -vv --ff --strict-markers"

ruff.toml

@@ -3,13 +3,14 @@
 import os
 import shlex
 import shutil
-import subprocess  # nosec
+import subprocess
 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
+from typing_extensions import LiteralString
 
 from databudgie.adapter.base import Adapter, QueryResult
 from databudgie.output import Console, default_console
@@ -46,24 +47,29 @@ 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:
+                    statement = sql.SQL(cast(LiteralString, f"COPY ({query}) TO STDOUT CSV HEADER"))
 
-                    cursor_.copy_expert(copy, buffer)
-                    result.row_count = cursor_.rowcount
+                    with cursor.copy(statement) as copy:
+                        while data := copy.read():
+                            buffer.write(data)
+                    result.row_count = cursor.rowcount
 
         return result
 
     def import_csv(self, csv_file: io.TextIOBase, table: str):
         engine: Engine = cast(Engine, self.session.get_bind())
 
         # 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))
+        columns = [sql.Identifier(c) for c in csv_file.readline().strip().split(",")]
+        statement = sql.SQL("COPY {table} ({columns}) FROM STDIN CSV").format(
+            table=sql.Identifier(*table.split(".")), columns=sql.SQL(", ").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(statement) as copy:
+                    copy.write(csv_file.read())
                 conn.commit()
 
     def export_schema_ddl(self, name: str, console: Console = default_console) -> bytes:
@@ -73,7 +79,10 @@ def export_schema_ddl(self, name: str, console: Console = default_console) -> by
 
         url = self.session.connection().engine.url
         result = pg_dump(url, f"--schema-only --schema={name} --exclude-table={name}.*")
-        return result.replace(f"CREATE SCHEMA {name};".encode(), f"CREATE SCHEMA IF NOT EXISTS {name};".encode())
+        return result.replace(
+            f"CREATE SCHEMA {name};".encode(),
+            f"CREATE SCHEMA IF NOT EXISTS {name};".encode(),
+        )
 
     def export_table_ddl(self, table_name: str, console: Console = default_console):
         if not shutil.which("pg_dump"):
@@ -205,7 +214,8 @@ def collect_table_dependencies(self, table_op: TableOp, console: Console = defau
         )
 
         results = self.session.execute(
-            collect_tables, params={"schema": table_op.schema, "table_name": table_op.table_name}
+            collect_tables,
+            params={"schema": table_op.schema, "table_name": table_op.table_name},
         )
 
         return [row[0] for row in results]
@@ -235,10 +245,18 @@ def collect_table_sequences(self) -> Dict[str, List[str]]:
         return result
 
     def collect_sequence_value(self, sequence_name: str) -> int:
-        return cast(int, self.session.execute(text(f"SELECT last_value from {sequence_name}")).scalar())  # noqa: S608
+        return cast(
+            int,
+            self.session.execute(
+                text(f"SELECT last_value from {sequence_name}")  # noqa: S608
+            ).scalar(),
+        )
 
     def restore_sequence_value(self, sequence_name: str, value: int) -> int:
-        return cast(int, self.session.execute(text(f"SELECT setval('{sequence_name}', {value})")).scalar())
+        return cast(
+            int,
+            self.session.execute(text(f"SELECT setval('{sequence_name}', {value})")).scalar(),
+        )
 
 
 def pg_dump(url: URL, rest: str = "", no_comments=True, clean=True) -> bytes:
@@ -254,8 +272,11 @@ def pg_dump(url: URL, rest: str = "", no_comments=True, clean=True) -> bytes:
     command = shlex.split(raw_command)
 
     try:
-        result = subprocess.run(  # nosec
-            command, capture_output=True, env={**os.environ, "PGPASSWORD": str(url.password or "")}, check=True
+        result = subprocess.run(  # noqa: S603
+            command,
+            capture_output=True,
+            env={**os.environ, "PGPASSWORD": str(url.password or "")},
+            check=True,
         )
     except subprocess.CalledProcessError as e:
         raise RuntimeError(e.stderr)

src/databudgie/adapter/postgres.py

@@ -1,13 +1,13 @@
 import gzip
 import io
-from typing import Dict, Optional, Type
+from typing import ClassVar, Dict, Optional, Type
 
 
 class Compressor:
     name: str = ""
     extension: str = ""
 
-    compressors: Dict[Optional[str], "Compressor"] = {}
+    compressors: ClassVar[Dict[Optional[str], "Compressor"]] = {}
 
     @classmethod
     def get_with_name(cls, name: Optional[str]) -> "Compressor":

src/databudgie/compression.py

@@ -17,14 +17,14 @@ def __init__(self, session: Session, table_name: str, action: str):
 
         self._transaction_id: Optional[int] = None
 
-    @functools.lru_cache()
+    @functools.lru_cache
     def manifest_table(self):
         schema, table = parse_table(self.table_name)
         self.metadata = MetaData()
         self.metadata.reflect(bind=self.session.get_bind(), schema=schema, only=[table])
         return Table(table, self.metadata, autoload=True, schema=schema)
 
-    @functools.lru_cache()
+    @functools.lru_cache
     def transaction_id(self):
         if self._transaction_id is None:
             table = self.manifest_table()

src/databudgie/manifest/manager.py

@@ -49,8 +49,8 @@ def update(self, task, description, advance=1):
 
 __all__ = [
     "Console",
-    "default_console",
     "Progress",
-    "Traceback",
     "Table",
+    "Traceback",
+    "default_console",
 ]