dycw-utilities 0.148.4__py3-none-any.whl → 0.174.12__py3-none-any.whl

utilities/polars_ols.py

@@ -6,8 +6,8 @@ from polars import Expr, Series, struct
 from polars_ols import RollingKwargs, compute_rolling_least_squares
 
 from utilities.errors import ImpossibleCaseError
-from utilities.functions import is_sequence_of
 from utilities.polars import concat_series, ensure_expr_or_series
+from utilities.typing import is_sequence_of
 
 if TYPE_CHECKING:
     from polars._typing import IntoExprColumn

utilities/postgres.py

@@ -9,10 +9,12 @@ from sqlalchemy import Table
 from sqlalchemy.orm import DeclarativeBase
 
 from utilities.asyncio import stream_command
+from utilities.docker import docker_exec_cmd
 from utilities.iterables import always_iterable
-from utilities.logging import get_logger
+from utilities.logging import to_logger
 from utilities.os import temp_environ
-from utilities.sqlalchemy import get_table_name
+from utilities.pathlib import ensure_suffix
+from utilities.sqlalchemy import extract_url, get_table_name
 from utilities.timer import Timer
 from utilities.types import PathLike
 
@@ -20,7 +22,12 @@ if TYPE_CHECKING:
     from sqlalchemy import URL
 
     from utilities.sqlalchemy import TableOrORMInstOrClass
-    from utilities.types import LoggerOrName, MaybeListStr, MaybeSequence, PathLike
+    from utilities.types import (
+        LoggerLike,
+        MaybeCollection,
+        MaybeCollectionStr,
+        PathLike,
+    )
 
 
 type _PGDumpFormat = Literal["plain", "custom", "directory", "tar"]
