dycw-utilities 0.150.0__py3-none-any.whl → 0.150.2__py3-none-any.whl

{dycw_utilities-0.150.0.dist-info → dycw_utilities-0.150.2.dist-info}/METADATA

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

{dycw_utilities-0.150.0.dist-info → dycw_utilities-0.150.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=-FKXlj49n14usgRRgsXp-6U1_cXzP-blAYQWGuSGKL8,60
+utilities/__init__.py,sha256=bBUlYYM1VrEWyFEkC5ZlwXnej-DEWPaWXPM1YmuRbVk,60
 utilities/altair.py,sha256=92E2lCdyHY4Zb-vCw6rEJIsWdKipuu-Tu2ab1ufUfAk,9079
 utilities/asyncio.py,sha256=2m2a2C-Qgc6OHTTHL332-t66A7xDITt_SORT7a1DJWo,16792
 utilities/atomicwrites.py,sha256=xcOWenTBRS0oat3kg7Sqe51AohNThMQ2ixPL7QCG8hw,5795
@@ -49,7 +49,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=C3TxpQymM7S5ZivxIN_AFTM5g5TtENptzx8-XYf7pow,12182
+utilities/postgres.py,sha256=_LJDtzw50gHaPyW8lu_Z9zBh7eXK7A10yx-hJZ0f4nc,13084
 utilities/pottery.py,sha256=u0uvyGgYyujxftEMlsv6ppYTKQoVVjHt5jnVxxYz9s4,6596
 utilities/pqdm.py,sha256=BTsYPtbKQWwX-iXF4qCkfPG7DPxIB54J989n83bXrIo,3092
 utilities/psutil.py,sha256=KUlu4lrUw9Zg1V7ZGetpWpGb9DB8l_SSDWGbANFNCPU,2104
@@ -66,7 +66,7 @@ utilities/sentinel.py,sha256=3jIwgpMekWgDAxPDA_hXMP2St43cPhciKN3LWiZ7kv0,1248
 utilities/shelve.py,sha256=4OzjQI6kGuUbJciqf535rwnao-_IBv66gsT6tRGiUt0,759
 utilities/slack_sdk.py,sha256=ppFBvKgfg5IRWiIoKPtpTyzBtBF4XmwEvU3I5wLJikM,2140
 utilities/socket.py,sha256=K77vfREvzoVTrpYKo6MZakol0EYu2q1sWJnnZqL0So0,118
-utilities/sqlalchemy.py,sha256=M__UwdIiXDc6W8oRHQ6u3h_zEb9uTKwfl506KJCrOl4,37231
+utilities/sqlalchemy.py,sha256=5TPqPvJUCuxIeiMD7IKTptYg11UXNy7ThGw81reCIRU,37290
 utilities/sqlalchemy_polars.py,sha256=18AoEbeNJUKF3-5hroNy9J5LQwS_QJAXbMfKc9sChtk,14250
 utilities/statsmodels.py,sha256=koyiBHvpMcSiBfh99wFUfSggLNx7cuAw3rwyfAhoKpQ,3410
 utilities/string.py,sha256=MB0X6UPTUc06JdAdj-PctZ238IXeCjE5dAJibNw6ZrU,587
@@ -89,8 +89,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.150.0.dist-info/METADATA,sha256=F2t8-Dhx3lTo-_ns45mIJ0vtLW9KTBUHrViS1bR94v0,1697
-dycw_utilities-0.150.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.150.0.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
-dycw_utilities-0.150.0.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.150.0.dist-info/RECORD,,
+dycw_utilities-0.150.2.dist-info/METADATA,sha256=Hp0ONkfWYpv5XkLgy5V-bKY3vVQ1kgMD0FwkS7xsCok,1697
+dycw_utilities-0.150.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.150.2.dist-info/entry_points.txt,sha256=BOD_SoDxwsfJYOLxhrSXhHP_T7iw-HXI9f2WVkzYxvQ,135
+dycw_utilities-0.150.2.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.150.2.dist-info/RECORD,,

utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.150.0"
+__version__ = "0.150.2"

utilities/postgres.py

@@ -42,6 +42,7 @@ async def pg_dump(
     table_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
+    role: str | None = None,
     docker: str | None = None,
     dry_run: bool = False,
     logger: LoggerOrName | None = None,
@@ -62,6 +63,7 @@ async def pg_dump(
         table_exc=table_exc,
         inserts=inserts,
         on_conflict_do_nothing=on_conflict_do_nothing,
+        role=role,
         docker=docker,
     )
     if dry_run:
@@ -110,6 +112,7 @@ def _build_pg_dump(
     table_exc: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
     inserts: bool = False,
     on_conflict_do_nothing: bool = False,
+    role: str | None = None,
     docker: str | None = None,
 ) -> str:
     database, host, port = _extract_url(url)
@@ -122,6 +125,7 @@ def _build_pg_dump(
         f"--format={format_}",
         "--verbose",
         # output options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
         "--large-objects",
         "--create",
         "--no-owner",
@@ -133,10 +137,6 @@ def _build_pg_dump(
     ]
     if (format_ == "directory") and (jobs is not None):
         parts.append(f"--jobs={jobs}")
-    if data_only:
-        parts.append("--data-only")
-    if clean:
-        parts.extend(["--clean", "--if-exists"])
     if schema is not None:
         parts.extend([f"--schema={s}" for s in always_iterable(schema)])
     if schema_exc is not None:
@@ -153,6 +153,8 @@ def _build_pg_dump(
         parts.append("--on-conflict-do-nothing")
     if url.username is not None:
         parts.append(f"--username={url.username}")
+    if role is not None:
+        parts.append(f"--role={role}")
     if docker is not None:
         parts = _wrap_docker(parts, docker)
     return " ".join(parts)
@@ -192,6 +194,7 @@ async def restore(
     schema: MaybeListStr | None = None,
     schema_exc: MaybeListStr | None = None,
     table: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
     docker: str | None = None,
     dry_run: bool = False,
     logger: LoggerOrName | None = None,
@@ -208,6 +211,7 @@ async def restore(
         schema=schema,
         schema_exc=schema_exc,
         table=table,
+        role=role,
         docker=docker,
     )
     if dry_run:
@@ -255,6 +259,7 @@ def _build_pg_restore_or_psql(
     schema: MaybeListStr | None = None,
     schema_exc: MaybeListStr | None = None,
     table: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
     docker: str | None = None,
 ) -> str:
     path = Path(path)
@@ -270,6 +275,7 @@ def _build_pg_restore_or_psql(
         schemas=schema,
         schemas_exc=schema_exc,
         tables=table,
+        role=role,
         docker=docker,
     )
 
@@ -286,6 +292,7 @@ def _build_pg_restore(
     schemas: MaybeListStr | None = None,
     schemas_exc: MaybeListStr | None = None,
     tables: MaybeSequence[TableOrORMInstOrClass | str] | None = None,
+    role: str | None = None,
     docker: str | None = None,
 ) -> str:
     """Run `pg_restore`."""
@@ -297,6 +304,8 @@ def _build_pg_restore(
         f"--dbname={database_use}",
         "--verbose",
         # restore options
+        *_resolve_data_only_and_clean(data_only=data_only, clean=clean),
+        "--create",
         "--exit-on-error",
         "--no-owner",
         "--no-privileges",
@@ -305,10 +314,6 @@ def _build_pg_restore(
         f"--port={port}",
         "--no-password",
     ]
-    if data_only:
-        parts.append("--data-only")
-    if clean:
-        parts.extend(["--clean", "--if-exists"])
     if jobs is not None:
         parts.append(f"--jobs={jobs}")
     if schemas is not None:
@@ -319,6 +324,8 @@ def _build_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 role is not None:
+        parts.append(f"--role={role}")
     if docker is not None:
         parts = _wrap_docker(parts, docker)
     parts.append(str(path))
@@ -402,6 +409,29 @@ def _get_table_name(obj: TableOrORMInstOrClass | str, /) -> str:
             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 _ as 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"
+
+
 def _wrap_docker(parts: list[str], container: str, /) -> list[str]:
     return ["docker", "exec", "-it", container, *parts]
 

utilities/sqlalchemy.py

@@ -57,7 +57,12 @@ from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
 from sqlalchemy.dialects.sqlite import Insert as sqlite_Insert
 from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-from sqlalchemy.exc import ArgumentError, DatabaseError, OperationalError
+from sqlalchemy.exc import (
+    ArgumentError,
+    DatabaseError,
+    OperationalError,
+    ProgrammingError,
+)
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 from sqlalchemy.orm import (
     DeclarativeBase,
@@ -115,7 +120,7 @@ def check_connect(engine: Engine, /) -> bool:
     try:
         with engine.connect() as conn:
             return bool(conn.execute(text("SELECT 1")).scalar_one())
-    except OperationalError:
+    except (OperationalError, ProgrammingError):
         return False