dycw-utilities 0.148.2__tar.gz → 0.148.3__tar.gz

{dycw_utilities-0.148.2 → dycw_utilities-0.148.3}/PKG-INFO

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

{dycw_utilities-0.148.2 → dycw_utilities-0.148.3}/pyproject.toml

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

{dycw_utilities-0.148.2 → dycw_utilities-0.148.3}/src/tests/test_postgres.py

@@ -26,10 +26,13 @@ if TYPE_CHECKING:
 
 
 @composite
-def tables(draw: DrawFn, /) -> list[Table]:
+def tables(draw: DrawFn, /) -> list[Table | str]:
     metadata = MetaData()
-    tables = draw(lists(text_ascii(min_size=1), max_size=5, unique=True))
-    return [Table(t, metadata, Column("id", Integer, primary_key=True)) for t in tables]
+    names = draw(lists(text_ascii(min_size=1), max_size=5, unique=True))
+    tables = [
+        Table(n, metadata, Column("id", Integer, primary_key=True)) for n in names
+    ]
+    return [draw(sampled_from([n, t])) for n, t in zip(names, tables, strict=True)]
 
 
 @composite
@@ -57,7 +60,10 @@ class TestPGDump:
         format_=sampled_from(get_literal_elements(_PGDumpFormat)),
         jobs=integers(min_value=0) | none(),
         schemas=lists(text_ascii(min_size=1)) | none(),
+        schemas_exc=lists(text_ascii(min_size=1)) | none(),
         tables=tables() | none(),
+        tables_exc=tables() | none(),
+        inserts=booleans(),
         logger=text_ascii(min_size=1) | none(),
     )
     async def test_main(
@@ -69,7 +75,10 @@ class TestPGDump:
         format_: _PGDumpFormat,
         jobs: int | None,
         schemas: list[str] | None,
-        tables: list[Table] | None,
+        schemas_exc: list[str] | None,
+        tables: list[Table | str] | None,
+        tables_exc: list[Table | str] | None,
+        inserts: bool,
         logger: str | None,
     ) -> None:
         _ = await pg_dump(
@@ -79,7 +88,10 @@ class TestPGDump:
             format_=format_,
             jobs=jobs,
             schemas=schemas,
+            schemas_exc=schemas_exc,
             tables=tables,
+            tables_exc=tables_exc,
+            inserts=inserts,
             logger=logger,
             dry_run=True,
         )
@@ -111,6 +123,7 @@ class TestPGRestore:
         data_only=booleans(),
         jobs=integers(min_value=0) | none(),
         schemas=lists(text_ascii(min_size=1)) | none(),
+        schemas_exc=lists(text_ascii(min_size=1)) | none(),
         tables=tables() | none(),
         logger=text_ascii(min_size=1) | none(),
     )
@@ -124,7 +137,8 @@ class TestPGRestore:
         data_only: bool,
         jobs: int | None,
         schemas: list[str] | None,
-        tables: list[Table] | None,
+        schemas_exc: list[str] | None,
+        tables: list[Table | str] | None,
         logger: str | None,
     ) -> None:
         _ = await pg_restore(
@@ -135,6 +149,7 @@ class TestPGRestore:
             data_only=data_only,
             jobs=jobs,
             schemas=schemas,
+            schemas_exc=schemas_exc,
             tables=tables,
             logger=logger,
             dry_run=True,

{dycw_utilities-0.148.2 → dycw_utilities-0.148.3}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.148.2"
+__version__ = "0.148.3"

{dycw_utilities-0.148.2 → dycw_utilities-0.148.3}/src/utilities/postgres.py

@@ -5,17 +5,21 @@ from pathlib import Path
 from shutil import rmtree
 from typing import TYPE_CHECKING, Literal, assert_never, override
 
+from sqlalchemy import Table
+from sqlalchemy.orm import DeclarativeBase
+
 from utilities.asyncio import stream_command
 from utilities.iterables import always_iterable
 from utilities.logging import get_logger
 from utilities.os import temp_environ
-from utilities.sqlalchemy import TableOrORMInstOrClass, get_table_name
+from utilities.sqlalchemy import get_table_name
 from utilities.timer import Timer
 from utilities.types import PathLike
 
 if TYPE_CHECKING:
     from sqlalchemy import URL
 
+    from utilities.sqlalchemy import TableOrORMInstOrClass
     from utilities.types import LoggerOrName, MaybeListStr, MaybeSequence, PathLike
 
 
@@ -31,7 +35,10 @@ async def pg_dump(
     format_: _PGDumpFormat = "plain",
     jobs: int | None = None,
     schemas: MaybeListStr | None = None,
-    tables: MaybeSequence[TableOrORMInstOrClass] | None = None,
+    schemas_exc: MaybeListStr | None = None,
+    tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    tables_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    inserts: bool = False,
     logger: LoggerOrName | None = None,
     dry_run: bool = False,
 ) -> None:
@@ -69,8 +76,16 @@ async def pg_dump(
         parts.append(f"--jobs={jobs}")
     if schemas is not None:
         parts.extend([f"--schema={s}" for s in always_iterable(schemas)])
+    if schemas_exc is not None:
+        parts.extend([f"--exclude-schema={s}" for s in always_iterable(schemas_exc)])
     if tables is not None:
-        parts.extend([f"--table={get_table_name(t)}" for t in always_iterable(tables)])
+        parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(tables)])
+    if tables_exc is not None:
+        parts.extend([
+            f"--exclude-table={_get_table_name(t)}" for t in always_iterable(tables_exc)
+        ])
+    if inserts:
+        parts.append("--inserts")
     if url.username is not None:
         parts.append(f"--username={url.username}")
     cmd = " ".join(parts)
@@ -144,7 +159,8 @@ async def pg_restore(
     data_only: bool = False,
     jobs: int | None = None,
     schemas: MaybeListStr | None = None,
-    tables: MaybeSequence[TableOrORMInstOrClass] | None = None,
+    schemas_exc: MaybeListStr | None = None,
+    tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     logger: LoggerOrName | None = None,
     dry_run: bool = False,
 ) -> None:
@@ -174,7 +190,6 @@ async def pg_restore(
         "--exit-on-error",
         "--no-owner",
         "--no-privileges",
-        "--if-exists",
         # connection options
         f"--host={url.host}",
         f"--port={url.port}",
@@ -183,13 +198,15 @@ async def pg_restore(
     if data_only:
         parts.append("--data-only")
     else:
-        parts.append("--clean")
+        parts.extend(["--clean", "--if-exists"])
     if jobs is not None:
         parts.append(f"--jobs={jobs}")
     if schemas is not None:
         parts.extend([f"--schema={s}" for s in always_iterable(schemas)])
+    if schemas_exc is not None:
+        parts.extend([f"--exclude-schema={s}" for s in always_iterable(schemas_exc)])
     if tables is not None:
-        parts.extend([f"--table={get_table_name(t)}" for t in always_iterable(tables)])
+        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}")
     parts.append(str(path))
@@ -249,4 +266,18 @@ class _PGRestorePortError(PGRestoreError):
         return f"Expected URL to contain a 'port'; got {self.url}"
 
 
+##
+
+
+def _get_table_name(obj: TableOrORMInstOrClass | str, /) -> str:
+    """Get the table name from a Table or mapped class."""
+    match obj:
+        case Table() | DeclarativeBase() | type() as table_or_orm:
+            return get_table_name(table_or_orm)
+        case str() as name:
+            return name
+        case _ as never:
+            assert_never(never)
+
+
 __all__ = ["PGDumpError", "PGRestoreError", "pg_dump", "pg_restore"]