@@ -31,174 +38,302 @@ async def pg_dump(
     path: PathLike,
     /,
     *,
-    docker: str | None = None,
+    docker_container: str | None = None,
     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,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    extension: MaybeCollectionStr | None = None,
+    extension_exc: MaybeCollectionStr | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    table_exc: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     inserts: bool = False,
-    logger: LoggerOrName | None = None,
+    on_conflict_do_nothing: bool = False,
+    role: str | None = None,
     dry_run: bool = False,
-) -> None:
+    logger: LoggerLike | None = None,
+) -> bool:
     """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])
+    cmd = _build_pg_dump(
+        url,
+        path,
+        docker_container=docker_container,
+        format_=format_,
+        jobs=jobs,
+        data_only=data_only,
+        clean=clean,
+        create=create,
+        extension=extension,
+        extension_exc=extension_exc,
+        schema=schema,
+        schema_exc=schema_exc,
+        table=table,
+        table_exc=table_exc,
+        inserts=inserts,
+        on_conflict_do_nothing=on_conflict_do_nothing,
+        role=role,
+    )
+    if dry_run:
+        if logger is not None:
+            to_logger(logger).info("Would run:\n\t%r", str(cmd))
+        return True
+    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:
+                to_logger(logger).info(
+                    "Cancelled backup to %r after %s", str(path), timer
+                )
+            rmtree(path, ignore_errors=True)
+            return False
+        if output.return_code != 0:
+            if logger is not None:
+                to_logger(logger).exception(
+                    "Backup to %r failed after %s\nstderr:\n%s",
+                    str(path),
+                    timer,
+                    output.stderr,
+                )
+            rmtree(path, ignore_errors=True)
+            return False
+    if logger is not None:  # pragma: no cover
+        to_logger(logger).info("Backup to %r finished after %s", str(path), timer)
+    return True  # pragma: no cover
+
+
+def _build_pg_dump(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    docker_container: str | None = None,
+    format_: _PGDumpFormat = "plain",
+    jobs: int | None = None,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    extension: MaybeCollectionStr | None = None,
+    extension_exc: MaybeCollectionStr | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    table_exc: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    inserts: bool = False,
+    on_conflict_do_nothing: bool = False,
+    role: str | None = None,
+) -> str:
+    extracted = extract_url(url)
+    path = _path_pg_dump(path, format_=format_)
+    parts: list[str] = ["pg_dump"]
+    if docker_container is not None:
+        parts = docker_exec_cmd(docker_container, *parts, PGPASSWORD=extracted.password)
     parts.extend([
-        "pg_dump",
         # general options
-        f"--dbname={url.database}",
         f"--file={str(path)!r}",
         f"--format={format_}",
         "--verbose",
         # output options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
         "--large-objects",
-        "--clean",
         "--no-owner",
         "--no-privileges",
-        "--if-exists",
         # connection options
-        f"--host={url.host}",
-        f"--port={url.port}",
+        f"--dbname={extracted.database}",
+        f"--host={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
         "--no-password",
     ])
     if (format_ == "directory") and (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)])
-    if tables_exc is not None:
+    if create:
+        parts.append("--create")
+    if extension is not None:
+        parts.extend([f"--extension={e}" for e in always_iterable(extension)])
+    if extension_exc is not None:
+        parts.extend([
+            f"--exclude-extension={e}" for e in always_iterable(extension_exc)
+        ])
+    if schema is not None:
+        parts.extend([f"--schema={s}" for s in always_iterable(schema)])
+    if schema_exc is not None:
+        parts.extend([f"--exclude-schema={s}" for s in always_iterable(schema_exc)])
+    if table is not None:
+        parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(table)])
+    if table_exc is not None:
         parts.extend([
-            f"--exclude-table={_get_table_name(t)}" for t in always_iterable(tables_exc)
+            f"--exclude-table={_get_table_name(t)}" for t in always_iterable(table_exc)
         ])
     if inserts:
         parts.append("--inserts")
-    if url.username is not None:
-        parts.append(f"--username={url.username}")
-    cmd = " ".join(parts)
+    if on_conflict_do_nothing:
+        parts.append("--on-conflict-do-nothing")
+    if role is not None:
+        parts.append(f"--role={role}")
+    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 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,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    jobs: int | None = None,
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
+    docker_container: str | None = None,
+    dry_run: bool = False,
+    logger: LoggerLike | None = None,
+) -> bool:
+    """Run `pg_restore`/`psql`."""
+    cmd = _build_pg_restore_or_psql(
+        url,
+        path,
+        psql=psql,
+        data_only=data_only,
+        clean=clean,
+        create=create,
+        jobs=jobs,
+        schema=schema,
+        schema_exc=schema_exc,
+        table=table,
+        role=role,
+        docker_container=docker_container,
+    )
     if dry_run:
         if logger is not None:
-            get_logger(logger=logger).info("Would run %r", str(cmd))
-        return
+            to_logger(logger).info("Would run:\n\t%r", str(cmd))
+        return True
     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
+                to_logger(logger).info(
+                    "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
-                        )
-                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)
-
-
-@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}"
+            return False
+        if output.return_code != 0:
+            if logger is not None:
+                to_logger(logger).exception(
+                    "Restore from %r failed after %s\nstderr:\n%s",
+                    str(path),
+                    timer,
+                    output.stderr,
+                )
+            return False
+    if logger is not None:  # pragma: no cover
+        to_logger(logger).info("Restore from %r finished after %s", str(path), timer)
+    return True  # pragma: no cover
 
 
 ##
 
 
-async def pg_restore(
+def _build_pg_restore_or_psql(
     url: URL,
     path: PathLike,
     /,
     *,
-    database: str | None = None,
-    docker: str | None = None,
+    psql: bool = False,
     data_only: bool = False,
+    clean: bool = False,
+    create: 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:
+    schema: MaybeCollectionStr | None = None,
+    schema_exc: MaybeCollectionStr | None = None,
+    table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
+    docker_container: str | None = None,
+) -> str:
+    path = Path(path)
+    if (path.suffix == ".sql") or psql:
+        return _build_psql(url, path, docker_container=docker_container)
+    return _build_pg_restore(
+        url,
+        path,
+        data_only=data_only,
+        clean=clean,
+        create=create,
+        jobs=jobs,
+        schemas=schema,
+        schemas_exc=schema_exc,
+        tables=table,
+        role=role,
+        docker_container=docker_container,
+    )
+
+
+def _build_pg_restore(
+    url: URL,
+    path: PathLike,
+    /,
+    *,
+    data_only: bool = False,
+    clean: bool = False,
+    create: bool = False,
+    jobs: int | None = None,
+    schemas: MaybeCollectionStr | None = None,
+    schemas_exc: MaybeCollectionStr | None = None,
+    tables: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
+    docker_container: 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])
+    extracted = extract_url(url)
+    parts: list[str] = ["pg_restore"]
+    if docker_container is not None:
+        parts = docker_exec_cmd(docker_container, *parts, PGPASSWORD=extracted.password)
     parts.extend([
-        "pg_restore",
         # general options
-        f"--dbname={database_use}",
         "--verbose",
         # restore options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
         "--exit-on-error",
         "--no-owner",
         "--no-privileges",
         # connection options
-        f"--host={url.host}",
-        f"--port={url.port}",
+        f"--host={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
+        f"--dbname={extracted.database}",
         "--no-password",
     ])
-    if data_only:
-        parts.append("--data-only")
-    else:
-        parts.extend(["--clean", "--if-exists"])
+    if create:
+        parts.append("--create")
     if jobs is not None:
         parts.append(f"--jobs={jobs}")
     if schemas is not None:
@@ -207,77 +342,67 @@ async def pg_restore(
         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)])
-    if url.username is not None:
-        parts.append(f"--username={url.username}")
+    if role is not None:
+        parts.append(f"--role={role}")
     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)
 
 
-@dataclass(kw_only=True, slots=True)
-class PGRestoreError(Exception):
-    url: URL
-
-
-@dataclass(kw_only=True, slots=True)
-class _PGRestoreDatabaseError(PGRestoreError):
-    @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):
-    @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):
-    @override
-    def __str__(self) -> str:
-        return f"Expected URL to contain a 'port'; got {self.url}"
+def _build_psql(
+    url: URL, path: PathLike, /, *, docker_container: str | None = None
+) -> str:
+    """Run `psql`."""
+    extracted = extract_url(url)
+    parts: list[str] = ["psql"]
+    if docker_container is not None:
+        parts = docker_exec_cmd(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
+        # general options
+        f"--dbname={extracted.database}",
+        f"--file={str(path)!r}",
+        # connection options
+        f"--host={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
+        "--no-password",
+    ])
+    return " ".join(parts)
 
 
 ##
 
 
 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:
+        case never:
+            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 never:
             assert_never(never)
 
 
-__all__ = ["PGDumpError", "PGRestoreError", "pg_dump", "pg_restore"]
+@dataclass(kw_only=True, slots=True)
+class _ResolveDataOnlyAndCleanError(Exception):
+    @override
+    def __str__(self) -> str:
+        return "Cannot use '--data-only' and '--clean' together"
+
+
+__all__ = ["pg_dump", "restore"]