dycw-utilities 0.150.0__tar.gz → 0.150.2__tar.gz

{dycw_utilities-0.150.0 → dycw_utilities-0.150.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.150.0
+Version: 0.150.2
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.150.0 → dycw_utilities-0.150.2}/pyproject.toml

@@ -103,7 +103,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.150.0"
+version = "0.150.2"
 
 [project.entry-points.pytest11]
 pytest-randomly = "utilities.pytest_plugins.pytest_randomly"
@@ -136,7 +136,7 @@ test = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.150.0"
+current_version = "0.150.2"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.150.0 → dycw_utilities-0.150.2}/src/tests/test_postgres.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from hypothesis import given
+from hypothesis import assume, given
 from hypothesis.strategies import DrawFn, booleans, composite, lists, none, sampled_from
 from pytest import raises
 from sqlalchemy import URL, Column, Integer, MetaData, Table
@@ -17,6 +17,8 @@ from utilities.postgres import (
     _ExtractURLPortError,
     _path_pg_dump,
     _PGDumpFormat,
+    _resolve_data_only_and_clean,
+    _ResolveDataOnlyAndCleanError,
     pg_dump,
     restore,
 )
@@ -63,14 +65,13 @@ class TestPGDump:
         path=temp_paths(),
         format_=sampled_from(get_literal_elements(_PGDumpFormat)),
         jobs=integers(min_value=0) | none(),
-        data_only=booleans(),
-        clean=booleans(),
         schema=lists(text_ascii(min_size=1)) | none(),
         schema_exc=lists(text_ascii(min_size=1)) | none(),
         table=tables() | none(),
         table_exc=tables() | none(),
         inserts=booleans(),
         on_conflict_do_nothing=booleans(),
+        role=text_ascii(min_size=1) | none(),
         docker=text_ascii(min_size=1) | none(),
     )
     def test_build(
@@ -80,14 +81,13 @@ class TestPGDump:
         path: Path,
         format_: _PGDumpFormat,
         jobs: int | None,
-        data_only: bool,
-        clean: bool,
         schema: list[str] | None,
         schema_exc: list[str] | None,
         table: list[Table | str] | None,
         table_exc: list[Table | str] | None,
         inserts: bool,
         on_conflict_do_nothing: bool,
+        role: str | None,
         docker: str | None,
     ) -> None:
         _ = _build_pg_dump(
@@ -95,14 +95,13 @@ class TestPGDump:
             path,
             format_=format_,
             jobs=jobs,
-            data_only=data_only,
-            clean=clean,
             schema=schema,
             schema_exc=schema_exc,
             table=table,
             table_exc=table_exc,
             inserts=inserts,
             on_conflict_do_nothing=on_conflict_do_nothing,
+            role=role,
             docker=docker,
         )
 
@@ -112,6 +111,20 @@ class TestPGDump:
         assert path.suffix in [".sql", ".pgdump", "", ".tar"]
 
 
+class TestResolveDataOnlyAndClean:
+    @given(data_only=booleans(), clean=booleans())
+    def test_main(self, *, data_only: bool, clean: bool) -> None:
+        _ = assume(not (data_only and clean))
+        _ = _resolve_data_only_and_clean(data_only=data_only, clean=clean)
+
+    def test_erorr(self) -> None:
+        with raises(
+            _ResolveDataOnlyAndCleanError,
+            match="Cannot use '--data-only' and '--clean' together",
+        ):
+            _ = _resolve_data_only_and_clean(data_only=True, clean=True)
+
+
 class TestRestore:
     @given(url=urls(), path=temp_paths(), logger=text_ascii(min_size=1) | none())
     async def test_main(self, *, url: URL, path: Path, logger: str | None) -> None:
@@ -122,12 +135,11 @@ class TestRestore:
         path=temp_paths(),
         psql=booleans(),
         database=text_ascii(min_size=1) | none(),
-        data_only=booleans(),
-        clean=booleans(),
         jobs=integers(min_value=0) | none(),
         schema=lists(text_ascii(min_size=1)) | none(),
         schema_exc=lists(text_ascii(min_size=1)) | none(),
         table=tables() | none(),
+        role=text_ascii(min_size=1) | none(),
         docker=text_ascii(min_size=1) | none(),
     )
     def test_build(
@@ -137,12 +149,11 @@ class TestRestore:
         path: Path,
         psql: bool,
         database: str | None,
-        data_only: bool,
-        clean: bool,
         jobs: int | None,
         schema: list[str] | None,
         schema_exc: list[str] | None,
         table: list[Table | str] | None,
+        role: str | None,
         docker: str | None,
     ) -> None:
         _ = _build_pg_restore_or_psql(
@@ -150,12 +161,11 @@ class TestRestore:
             path,
             psql=psql,
             database=database,
-            data_only=data_only,
-            clean=clean,
             jobs=jobs,
             schema=schema,
             schema_exc=schema_exc,
             table=table,
+            role=role,
             docker=docker,
         )
 

{dycw_utilities-0.150.0 → dycw_utilities-0.150.2}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.150.0"
+__version__ = "0.150.2"

{dycw_utilities-0.150.0 → dycw_utilities-0.150.2}/src/utilities/postgres.py

