Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Use psycopg instead of psycopg2 for postgres. #120

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
1,196 changes: 586 additions & 610 deletions poetry.lock

Large diffs are not rendered by default.

12 changes: 5 additions & 7 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "databudgie"
version = "2.8.5"
version = "2.9.0"
packages = [
{ include = "databudgie", from = "src" },
]
Expand All @@ -27,7 +27,7 @@ repository = "https://github.com/schireson/databudgie"
databudgie = "databudgie.__main__:run"

[tool.poetry.dependencies]
python = ">=3.7,<4"
python = ">=3.8,<4"

rich = "*"
configly = {version = ">=1.0.0", extras = ["yaml"]}
Expand All @@ -38,21 +38,19 @@ 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"
black = "22.3.0"
boto3-stubs = {extras = ["s3"], version = "^1.18.38"}
coverage = ">=5"
freezegun = "*"
mypy = "^0.991"
mypy = "^1.10.0"
pytest = ">=6.2.4"
pytest-mock-resources = {version = ">=2.1.10", extras = ["docker"]}
responses = "^0.10.9"
Expand Down
21 changes: 13 additions & 8 deletions src/databudgie/adapter/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand All @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion src/databudgie/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
4 changes: 2 additions & 2 deletions tests/cli/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down
Loading