dycw-utilities 0.148.5__tar.gz → 0.149.1__tar.gz

{dycw_utilities-0.148.5 → dycw_utilities-0.149.1}/PKG-INFO

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

{dycw_utilities-0.148.5 → dycw_utilities-0.149.1}/pyproject.toml

@@ -102,7 +102,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.148.5"
+version = "0.149.1"
 
 [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.5"
+current_version = "0.149.1"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.148.5 → dycw_utilities-0.149.1}/src/tests/test_postgres.py

@@ -9,15 +9,16 @@ from sqlalchemy import URL, Column, Integer, MetaData, Table
 
 from utilities.hypothesis import integers, temp_paths, text_ascii
 from utilities.postgres import (
-    _PGDumpDatabaseError,
+    _build_pg_dump,
+    _build_pg_restore_or_psql,
+    _extract_url,
+    _ExtractURLDatabaseError,
+    _ExtractURLHostError,
+    _ExtractURLPortError,
+    _path_pg_dump,
     _PGDumpFormat,
-    _PGDumpHostError,
-    _PGDumpPortError,
-    _PGRestoreDatabaseError,
-    _PGRestoreHostError,
-    _PGRestorePortError,
     pg_dump,
-    pg_restore,
+    restore,
 )
 from utilities.typing import get_literal_elements
 
@@ -53,11 +54,13 @@ def urls(draw: DrawFn, /) -> URL:
 
 
 class TestPGDump:
+    @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:
+        _ = await pg_dump(url, path, dry_run=True, logger=logger)
+
     @given(
         url=urls(),
         path=temp_paths(),
-        docker=text_ascii(min_size=1) | none(),
-        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(),
@@ -65,15 +68,13 @@ class TestPGDump:
         tables_exc=tables() | none(),
         inserts=booleans(),
         on_conflict_do_nothing=booleans(),
-        logger=text_ascii(min_size=1) | none(),
+        docker=text_ascii(min_size=1) | none(),
     )
