dbx-tools-postgres 0.6.78__py3-none-any.whl

dbx_tools/postgres/__init__.py

@@ -0,0 +1,109 @@
+from .address import (
+    SSL_MODES,
+    LakebaseConnectionInputs,
+    ParsedAddress,
+    SslMode,
+    parse_address,
+    parse_resource_path,
+    parseAddress,
+    parseResourcePath,
+)
+from .advisory_lock import (
+    ExplicitAdvisoryLockId,
+    acquire_advisory_lock,
+    acquire_advisory_lock_async,
+    advisory_lock,
+    advisory_lock_async,
+    advisory_lock_id,
+    advisory_transaction_lock,
+    advisory_transaction_lock_async,
+    explicit_advisory_lock_id,
+    release_advisory_lock,
+    release_advisory_lock_async,
+    try_advisory_lock,
+    try_advisory_lock_async,
+    try_advisory_transaction_lock,
+    try_advisory_transaction_lock_async,
+)
+from .engine import (
+    CredentialProvider,
+    PostgresEngineConfig,
+    ResolvedPostgresConnection,
+    WorkspaceClientLike,
+    autoscaling_credential_provider,
+    autoscalingCredentialProvider,
+    create_async_engine,
+    create_engine,
+    createAsyncEngine,
+    createEngine,
+    install_credential_injection,
+    installCredentialInjection,
+    resolve_postgres_connection,
+    resolvePostgresConnection,
+    workspace_credential_provider,
+    workspaceCredentialProvider,
+)
+from .topic_bus import (
+    PostgresTopicBus,
+    PostgresTopicBusOptions,
+    SerializableValue,
+    TopicListener,
+    TopicMessage,
+    TopicMetadata,
+    TopicMetadataProvider,
+    TopicPublishInput,
+    channel_name,
+    channelName,
+)
+
+__all__ = [
+    "SSL_MODES",
+    "CredentialProvider",
+    "ExplicitAdvisoryLockId",
+    "LakebaseConnectionInputs",
+    "ParsedAddress",
+    "PostgresEngineConfig",
+    "PostgresTopicBus",
+    "PostgresTopicBusOptions",
+    "ResolvedPostgresConnection",
+    "SerializableValue",
+    "SslMode",
+    "TopicListener",
+    "TopicMessage",
+    "TopicMetadata",
+    "TopicMetadataProvider",
+    "TopicPublishInput",
+    "WorkspaceClientLike",
+    "acquire_advisory_lock",
+    "acquire_advisory_lock_async",
+    "advisory_lock",
+    "advisory_lock_async",
+    "advisory_lock_id",
+    "advisory_transaction_lock",
+    "advisory_transaction_lock_async",
+    "autoscalingCredentialProvider",
+    "autoscaling_credential_provider",
+    "channelName",
+    "channel_name",
+    "createAsyncEngine",
+    "createEngine",
+    "create_async_engine",
+    "create_engine",
+    "explicit_advisory_lock_id",
+    "installCredentialInjection",
+    "install_credential_injection",
+    "parseAddress",
+    "parseResourcePath",
+    "parse_address",
+    "parse_resource_path",
+    "release_advisory_lock",
+    "release_advisory_lock_async",
+    "resolvePostgresConnection",
+    "resolve_postgres_connection",
+    "try_advisory_lock",
+    "try_advisory_lock_async",
+    "try_advisory_transaction_lock",
+    "try_advisory_transaction_lock_async",
+    "workspaceCredentialProvider",
+    "workspace_credential_provider",
+]

