dycw-utilities 0.148.5__py3-none-any.whl → 0.149.1__py3-none-any.whl

{dycw_utilities-0.148.5.dist-info → dycw_utilities-0.149.1.dist-info}/METADATA

@@ -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.dist-info → dycw_utilities-0.149.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=LCky1VazOaGgD8qf6DVgou1be6BJARw1j7YB_Kwfbgw,60
+utilities/__init__.py,sha256=JP1_vDjx7triKVpffufyRF4tTLDTrHerAtlEopa1IYA,60
 utilities/altair.py,sha256=92E2lCdyHY4Zb-vCw6rEJIsWdKipuu-Tu2ab1ufUfAk,9079
 utilities/asyncio.py,sha256=z0w3fb-U5Ml5YXVaFFPClizXaQmjDO6YgZg-V9QL0VQ,16021
 utilities/atomicwrites.py,sha256=xcOWenTBRS0oat3kg7Sqe51AohNThMQ2ixPL7QCG8hw,5795
@@ -48,7 +48,7 @@ utilities/pickle.py,sha256=MBT2xZCsv0pH868IXLGKnlcqNx2IRVKYNpRcqiQQqxw,653
 utilities/platform.py,sha256=Ue9LSxYvg9yUXGKuz5aZoy_qkUEXde-v6B09exgSctU,2813
 utilities/polars.py,sha256=BgiDryAVOapi41ddfJqN0wYh_sDj8BNEYtPB36LaHdo,71824
 utilities/polars_ols.py,sha256=Uc9V5kvlWZ5cU93lKZ-cfAKdVFFw81tqwLW9PxtUvMs,5618
-utilities/postgres.py,sha256=zQx4rNMg65znDdcj_NC1ZPTtW-XkgSorlTe9CTSmrHk,9115
+utilities/postgres.py,sha256=rMJ2Bp5IzPz365qKy0r80usa7p_3F0bhNjgFDUHA8CY,11847
 utilities/pottery.py,sha256=w2X80PXWwzdHdqSYJP6ESrPNNDP3xzpyuJn-fp-Vt3M,5969
 utilities/pqdm.py,sha256=BTsYPtbKQWwX-iXF4qCkfPG7DPxIB54J989n83bXrIo,3092
 utilities/psutil.py,sha256=KUlu4lrUw9Zg1V7ZGetpWpGb9DB8l_SSDWGbANFNCPU,2104
@@ -88,8 +88,8 @@ utilities/zoneinfo.py,sha256=oEH-nL3t4h9uawyZqWDtNtDAl6M-CLpLYGI_nI6DulM,1971
 utilities/pytest_plugins/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 utilities/pytest_plugins/pytest_randomly.py,sha256=NXzCcGKbpgYouz5yehKb4jmxmi2SexKKpgF4M65bi10,414
 utilities/pytest_plugins/pytest_regressions.py,sha256=Iwhfv_OJH7UCPZCfoh7ugZ2Xjqjil-BBBsOb8sDwiGI,1471
-dycw_utilities-0.148.5.dist-info/METADATA,sha256=zKtIyZ---2ibI4hD6IeeSvJG2Gsx8CL31o-o1PxcWME,1697
-dycw_utilities-0.148.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.148.5.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.148.5.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.148.5.dist-info/RECORD,,
+dycw_utilities-0.149.1.dist-info/METADATA,sha256=4gXdqC0JM7T2UI9VWjQkJxq3-eqcaDG3SAEGsoRO01c,1697
+dycw_utilities-0.149.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.149.1.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.149.1.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.149.1.dist-info/RECORD,,

utilities/__init__.py

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

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"]