-    async def test_main(
+    def test_build(
         self,
         *,
         url: URL,
         path: Path,
-        docker: str | None,
-        format_: _PGDumpFormat,
         jobs: int | None,
         schemas: list[str] | None,
         schemas_exc: list[str] | None,
@@ -81,13 +82,11 @@ class TestPGDump:
         tables_exc: list[Table | str] | None,
         inserts: bool,
         on_conflict_do_nothing: bool,
-        logger: str | None,
+        docker: str | None,
     ) -> None:
-        _ = await pg_dump(
+        _ = _build_pg_dump(
             url,
             path,
-            docker=docker,
-            format_=format_,
             jobs=jobs,
             schemas=schemas,
             schemas_exc=schemas_exc,
@@ -95,87 +94,79 @@ class TestPGDump:
             tables_exc=tables_exc,
             inserts=inserts,
             on_conflict_do_nothing=on_conflict_do_nothing,
-            logger=logger,
-            dry_run=True,
+            docker=docker,
         )
 
-    async def test_error_database(self, *, tmp_path: Path) -> None:
-        url = URL.create("postgres")
-        with raises(
-            _PGDumpDatabaseError, match="Expected URL to contain a 'database'; got .*"
-        ):
-            _ = await pg_dump(url, tmp_path, dry_run=True)
-
-    async def test_error_host(self, *, tmp_path: Path) -> None:
-        url = URL.create("postgres", database="database")
-        with raises(_PGDumpHostError, match="Expected URL to contain a 'host'; got .*"):
-            _ = await pg_dump(url, tmp_path, dry_run=True)
+    @given(path=temp_paths(), format_=sampled_from(get_literal_elements(_PGDumpFormat)))
+    def test_path(self, *, path: Path, format_: _PGDumpFormat) -> None:
+        path = _path_pg_dump(path, format_=format_)
+        assert path.suffix in [".sql", ".pgdump", "", ".tar"]
 
-    async def test_error_port(self, *, tmp_path: Path) -> None:
-        url = URL.create("postgres", database="database", host="host")
-        with raises(_PGDumpPortError, match="Expected URL to contain a 'port'; got .*"):
-            _ = await pg_dump(url, tmp_path, dry_run=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:
+        _ = await restore(url, path, dry_run=True, logger=logger)
 
-class TestPGRestore:
     @given(
         url=urls(),
         path=temp_paths(),
+        psql=booleans(),
         database=text_ascii(min_size=1) | none(),
-        docker=text_ascii(min_size=1) | none(),
         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(),
+        docker=text_ascii(min_size=1) | none(),
     )
-    async def test_main(
+    def test_build(
         self,
         *,
         url: URL,
         path: Path,
+        psql: bool,
         database: str | None,
-        docker: str | None,
         data_only: bool,
         jobs: int | None,
         schemas: list[str] | None,
         schemas_exc: list[str] | None,
         tables: list[Table | str] | None,
-        logger: str | None,
+        docker: str | None,
     ) -> None:
-        _ = await pg_restore(
+        _ = _build_pg_restore_or_psql(
             url,
             path,
+            psql=psql,
             database=database,
-            docker=docker,
             data_only=data_only,
             jobs=jobs,
             schemas=schemas,
             schemas_exc=schemas_exc,
             tables=tables,
-            logger=logger,
-            dry_run=True,
+            docker=docker,
         )
 
-    async def test_error_database(self, *, tmp_path: Path) -> None:
+
+class TestExtractURL:
+    def test_database(self) -> None:
         url = URL.create("postgres")
         with raises(
-            _PGRestoreDatabaseError,
+            _ExtractURLDatabaseError,
             match="Expected URL to contain a 'database'; got .*",
         ):
-            _ = await pg_restore(url, tmp_path, dry_run=True)
+            _ = _extract_url(url)
 
-    async def test_error_host(self, *, tmp_path: Path) -> None:
+    def test_host(self) -> None:
         url = URL.create("postgres", database="database")
         with raises(
-            _PGRestoreHostError, match="Expected URL to contain a 'host'; got .*"
+            _ExtractURLHostError, match="Expected URL to contain a 'host'; got .*"
         ):
-            _ = await pg_restore(url, tmp_path, dry_run=True)
+            _ = _extract_url(url)
 
-    async def test_error_port(self, *, tmp_path: Path) -> None:
+    def test_port(self) -> None:
         url = URL.create("postgres", database="database", host="host")
         with raises(
-            _PGRestorePortError, match="Expected URL to contain a 'port'; got .*"
+            _ExtractURLPortError, match="Expected URL to contain a 'port'; got .*"
         ):
-            _ = await pg_restore(url, tmp_path, dry_run=True)
+            _ = _extract_url(url)

{dycw_utilities-0.148.5 → dycw_utilities-0.149.1}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.148.5"
+__version__ = "0.149.1"

{dycw_utilities-0.148.5 → dycw_utilities-0.149.1}/src/utilities/postgres.py

@@ -12,6 +12,7 @@ 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.pathlib import ensure_suffix
 from utilities.sqlalchemy import get_table_name
 from utilities.timer import Timer
 from utilities.types import PathLike
@@ -31,7 +32,6 @@ async def pg_dump(
     path: PathLike,
     /,
     *,
-    docker: str | None = None,
     format_: _PGDumpFormat = "plain",
     jobs: int | None = None,
     schemas: MaybeListStr | None = None,
@@ -40,25 +40,78 @@ async def pg_dump(
     tables_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
-    logger: LoggerOrName | None = None,
+    docker: str | None = None,
     dry_run: bool = False,
+    logger: LoggerOrName | None = None,
 ) -> None:
     """Run `pg_dump`."""
-    path = Path(path)
+    path = _path_pg_dump(path, format_=format_)
     path.parent.mkdir(parents=True, exist_ok=True)
-    if url.database is None:
-        raise _PGDumpDatabaseError(url=url)
-    if url.host is None:
-        raise _PGDumpHostError(url=url)
-    if url.port is None:
-        raise _PGDumpPortError(url=url)
-    parts: list[str] = []
-    if docker is not None:
-        parts.extend(["docker", "exec", "-it", docker])
-    parts.extend([
+    cmd = _build_pg_dump(
+        url,
+        path,
+        format_=format_,
+        jobs=jobs,
+        schemas=schemas,
+        schemas_exc=schemas_exc,
+        tables=tables,
+        tables_exc=tables_exc,
+        inserts=inserts,
+        on_conflict_do_nothing=on_conflict_do_nothing,
+        docker=docker,
+    )
+    if dry_run:
+        if logger is not None:
+            get_logger(logger=logger).info("Would run %r", str(cmd))
+        return
+    with temp_environ(PGPASSWORD=url.password), Timer() as timer:  # pragma: no cover
+        try:
+            output = await stream_command(cmd)
+        except KeyboardInterrupt:
+            if logger is not None:
+                get_logger(logger=logger).info(
+                    "Cancelled backup to %r after %s", str(path), timer
+                )
+            rmtree(path, ignore_errors=True)
+        else:
+            match output.return_code:
+                case 0:
+                    if logger is not None:
+                        get_logger(logger=logger).info(
+                            "Backup to %r finished after %s", str(path), timer
+                        )
+                case _:
+                    if logger is not None:
+                        get_logger(logger=logger).exception(
+                            "Backup to %r failed after %s\nstderr:\n%s",
+                            str(path),
+                            timer,
+                            output.stderr,
+                        )
+                    rmtree(path, ignore_errors=True)
+
+
+def _build_pg_dump(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    format_: _PGDumpFormat = "plain",
+    jobs: int | None = None,
+    schemas: MaybeListStr | None = None,
+    schemas_exc: MaybeListStr | None = None,
+    tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    tables_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    inserts: bool = False,
+    on_conflict_do_nothing: bool = False,
+    docker: str | None = None,
+) -> str:
+    database, host, port = _extract_url(url)
+    path = _path_pg_dump(path, format_=format_)
+    parts: list[str] = [
         "pg_dump",
         # general options
-        f"--dbname={url.database}",
+        f"--dbname={database}",
         f"--file={str(path)!r}",
         f"--format={format_}",
         "--verbose",
@@ -69,10 +122,10 @@ async def pg_dump(
         "--no-privileges",
         "--if-exists",
         # connection options
-        f"--host={url.host}",
-        f"--port={url.port}",
+        f"--host={host}",
+        f"--port={port}",
         "--no-password",
-    ])
+    ]
     if (format_ == "directory") and (jobs is not None):
         parts.append(f"--jobs={jobs}")
     if schemas is not None:
@@ -91,7 +144,61 @@ async def pg_dump(
         parts.append("--on-conflict-do-nothing")
     if url.username is not None:
         parts.append(f"--username={url.username}")
-    cmd = " ".join(parts)
+    if docker is not None:
+        parts = _wrap_docker(parts, docker)
+    return " ".join(parts)
+
+
+def _path_pg_dump(path: PathLike, /, *, format_: _PGDumpFormat = "plain") -> Path:
+    match format_:
+        case "plain":
+            suffix = ".sql"
+        case "custom":
+            suffix = ".pgdump"
+        case "directory":
+            suffix = None
+        case "tar":
+            suffix = ".tar"
+        case _ as never:
+            assert_never(never)
+    path = Path(path)
+    if suffix is not None:
+        path = ensure_suffix(path, suffix)
+    return path
+
+
+##
+
+
+async def restore(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    psql: bool = False,
+    database: str | None = None,
+    data_only: bool = False,
+    jobs: int | None = None,
+    schemas: MaybeListStr | None = None,
+    schemas_exc: MaybeListStr | None = None,
+    tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    docker: str | None = None,
+    dry_run: bool = False,
+    logger: LoggerOrName | None = None,
+) -> None:
+    """Run `pg_restore`/`psql`."""
+    cmd = _build_pg_restore_or_psql(
+        url,
+        path,
+        psql=psql,
+        database=database,
+        data_only=data_only,
+        jobs=jobs,
+        schemas=schemas,
+        schemas_exc=schemas_exc,
+        tables=tables,
+        docker=docker,
+    )
     if dry_run:
         if logger is not None:
             get_logger(logger=logger).info("Would run %r", str(cmd))
@@ -102,89 +209,75 @@ async def pg_dump(
         except KeyboardInterrupt:
             if logger is not None:
                 get_logger(logger=logger).info(
-                    "Cancelled backup to %r after %s", str(path), timer
+                    "Cancelled restore from %r after %s", str(path), timer
                 )
-            rmtree(path, ignore_errors=True)
         else:
             match output.return_code:
                 case 0:
                     if logger is not None:
                         get_logger(logger=logger).info(
-                            "Backup to %r finished after %s", str(path), timer
+                            "Restore from %r finished after %s", str(path), timer
                         )
                 case _:
                     if logger is not None:
                         get_logger(logger=logger).exception(
-                            "Backup to %r failed after %s\nstderr:\n%s",
+                            "Restore from %r failed after %s\nstderr:\n%s",
                             str(path),
                             timer,
                             output.stderr,
                         )
-                    rmtree(path, ignore_errors=True)
-
-
-@dataclass(kw_only=True, slots=True)
-class PGDumpError(Exception):
-    url: URL
-
-
-@dataclass(kw_only=True, slots=True)
-class _PGDumpDatabaseError(PGDumpError):
-    @override
-    def __str__(self) -> str:
-        return f"Expected URL to contain a 'database'; got {self.url}"
-
-
-@dataclass(kw_only=True, slots=True)
-class _PGDumpHostError(PGDumpError):
-    @override
-    def __str__(self) -> str:
-        return f"Expected URL to contain a 'host'; got {self.url}"
-
-
-@dataclass(kw_only=True, slots=True)
-class _PGDumpPortError(PGDumpError):
-    @override
-    def __str__(self) -> str:
-        return f"Expected URL to contain a 'port'; got {self.url}"
 
 
 ##
 
 
-async def pg_restore(
+def _build_pg_restore_or_psql(
     url: URL,
     path: PathLike,
     /,
     *,
+    psql: bool = False,
     database: str | None = None,
+    data_only: bool = False,
+    jobs: int | None = None,
+    schemas: MaybeListStr | None = None,
+    schemas_exc: MaybeListStr | None = None,
+    tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     docker: str | None = None,
+) -> str:
+    path = Path(path)
+    if (path.suffix == ".sql") or psql:
+        return _build_psql(url, path, database=database, docker=docker)
+    return _build_pg_restore(
+        url,
+        path,
+        database=database,
+        data_only=data_only,
+        jobs=jobs,
+        schemas=schemas,
+        schemas_exc=schemas_exc,
+        tables=tables,
+        docker=docker,
+    )
+
+
+def _build_pg_restore(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    database: str | None = None,
     data_only: bool = False,
     jobs: int | None = None,
     schemas: MaybeListStr | None = None,
     schemas_exc: MaybeListStr | None = None,
     tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
-    logger: LoggerOrName | None = None,
-    dry_run: bool = False,
-) -> None:
+    docker: str | None = None,
+) -> str:
     """Run `pg_restore`."""
-    match database, url.database:
-        case str() as database_use, _:
-            ...
-        case None, str() as database_use:
-            ...
-        case None, None:
-            raise _PGRestoreDatabaseError(url=url)
-        case _ as never:
-            assert_never(never)
-    if url.host is None:
-        raise _PGRestoreHostError(url=url)
-    if url.port is None:
-        raise _PGRestorePortError(url=url)
-    parts: list[str] = []
-    if docker is not None:
-        parts.extend(["docker", "exec", "-it", docker])
-    parts.extend([
+    url_database, host, port = _extract_url(url)
+    database_use = url_database if database is None else database
+    parts: list[str] = [
         "pg_restore",
         # general options
         f"--dbname={database_use}",
@@ -194,10 +287,10 @@ async def pg_restore(
         "--no-owner",
         "--no-privileges",
         # connection options
-        f"--host={url.host}",
-        f"--port={url.port}",
+        f"--host={host}",
+        f"--port={port}",
         "--no-password",
-    ])
+    ]
     if data_only:
         parts.append("--data-only")
     else:
@@ -212,68 +305,80 @@ async def 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 docker is not None:
+        parts = _wrap_docker(parts, docker)
     parts.append(str(path))
-    cmd = " ".join(parts)
-    if dry_run:
-        if logger is not None:
-            get_logger(logger=logger).info("Would run %r", str(cmd))
-        return
-    with temp_environ(PGPASSWORD=url.password), Timer() as timer:  # pragma: no cover
-        try:
-            output = await stream_command(cmd)
-        except KeyboardInterrupt:
-            if logger is not None:
-                get_logger(logger=logger).info(
-                    "Cancelled restore from %r after %s", str(path), timer
-                )
-        else:
-            match output.return_code:
-                case 0:
-                    if logger is not None:
-                        get_logger(logger=logger).info(
-                            "Restore from %r finished after %s", str(path), timer
-                        )
-                case _:
-                    if logger is not None:
-                        get_logger(logger=logger).exception(
-                            "Restore from %r failed after %s\nstderr:\n%s",
-                            str(path),
-                            timer,
-                            output.stderr,
-                        )
+    return " ".join(parts)
+
+
+def _build_psql(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    database: str | None = None,
+    docker: str | None = None,
+) -> str:
+    """Run `psql`."""
+    url_database, host, port = _extract_url(url)
+    database_use = url_database if database is None else database
+    parts: list[str] = [
+        "psql",
+        # general options
+        f"--dbname={database_use}",
+        f"--file={str(path)!r}",
+        # connection options
+        f"--host={host}",
+        f"--port={port}",
+        "--no-password",
+    ]
+    if url.username is not None:
+        parts.append(f"--username={url.username}")
+    if docker is not None:
+        parts = _wrap_docker(parts, docker)
+    return " ".join(parts)
+
+
+##
+
+
+def _extract_url(url: URL, /) -> tuple[str, str, int]:
+    if url.database is None:
+        raise _ExtractURLDatabaseError(url=url)
+    if url.host is None:
+        raise _ExtractURLHostError(url=url)
+    if url.port is None:
+        raise _ExtractURLPortError(url=url)
+    return url.database, url.host, url.port
 
 
 @dataclass(kw_only=True, slots=True)
-class PGRestoreError(Exception):
+class ExtractURLError(Exception):
     url: URL
 
 
 @dataclass(kw_only=True, slots=True)
-class _PGRestoreDatabaseError(PGRestoreError):
+class _ExtractURLDatabaseError(ExtractURLError):
     @override
     def __str__(self) -> str:
         return f"Expected URL to contain a 'database'; got {self.url}"
 
 
 @dataclass(kw_only=True, slots=True)
-class _PGRestoreHostError(PGRestoreError):
+class _ExtractURLHostError(ExtractURLError):
     @override
     def __str__(self) -> str:
         return f"Expected URL to contain a 'host'; got {self.url}"
 
 
 @dataclass(kw_only=True, slots=True)
-class _PGRestorePortError(PGRestoreError):
+class _ExtractURLPortError(ExtractURLError):
     @override
     def __str__(self) -> str:
         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)
@@ -283,4 +388,8 @@ def _get_table_name(obj: TableOrORMInstOrClass | str, /) -> str:
             assert_never(never)
 
 
-__all__ = ["PGDumpError", "PGRestoreError", "pg_dump", "pg_restore"]
+def _wrap_docker(parts: list[str], container: str, /) -> list[str]:
+    return ["docker", "exec", "-it", container, *parts]
+
+
+__all__ = ["ExtractURLError", "pg_dump", "restore"]