dbx_tools/postgres/address.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import re
+from dataclasses import asdict, dataclass
+from typing import Literal
+from urllib.parse import parse_qs, unquote, urlparse
+
+SslMode = Literal["require", "disable", "prefer"]
+SSL_MODES: tuple[SslMode, ...] = ("require", "disable", "prefer")
+
+_URL_SCHEME_RE = re.compile(r"^(postgres|postgresql)://", re.IGNORECASE)
+_PROJECT_ID_RE = re.compile(r"^[a-z][a-z0-9-]{0,61}[a-z0-9]$|^[a-z]$")
+_HOSTNAME_HINT_RE = re.compile(r"^[a-z0-9][a-z0-9-]*(\.[a-z0-9][a-z0-9-]*)+$", re.IGNORECASE)
+
+
+@dataclass(frozen=True, slots=True)
+class LakebaseConnectionInputs:
+    project: str | None = None
+    branch: str | None = None
+    endpoint: str | None = None
+    database: str | None = None
+    host: str | None = None
+    port: int | None = None
+    ssl_mode: SslMode | None = None
+
+    def as_dict(self) -> dict[str, object]:
+        return {key: value for key, value in asdict(self).items() if value is not None}
+
+
+@dataclass(frozen=True, slots=True)
+class ParsedAddress(LakebaseConnectionInputs):
+    endpoint_id: str | None = None
+    database_resource_id: str | None = None
+    user: str | None = None
+
+
+def parse_address(value: str | None) -> ParsedAddress:
+    if not value or not (address := value.strip()):
+        return ParsedAddress()
+    if _URL_SCHEME_RE.match(address):
+        return _parse_uri(address)
+    if address.startswith("projects/"):
+        return _parse_resource_path_segments(address)
+    if "." in address and _HOSTNAME_HINT_RE.fullmatch(address):
+        return ParsedAddress(host=address)
+    if _PROJECT_ID_RE.fullmatch(address):
+        return ParsedAddress(project=address)
+    return ParsedAddress()
+
+
+def parse_resource_path(value: str | None) -> ParsedAddress:
+    if not value or not (address := value.strip()).startswith("projects/"):
+        return ParsedAddress()
+    return _parse_resource_path_segments(address)
+
+
+def _parse_uri(address: str) -> ParsedAddress:
+    try:
+        parsed = urlparse(address)
+        port = parsed.port
+    except ValueError:
+        return ParsedAddress()
+    if parsed.scheme.lower() not in {"postgres", "postgresql"}:
+        return ParsedAddress()
+    ssl_mode_value = parse_qs(parsed.query).get("sslmode") or parse_qs(parsed.query).get("sslMode")
+    ssl_mode = ssl_mode_value[0].lower() if ssl_mode_value else None
+    return ParsedAddress(
+        host=parsed.hostname or None,
+        port=port,
+        user=unquote(parsed.username) if parsed.username else None,
+        database=unquote(parsed.path.removeprefix("/")) or None,
+        ssl_mode=ssl_mode if ssl_mode in SSL_MODES else None,
+    )
+
+
+def _parse_resource_path_segments(address: str) -> ParsedAddress:
+    parts = address.split("/")
+    if len(parts) < 2 or parts[0] != "projects" or not parts[1]:
+        return ParsedAddress()
+    project = parts[1]
+    if len(parts) == 2:
+        return ParsedAddress(project=project)
+    if len(parts) == 4 and parts[2] == "branches" and parts[3]:
+        return ParsedAddress(project=project, branch=parts[3])
+    if (
+        len(parts) == 6
+        and parts[2] == "branches"
+        and parts[3]
+        and parts[4] == "endpoints"
+        and parts[5]
+    ):
+        return ParsedAddress(
+            project=project,
+            branch=parts[3],
+            endpoint=address,
+            endpoint_id=parts[5],
+        )
+    if (
+        len(parts) == 6
+        and parts[2] == "branches"
+        and parts[3]
+        and parts[4] == "databases"
+        and parts[5]
+    ):
+        return ParsedAddress(
+            project=project,
+            branch=parts[3],
+            database_resource_id=parts[5],
+        )
+    return ParsedAddress()
+
+
+parseAddress = parse_address
+parseResourcePath = parse_resource_path

