dycw-utilities 0.166.37__py3-none-any.whl → 0.167.1__py3-none-any.whl

{dycw_utilities-0.166.37.dist-info → dycw_utilities-0.167.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.166.37
+Version: 0.167.1
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12
@@ -12,7 +12,7 @@ Provides-Extra: logging
 Requires-Dist: coloredlogs<15.1,>=15.0.1; extra == 'logging'
 Provides-Extra: test
 Requires-Dist: dycw-pytest-only<2.2,>=2.1.1; extra == 'test'
-Requires-Dist: hypothesis<6.140,>=6.139.1; extra == 'test'
+Requires-Dist: hypothesis<6.140,>=6.139.2; extra == 'test'
 Requires-Dist: pytest-asyncio<1.3,>=1.2.0; extra == 'test'
 Requires-Dist: pytest-cov<7.1,>=7.0.0; extra == 'test'
 Requires-Dist: pytest-instafail<0.6,>=0.5.0; extra == 'test'

{dycw_utilities-0.166.37.dist-info → dycw_utilities-0.167.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=5A2yZ5q_GCT119LdZKfbPBpyhrtF0D3ytDUhvCLLd1s,61
+utilities/__init__.py,sha256=5n59LyIEvqGZR7Atiaj1WopkM0tKCDhQyUb2l-MH_qc,60
 utilities/aeventkit.py,sha256=ddoleSwW9zdc2tjX5Ge0pMKtYwV_JMxhHYOxnWX2AGM,12609
 utilities/altair.py,sha256=nHdpWt8ZwdUwRQN970MvHd5bRWokNqzHcZQEdSHKRuE,9033
 utilities/asyncio.py,sha256=60l1IwjnRGeaVphAFiwDIHyfKoZYKY-XGpptUxGiU-M,17034
@@ -12,6 +12,7 @@ utilities/contextvars.py,sha256=J8OhC7jqozAGYOCe2KUWysbPXNGe5JYz3HfaY_mIs08,883
 utilities/cryptography.py,sha256=5PFrzsNUGHay91dFgYnDKwYprXxahrBqztmUqViRzBk,956
 utilities/cvxpy.py,sha256=Rv1-fD-XYerosCavRF8Pohop2DBkU3AlFaGTfD8AEAA,13776
 utilities/dataclasses.py,sha256=xbU3QN1GFy7RC6hIJRZIeUZm7YRlodrgEWmahWG6k2g,32465
+utilities/docker.py,sha256=AahXNwtu4ojdVeTdAa-VBR8h7GWNAE3ChHeVBa0oZAE,603
 utilities/enum.py,sha256=5l6pwZD1cjSlVW4ss-zBPspWvrbrYrdtJWcg6f5_J5w,5781
 utilities/errors.py,sha256=mFlDGSM0LI1jZ1pbqwLAH3ttLZ2JVIxyZLojw8tGVZU,1479
 utilities/fastapi.py,sha256=TqyKvBjiMS594sXPjrz-KRTLMb3l3D3rZ1zAYV7GfOk,1454
@@ -49,7 +50,7 @@ utilities/pickle.py,sha256=MBT2xZCsv0pH868IXLGKnlcqNx2IRVKYNpRcqiQQqxw,653
 utilities/platform.py,sha256=pTn7gw6N4T6LdKrf0virwarof_mze9WtoQlrGMzhGVI,2798
 utilities/polars.py,sha256=qsiYY9p_41fORGnc7HNkA4zhlsycK7sgD74xuigMDAc,87466
 utilities/polars_ols.py,sha256=LNTFNLPuYW7fcAHymlbnams_DhitToblYvib3mhKbwI,5615
-utilities/postgres.py,sha256=ynCTTaF-bVEOSW-KEAR-dlLh_hYjeVVjm__-4pEU8Zk,12269
+utilities/postgres.py,sha256=qQzXJngzs2jyZ4rkzL6uBdhXPwiqpPtYQiIonNcGIwI,12516
 utilities/pottery.py,sha256=nA0SsF9irvfC0tk68YAr08tuL9lGRSlBKihSx7Ibk84,3963
 utilities/pqdm.py,sha256=idv2seRVP2f6NeSfpeEnT5A-tQezaHZKDyeu16g2-0E,3091
 utilities/psutil.py,sha256=KUlu4lrUw9Zg1V7ZGetpWpGb9DB8l_SSDWGbANFNCPU,2104
@@ -68,7 +69,7 @@ utilities/sentinel.py,sha256=A_p5jX2K0Yc5XBfoYHyBLqHsEWzE1ByOdDuzzA2pZnE,1434
 utilities/shelve.py,sha256=4OzjQI6kGuUbJciqf535rwnao-_IBv66gsT6tRGiUt0,759
 utilities/slack_sdk.py,sha256=76-DYtcGiUhEvl-voMamc5OjfF7Y7nCq54Bys1arqzw,2233
 utilities/socket.py,sha256=K77vfREvzoVTrpYKo6MZakol0EYu2q1sWJnnZqL0So0,118
-utilities/sqlalchemy.py,sha256=X_F4Vq3t0ftgmhYdrs1DOMrI4ls5Tw_ddcW7dmhQaPY,36407
+utilities/sqlalchemy.py,sha256=HQYpd7LFxdTF5WYVWYtCJeEBI71EJm7ytvCGyAH9B-U,37163
 utilities/sqlalchemy_polars.py,sha256=JCGhB37raSR7fqeWV5dTsciRTMVzIdVT9YSqKT0piT0,13370
 utilities/statsmodels.py,sha256=koyiBHvpMcSiBfh99wFUfSggLNx7cuAw3rwyfAhoKpQ,3410
 utilities/string.py,sha256=shmBK87zZwzGyixuNuXCiUbqzfeZ9xlrFwz6JTaRvDk,582
@@ -92,8 +93,8 @@ utilities/zoneinfo.py,sha256=tdIScrTB2-B-LH0ukb1HUXKooLknOfJNwHk10MuMYvA,3619
 utilities/pytest_plugins/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 utilities/pytest_plugins/pytest_randomly.py,sha256=B1qYVlExGOxTywq2r1SMi5o7btHLk2PNdY_b1p98dkE,409
 utilities/pytest_plugins/pytest_regressions.py,sha256=mnHYBfdprz50UGVkVzV1bZERZN5CFfoF8YbokGxdFwU,1639
-dycw_utilities-0.166.37.dist-info/METADATA,sha256=rz6vDyAduIOHCHNE2PwjwCjAi6Oj886M54iwgmo7Bxk,1699
-dycw_utilities-0.166.37.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.166.37.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.166.37.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.166.37.dist-info/RECORD,,
+dycw_utilities-0.167.1.dist-info/METADATA,sha256=YDvpCOu6_rWr8_UnnjOzy0UM4JQABa4nltfz7kkNkLU,1698
+dycw_utilities-0.167.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.167.1.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.167.1.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.167.1.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.166.37"
+__version__ = "0.167.1"

utilities/docker.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+def docker_exec(
+    container: str,
+    /,
+    *cmd: str,
+    env: Mapping[str, str] | None = None,
+    **env_kwargs: str | None,
+) -> list[str]:
+    """Run a command through `docker exec`."""
+    full = ["docker", "exec"]
+    mapping: dict[str, str | None] = ({} if env is None else dict(env)) | env_kwargs
+    for key, value in mapping.items():
+        full.append(f"--env={key}={value}")
+    return [*full, "--interactive", container, *cmd]
+
+
+__all__ = ["docker_exec"]

utilities/postgres.py

@@ -9,6 +9,7 @@ from sqlalchemy import Table
 from sqlalchemy.orm import DeclarativeBase
 
 from utilities.asyncio import stream_command
+from utilities.docker import docker_exec
 from utilities.iterables import always_iterable
 from utilities.logging import to_logger
 from utilities.os import temp_environ
@@ -37,6 +38,7 @@ async def pg_dump(
     path: PathLike,
     /,
     *,
+    docker_container: str | None = None,
     format_: _PGDumpFormat = "plain",
     jobs: int | None = None,
     data_only: bool = False,
@@ -51,7 +53,6 @@ async def pg_dump(
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
     role: str | None = None,
-    docker: str | None = None,
     dry_run: bool = False,
     logger: LoggerLike | None = None,
 ) -> bool:
@@ -61,6 +62,7 @@ async def pg_dump(
     cmd = _build_pg_dump(
         url,
         path,
+        docker_container=docker_container,
         format_=format_,
         jobs=jobs,
         data_only=data_only,
@@ -75,7 +77,6 @@ async def pg_dump(
         inserts=inserts,
         on_conflict_do_nothing=on_conflict_do_nothing,
         role=role,
-        docker=docker,
     )
     if dry_run:
         if logger is not None:
@@ -111,6 +112,7 @@ def _build_pg_dump(
     path: PathLike,
     /,
     *,
+    docker_container: str | None = None,
     format_: _PGDumpFormat = "plain",
     jobs: int | None = None,
     data_only: bool = False,
@@ -125,12 +127,13 @@ def _build_pg_dump(
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
     role: str | None = None,
-    docker: str | None = None,
 ) -> str:
     extracted = extract_url(url)
     path = _path_pg_dump(path, format_=format_)
-    parts: list[str] = [
-        "pg_dump",
+    parts: list[str] = ["pg_dump"]
+    if docker_container is not None:
+        parts = docker_exec(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
         # general options
         f"--file={str(path)!r}",
         f"--format={format_}",
@@ -146,7 +149,7 @@ def _build_pg_dump(
         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:
@@ -173,8 +176,6 @@ def _build_pg_dump(
         parts.append("--on-conflict-do-nothing")
     if role is not None:
         parts.append(f"--role={role}")
-    if docker is not None:
-        parts = _wrap_docker(parts, docker)
     return " ".join(parts)
 
 
@@ -213,7 +214,7 @@ async def restore(
     schema_exc: MaybeCollectionStr | None = None,
     table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     role: str | None = None,
-    docker: str | None = None,
+    docker_container: str | None = None,
     dry_run: bool = False,
     logger: LoggerLike | None = None,
 ) -> bool:
@@ -230,7 +231,7 @@ async def restore(
         schema_exc=schema_exc,
         table=table,
         role=role,
-        docker=docker,
+        docker_container=docker_container,
     )
     if dry_run:
         if logger is not None:
@@ -276,11 +277,11 @@ def _build_pg_restore_or_psql(
     schema_exc: MaybeCollectionStr | None = None,
     table: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     role: str | None = None,
-    docker: str | None = None,
+    docker_container: str | None = None,
 ) -> str:
     path = Path(path)
     if (path.suffix == ".sql") or psql:
-        return _build_psql(url, path, docker=docker)
+        return _build_psql(url, path, docker_container=docker_container)
     return _build_pg_restore(
         url,
         path,
@@ -292,7 +293,7 @@ def _build_pg_restore_or_psql(
         schemas_exc=schema_exc,
         tables=table,
         role=role,
-        docker=docker,
+        docker_container=docker_container,
     )
 
 
@@ -309,12 +310,14 @@ def _build_pg_restore(
     schemas_exc: MaybeCollectionStr | None = None,
     tables: MaybeCollection[TableOrORMInstOrClass | str] | None = None,
     role: str | None = None,
-    docker: str | None = None,
+    docker_container: str | None = None,
 ) -> str:
     """Run `pg_restore`."""
     extracted = extract_url(url)
-    parts: list[str] = [
-        "pg_restore",
+    parts: list[str] = ["pg_restore"]
+    if docker_container is not None:
+        parts = docker_exec(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
         # general options
         "--verbose",
         # restore options
@@ -328,7 +331,7 @@ def _build_pg_restore(
         f"--username={extracted.username}",
         f"--dbname={extracted.database}",
         "--no-password",
-    ]
+    ])
     if create:
         parts.append("--create")
     if jobs is not None:
@@ -341,17 +344,19 @@ def _build_pg_restore(
         parts.extend([f"--table={_get_table_name(t)}" for t in always_iterable(tables)])
     if role is not None:
         parts.append(f"--role={role}")
-    if docker is not None:
-        parts = _wrap_docker(parts, docker)
     parts.append(str(path))
     return " ".join(parts)
 
 
-def _build_psql(url: URL, path: PathLike, /, *, docker: str | None = None) -> str:
+def _build_psql(
+    url: URL, path: PathLike, /, *, docker_container: str | None = None
+) -> str:
     """Run `psql`."""
     extracted = extract_url(url)
-    parts: list[str] = [
-        "psql",
+    parts: list[str] = ["psql"]
+    if docker_container is not None:
+        parts = docker_exec(docker_container, *parts, PGPASSWORD=extracted.password)
+    parts.extend([
         # general options
         f"--dbname={extracted.database}",
         f"--file={str(path)!r}",
@@ -360,9 +365,7 @@ def _build_psql(url: URL, path: PathLike, /, *, docker: str | None = None) -> st
         f"--port={extracted.port}",
         f"--username={extracted.username}",
         "--no-password",
-    ]
-    if docker is not None:
-        parts = _wrap_docker(parts, docker)
+    ])
     return " ".join(parts)
 
 
@@ -402,8 +405,4 @@ class _ResolveDataOnlyAndCleanError(Exception):
         return "Cannot use '--data-only' and '--clean' together"
 
 
-def _wrap_docker(parts: list[str], container: str, /) -> list[str]:
-    return ["docker", "exec", "-it", container, *parts]
-
-
 __all__ = ["pg_dump", "restore"]

utilities/sqlalchemy.py

@@ -331,6 +331,20 @@ async def ensure_database_dropped(super_: URL, database: str, /) -> None:
         _ = await conn.execute(text(f"DROP DATABASE IF EXISTS {database}"))
 
 
+async def ensure_database_users_disconnected(super_: URL, database: str, /) -> None:
+    """Ensure a databases' users are disconnected."""
+    engine = create_async_engine(super_, isolation_level="AUTOCOMMIT")
+    match dialect := _get_dialect(engine):
+        case "postgresql":  # skipif-ci-and-not-linux
+            query = f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = {database!r} AND pid <> pg_backend_pid()"  # noqa: S608
+        case "mssql" | "mysql" | "oracle" | "sqlite":  # pragma: no cover
+            raise NotImplementedError(dialect)
+        case never:
+            assert_never(never)
+    async with engine.begin() as conn:
+        _ = await conn.execute(text(query))
+
+
 ##
 
 
@@ -1166,6 +1180,7 @@ __all__ = [
     "create_engine",
     "ensure_database_created",
     "ensure_database_dropped",
+    "ensure_database_users_disconnected",
     "ensure_tables_created",
     "ensure_tables_dropped",
     "enum_name",