dycw-utilities 0.129.10__py3-none-any.whl → 0.175.17__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

@@ -0,0 +1,408 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+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.docker import docker_exec_cmd
+from utilities.iterables import always_iterable
+from utilities.logging import to_logger
+from utilities.os import temp_environ
+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
+
+if TYPE_CHECKING:
+    from sqlalchemy import URL
+
+    from utilities.sqlalchemy import TableOrORMInstOrClass
+    from utilities.types import (
+        LoggerLike,
+        MaybeCollection,
+        MaybeCollectionStr,
+        PathLike,
+    )
+
+
+type _PGDumpFormat = Literal["plain", "custom", "directory", "tar"]
+
+
+async def 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,
+    dry_run: bool = False,
+    logger: LoggerLike | None = None,
+) -> bool:
+    """Run `pg_dump`."""
+    path = _path_pg_dump(path, format_=format_)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    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([
+        # general options
+        f"--file={str(path)!r}",
+        f"--format={format_}",
+        "--verbose",
+        # output options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
+        "--large-objects",
+        "--no-owner",
+        "--no-privileges",
+        # connection options
+        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 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(table_exc)
+        ])
+    if inserts:
+        parts.append("--inserts")
+    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:
+            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 restore from %r after %s", str(path), timer
+                )
+            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
+
+
+##
+
+
+def _build_pg_restore_or_psql(
+    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,
+) -> 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`."""
+    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([
+        # general options
+        "--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={extracted.host}",
+        f"--port={extracted.port}",
+        f"--username={extracted.username}",
+        f"--dbname={extracted.database}",
+        "--no-password",
+    ])
+    if create:
+        parts.append("--create")
+    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)])
+    if role is not None:
+        parts.append(f"--role={role}")
+    parts.append(str(path))
+    return " ".join(parts)
+
+
+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:
+    match obj:
+        case Table() | DeclarativeBase() | type() as table_or_orm:
+            return get_table_name(table_or_orm)
+        case str() as name:
+            return name
+        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)
+
+
+@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"]

utilities/pottery.py

@@ -1,50 +1,137 @@
 from __future__ import annotations
 
-import asyncio
-from contextlib import asynccontextmanager, suppress
-from typing import TYPE_CHECKING
+from contextlib import suppress
+from dataclasses import dataclass
+from sys import maxsize
+from typing import TYPE_CHECKING, override
 
 from pottery import AIORedlock
 from pottery.exceptions import ReleaseUnlockedLock
 from redis.asyncio import Redis
 
-from utilities.datetime import MILLISECOND, SECOND, datetime_duration_to_float
+from utilities.asyncio import sleep_td, timeout_td
+from utilities.contextlib import enhanced_async_context_manager
 from utilities.iterables import always_iterable
+from utilities.whenever import MILLISECOND, SECOND, to_nanoseconds
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator
+    from collections.abc import AsyncIterator, Iterable
 
-    from utilities.types import Duration, MaybeIterable
+    from whenever import Delta
 
+    from utilities.types import MaybeIterable
 
-@asynccontextmanager
-async def yield_locked_resource(
+_NUM: int = 1
+_TIMEOUT_RELEASE: Delta = 10 * SECOND
+_SLEEP: Delta = MILLISECOND
+
+
+##
+
+
+async def extend_lock(
+    *, lock: AIORedlock | None = None, raise_on_redis_errors: bool | None = None
+) -> None:
+    """Extend a lock."""
+    if lock is not None:
+        await lock.extend(raise_on_redis_errors=raise_on_redis_errors)
+
+
+##
+
+
+@enhanced_async_context_manager
+async def yield_access(
     redis: MaybeIterable[Redis],
     key: str,
     /,
     *,
-    duration: Duration = 10 * SECOND,
-    sleep: Duration = MILLISECOND,
-) -> AsyncIterator[None]:
-    """Yield a locked resource."""
+    num: int = _NUM,
+    timeout_release: Delta = _TIMEOUT_RELEASE,
+    num_extensions: int | None = None,
+    timeout_acquire: Delta | None = None,
+    sleep: Delta = _SLEEP,
+    throttle: Delta | None = None,
+) -> AsyncIterator[AIORedlock]:
+    """Acquire access to a locked resource."""
+    if num <= 0:
+        raise _YieldAccessNumLocksError(key=key, num=num)
     masters = (  # skipif-ci-and-not-linux
         {redis} if isinstance(redis, Redis) else set(always_iterable(redis))
     )
-    duration_use = datetime_duration_to_float(duration)  # skipif-ci-and-not-linux
-    lock = AIORedlock(  # skipif-ci-and-not-linux
-        key=key,
-        masters=masters,
-        auto_release_time=duration_use,
-        context_manager_timeout=duration_use,
-    )
-    sleep_use = datetime_duration_to_float(sleep)  # skipif-ci-and-not-linux
-    while not await lock.acquire():  # pragma: no cover
-        _ = await asyncio.sleep(sleep_use)
+    locks = [  # skipif-ci-and-not-linux
+        AIORedlock(
+            key=f"{key}_{i}_of_{num}",
+            masters=masters,
+            auto_release_time=to_nanoseconds(timeout_release) / 1e9,
+            num_extensions=maxsize if num_extensions is None else num_extensions,
+        )
+        for i in range(1, num + 1)
+    ]
+    lock: AIORedlock | None = None  # skipif-ci-and-not-linux
     try:  # skipif-ci-and-not-linux
-        yield
+        lock = await _get_first_available_lock(
+            key, locks, num=num, timeout=timeout_acquire, sleep=sleep
+        )
+        yield lock
     finally:  # skipif-ci-and-not-linux
-        with suppress(ReleaseUnlockedLock):
-            await lock.release()
+        await sleep_td(throttle)
+        if lock is not None:
+            with suppress(ReleaseUnlockedLock):
+                await lock.release()
+
+
+async def _get_first_available_lock(
+    key: str,
+    locks: Iterable[AIORedlock],
+    /,
+    *,
+    num: int = _NUM,
+    timeout: Delta | None = None,
+    sleep: Delta | None = _SLEEP,
+) -> AIORedlock:
+    locks = list(locks)  # skipif-ci-and-not-linux
+    error = _YieldAccessUnableToAcquireLockError(  # skipif-ci-and-not-linux
+        key=key, num=num, timeout=timeout
+    )
+    async with timeout_td(timeout, error=error):  # skipif-ci-and-not-linux
+        while True:
+            if (result := await _get_first_available_lock_if_any(locks)) is not None:
+                return result
+            await sleep_td(sleep)
+
+
+async def _get_first_available_lock_if_any(
+    locks: Iterable[AIORedlock], /
+) -> AIORedlock | None:
+    for lock in locks:  # skipif-ci-and-not-linux
+        if await lock.acquire(blocking=False):
+            return lock
+    return None  # skipif-ci-and-not-linux
+
+
+@dataclass(kw_only=True, slots=True)
+class YieldAccessError(Exception):
+    key: str
+
+
+@dataclass(kw_only=True, slots=True)
+class _YieldAccessNumLocksError(YieldAccessError):
+    num: int
+
+    @override
+    def __str__(self) -> str:
+        return f"Number of locks for {self.key!r} must be positive; got {self.num}"
+
+
+@dataclass(kw_only=True, slots=True)
+class _YieldAccessUnableToAcquireLockError(YieldAccessError):
+    num: int
+    timeout: Delta | None
+
+    @override
+    def __str__(self) -> str:
+        return f"Unable to acquire any 1 of {self.num} locks for {self.key!r} after {self.timeout}"  # skipif-ci-and-not-linux
 
 
-__all__ = ["yield_locked_resource"]
+__all__ = ["YieldAccessError", "extend_lock", "yield_access"]

utilities/pqdm.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, assert_never
+from typing import TYPE_CHECKING, Any, Literal, assert_never
 
 from pqdm import processes, threads
 from tqdm.auto import tqdm as tqdm_auto
@@ -9,7 +9,7 @@ from tqdm.auto import tqdm as tqdm_auto
 from utilities.functions import get_func_name
 from utilities.iterables import apply_to_varargs
 from utilities.os import get_cpu_use
-from utilities.sentinel import Sentinel, sentinel
+from utilities.sentinel import Sentinel, is_sentinel, sentinel
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable
@@ -20,12 +20,11 @@ if TYPE_CHECKING:
     from utilities.types import Parallelism
 
 
-_T = TypeVar("_T")
 type _ExceptionBehaviour = Literal["ignore", "immediate", "deferred"]
 
 
-def pqdm_map(
-    func: Callable[..., _T],
+def pqdm_map[T](
+    func: Callable[..., T],
     /,
     *iterables: Iterable[Any],
     parallelism: Parallelism = "processes",
@@ -35,7 +34,7 @@ def pqdm_map(
     tqdm_class: tqdm_type = tqdm_auto,  # pyright: ignore[reportArgumentType]
     desc: str | None | Sentinel = sentinel,
     **kwargs: Any,
-) -> list[_T]:
+) -> list[T]:
     """Parallel map, powered by `pqdm`."""
     return pqdm_starmap(
         func,
@@ -50,8 +49,8 @@ def pqdm_map(
     )
 
 
-def pqdm_starmap(
-    func: Callable[..., _T],
+def pqdm_starmap[T](
+    func: Callable[..., T],
     iterable: Iterable[tuple[Any, ...]],
     /,
     *,
@@ -62,7 +61,7 @@ def pqdm_starmap(
     tqdm_class: tqdm_type = tqdm_auto,  # pyright: ignore[reportArgumentType]
     desc: str | None | Sentinel = sentinel,
     **kwargs: Any,
-) -> list[_T]:
+) -> list[T]:
     """Parallel starmap, powered by `pqdm`."""
     apply = partial(apply_to_varargs, func)
     n_jobs_use = get_cpu_use(n=n_jobs)
@@ -91,7 +90,7 @@ def pqdm_starmap(
                 **_get_desc(desc, func),
                 **kwargs,
             )
-        case _ as never:
+        case never:
             assert_never(never)
     return list(result)
 
@@ -99,7 +98,7 @@ def pqdm_starmap(
 def _get_desc(
     desc: str | None | Sentinel, func: Callable[..., Any], /
 ) -> dict[str, str]:
-    desc_use = get_func_name(func) if isinstance(desc, Sentinel) else desc
+    desc_use = get_func_name(func) if is_sentinel(desc) else desc
     return {} if desc_use is None else {"desc": desc_use}