dbx_tools/postgres/advisory_lock.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import hashlib
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass
+from typing import Protocol
+
+from dbx_tools.core import to_stable_key
+from sqlalchemy import text
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+
+@dataclass(frozen=True, slots=True)
+class ExplicitAdvisoryLockId:
+    value: int
+
+
+class SyncQueryable(Protocol):
+    def execute(self, statement: object, parameters: dict[str, object] | None = None) -> object: ...
+
+
+class AsyncQueryable(Protocol):
+    async def execute(
+        self, statement: object, parameters: dict[str, object] | None = None
+    ) -> object: ...
+
+
+def advisory_lock_id(key: object) -> int:
+    if isinstance(key, ExplicitAdvisoryLockId):
+        return _signed_64(key.value)
+    parts = key if isinstance(key, (list, tuple)) else [key]
+    canonical = "\0".join(to_stable_key(part) for part in parts)
+    return int.from_bytes(hashlib.sha256(canonical.encode()).digest()[:8], "big", signed=True)
+
+
+def explicit_advisory_lock_id(value: int) -> ExplicitAdvisoryLockId:
+    return ExplicitAdvisoryLockId(value)
+
+
+def acquire_advisory_lock(
+    connection: SyncQueryable,
+    key: object,
+    *,
+    transaction: bool = False,
+    wait: bool = True,
+) -> bool:
+    function = _function_name(transaction=transaction, wait=wait)
+    result = connection.execute(
+        text(f"SELECT {function}(:lock_id)"), {"lock_id": advisory_lock_id(key)}
+    )
+    return True if wait else bool(result.scalar_one())
+
+
+async def acquire_advisory_lock_async(
+    connection: AsyncQueryable,
+    key: object,
+    *,
+    transaction: bool = False,
+    wait: bool = True,
+) -> bool:
+    function = _function_name(transaction=transaction, wait=wait)
+    result = await connection.execute(
+        text(f"SELECT {function}(:lock_id)"), {"lock_id": advisory_lock_id(key)}
+    )
+    return True if wait else bool(result.scalar_one())
+
+
+def release_advisory_lock(connection: SyncQueryable, key: object) -> None:
+    lock_id = advisory_lock_id(key)
+    result = connection.execute(text("SELECT pg_advisory_unlock(:lock_id)"), {"lock_id": lock_id})
+    if result.scalar_one() is not True:
+        raise RuntimeError(f"Postgres advisory lock {lock_id} was not held by this connection")
+
+
+async def release_advisory_lock_async(connection: AsyncQueryable, key: object) -> None:
+    lock_id = advisory_lock_id(key)
+    result = await connection.execute(
+        text("SELECT pg_advisory_unlock(:lock_id)"), {"lock_id": lock_id}
+    )
+    if result.scalar_one() is not True:
+        raise RuntimeError(f"Postgres advisory lock {lock_id} was not held by this connection")
+
+
+@contextmanager
+def advisory_lock(engine: Engine, key: object) -> Iterator[Connection]:
+    with engine.connect() as connection:
+        acquire_advisory_lock(connection, key)
+        failure: Exception | None = None
+        try:
+            yield connection
+        except Exception as error:  # noqa: BLE001
+            failure = error
+        unlock_failure = _capture_release(connection, key)
+        if failure is not None:
+            if unlock_failure is not None:
+                failure.add_note(f"Advisory lock release also failed: {unlock_failure}")
+            raise failure
+        if unlock_failure is not None:
+            raise unlock_failure
+
+
+@contextmanager
+def try_advisory_lock(engine: Engine, key: object) -> Iterator[Connection | None]:
+    with engine.connect() as connection:
+        if not acquire_advisory_lock(connection, key, wait=False):
+            yield None
+            return
+        failure: Exception | None = None
+        try:
+            yield connection
+        except Exception as error:  # noqa: BLE001
+            failure = error
+        unlock_failure = _capture_release(connection, key)
+        if failure is not None:
+            if unlock_failure is not None:
+                failure.add_note(f"Advisory lock release also failed: {unlock_failure}")
+            raise failure
+        if unlock_failure is not None:
+            raise unlock_failure
+
+
+@contextmanager
+def advisory_transaction_lock(engine: Engine, key: object) -> Iterator[Connection]:
+    with engine.begin() as connection:
+        acquire_advisory_lock(connection, key, transaction=True)
+        yield connection
+
+
+@contextmanager
+def try_advisory_transaction_lock(engine: Engine, key: object) -> Iterator[Connection | None]:
+    with engine.begin() as connection:
+        if not acquire_advisory_lock(connection, key, transaction=True, wait=False):
+            yield None
+            return
+        yield connection
+
+
+@asynccontextmanager
+async def advisory_lock_async(engine: AsyncEngine, key: object) -> AsyncIterator[AsyncConnection]:
+    async with engine.connect() as connection:
+        await acquire_advisory_lock_async(connection, key)
+        failure: Exception | None = None
+        try:
+            yield connection
+        except Exception as error:  # noqa: BLE001
+            failure = error
+        unlock_failure = await _capture_release_async(connection, key)
+        if failure is not None:
+            if unlock_failure is not None:
+                failure.add_note(f"Advisory lock release also failed: {unlock_failure}")
+            raise failure
+        if unlock_failure is not None:
+            raise unlock_failure
+
+
+@asynccontextmanager
+async def try_advisory_lock_async(
+    engine: AsyncEngine, key: object
+) -> AsyncIterator[AsyncConnection | None]:
+    async with engine.connect() as connection:
+        if not await acquire_advisory_lock_async(connection, key, wait=False):
+            yield None
+            return
+        failure: Exception | None = None
+        try:
+            yield connection
+        except Exception as error:  # noqa: BLE001
+            failure = error
+        unlock_failure = await _capture_release_async(connection, key)
+        if failure is not None:
+            if unlock_failure is not None:
+                failure.add_note(f"Advisory lock release also failed: {unlock_failure}")
+            raise failure
+        if unlock_failure is not None:
+            raise unlock_failure
+
+
+@asynccontextmanager
+async def advisory_transaction_lock_async(
+    engine: AsyncEngine, key: object
+) -> AsyncIterator[AsyncConnection]:
+    async with engine.begin() as connection:
+        await acquire_advisory_lock_async(connection, key, transaction=True)
+        yield connection
+
+
+@asynccontextmanager
+async def try_advisory_transaction_lock_async(
+    engine: AsyncEngine, key: object
+) -> AsyncIterator[AsyncConnection | None]:
+    async with engine.begin() as connection:
+        if not await acquire_advisory_lock_async(connection, key, transaction=True, wait=False):
+            yield None
+            return
+        yield connection
+
+
+def _function_name(*, transaction: bool, wait: bool) -> str:
+    if transaction:
+        return "pg_advisory_xact_lock" if wait else "pg_try_advisory_xact_lock"
+    return "pg_advisory_lock" if wait else "pg_try_advisory_lock"
+
+
+def _signed_64(value: int) -> int:
+    return ((value + 2**63) % 2**64) - 2**63
+
+
+def _capture_release(connection: SyncQueryable, key: object) -> Exception | None:
+    try:
+        release_advisory_lock(connection, key)
+    except Exception as error:  # noqa: BLE001
+        return error
+    return None
+
+
+async def _capture_release_async(connection: AsyncQueryable, key: object) -> Exception | None:
+    try:
+        await release_advisory_lock_async(connection, key)
+    except Exception as error:  # noqa: BLE001
+        return error
+    return None