@@ -42,6 +42,7 @@ async def pg_dump(
     table_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
+    role: str | None = None,
     docker: str | None = None,
     dry_run: bool = False,
     logger: LoggerOrName | None = None,
@@ -62,6 +63,7 @@ async def pg_dump(
         table_exc=table_exc,
         inserts=inserts,
         on_conflict_do_nothing=on_conflict_do_nothing,
+        role=role,
         docker=docker,
     )
     if dry_run:
@@ -110,6 +112,7 @@ def _build_pg_dump(
     table_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
+    role: str | None = None,
     docker: str | None = None,
 ) -> str:
     database, host, port = _extract_url(url)
@@ -122,6 +125,7 @@ def _build_pg_dump(
         f"--format={format_}",
         "--verbose",
         # output options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
         "--large-objects",
         "--create",
         "--no-owner",
@@ -133,10 +137,6 @@ def _build_pg_dump(
     ]
     if (format_ == "directory") and (jobs is not None):
         parts.append(f"--jobs={jobs}")
-    if data_only:
-        parts.append("--data-only")
-    if clean:
-        parts.extend(["--clean", "--if-exists"])
     if schema is not None:
         parts.extend([f"--schema={s}" for s in always_iterable(schema)])
     if schema_exc is not None:
@@ -153,6 +153,8 @@ def _build_pg_dump(
         parts.append("--on-conflict-do-nothing")
     if url.username is not None:
         parts.append(f"--username={url.username}")
+    if role is not None:
+        parts.append(f"--role={role}")
     if docker is not None:
         parts = _wrap_docker(parts, docker)
     return " ".join(parts)
@@ -192,6 +194,7 @@ async def restore(
     schema: MaybeListStr | None = None,
     schema_exc: MaybeListStr | None = None,
     table: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
     docker: str | None = None,
     dry_run: bool = False,
     logger: LoggerOrName | None = None,
@@ -208,6 +211,7 @@ async def restore(
         schema=schema,
         schema_exc=schema_exc,
         table=table,
+        role=role,
         docker=docker,
     )
     if dry_run:
@@ -255,6 +259,7 @@ def _build_pg_restore_or_psql(
     schema: MaybeListStr | None = None,
     schema_exc: MaybeListStr | None = None,
     table: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
     docker: str | None = None,
 ) -> str:
     path = Path(path)
@@ -270,6 +275,7 @@ def _build_pg_restore_or_psql(
         schemas=schema,
         schemas_exc=schema_exc,
         tables=table,
+        role=role,
         docker=docker,
     )
 
@@ -286,6 +292,7 @@ def _build_pg_restore(
     schemas: MaybeListStr | None = None,
     schemas_exc: MaybeListStr | None = None,
     tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
     docker: str | None = None,
 ) -> str:
     """Run `pg_restore`."""
@@ -297,6 +304,8 @@ def _build_pg_restore(
         f"--dbname={database_use}",
         "--verbose",
         # restore options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
+        "--create",
         "--exit-on-error",
         "--no-owner",
         "--no-privileges",
@@ -305,10 +314,6 @@ def _build_pg_restore(
         f"--port={port}",
         "--no-password",
     ]
-    if data_only:
-        parts.append("--data-only")
-    if clean:
-        parts.extend(["--clean", "--if-exists"])
     if jobs is not None:
         parts.append(f"--jobs={jobs}")
     if schemas is not None:
@@ -319,6 +324,8 @@ def _build_pg_restore(
         parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(tables)])
     if url.username is not None:
         parts.append(f"--username={url.username}")
+    if role is not None:
+        parts.append(f"--role={role}")
     if docker is not None:
         parts = _wrap_docker(parts, docker)
     parts.append(str(path))
@@ -402,6 +409,29 @@ def _get_table_name(obj: TableOrORMInstOrClass | str, /) -> str:
             assert_never(never)
 
 
+def _resolve_data_only_and_clean(
+    *, data_only: bool = False, clean: bool = False
+) -> list[str]:
+    match data_only, clean:
+        case False, False:
+            return []
+        case True, False:
+            return ["--data-only"]
+        case False, True:
+            return ["--clean", "--if-exists"]
+        case True, True:
+            raise _ResolveDataOnlyAndCleanError
+        case _ as never:
+            assert_never(never)
+
+
+@dataclass(kw_only=True, slots=True)
+class _ResolveDataOnlyAndCleanError(Exception):
+    @override
+    def __str__(self) -> str:
+        return "Cannot use '--data-only' and '--clean' together"
+
+
 def _wrap_docker(parts: list[str], container: str, /) -> list[str]:
     return ["docker", "exec", "-it", container, *parts]
 

{dycw_utilities-0.150.0 → dycw_utilities-0.150.2}/src/utilities/sqlalchemy.py

@@ -57,7 +57,12 @@ from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
 from sqlalchemy.dialects.sqlite import Insert as sqlite_Insert
 from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-from sqlalchemy.exc import ArgumentError, DatabaseError, OperationalError
+from sqlalchemy.exc import (
+    ArgumentError,
+    DatabaseError,
+    OperationalError,
+    ProgrammingError,
+)
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 from sqlalchemy.orm import (
     DeclarativeBase,
@@ -115,7 +120,7 @@ def check_connect(engine: Engine, /) -> bool:
     try:
         with engine.connect() as conn:
             return bool(conn.execute(text("SELECT 1")).scalar_one())
-    except OperationalError:
+    except (OperationalError, ProgrammingError):
         return False