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

dbx_tools/postgres/topic_bus.py

@@ -0,0 +1,380 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import math
+import os
+import platform
+import re
+import socket
+import uuid
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Protocol, TypeAlias
+
+from dbx_tools.core import fnv_hash, to_identifier, to_stable_key
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+SerializableValue: TypeAlias = (
+    str | int | float | bool | None | list["SerializableValue"] | dict[str, "SerializableValue"]
+)
+TopicMetadata: TypeAlias = dict[str, SerializableValue]
+TopicListener: TypeAlias = Callable[["TopicMessage"], Awaitable[None] | None]
+TopicMetadataProvider: TypeAlias = Callable[[], Awaitable[TopicMetadata] | TopicMetadata]
+
+_DEFAULT_CHANNEL = "dbx_tools_topic_bus"
+_MAX_CHANNEL_LENGTH = 63
+_CHANNEL_HASH_LENGTH = 6
+_CHANNEL_FALLBACK = "bus"
+_MAX_NOTIFY_BYTES = 7_900
+_MIN_RECONNECT_DELAY = 0.25
+_MAX_RECONNECT_DELAY = 5.0
+
+
+class AsyncEngineLike(Protocol):
+    def begin(self) -> Any: ...
+
+    async def raw_connection(self) -> Any: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TopicPublishInput:
+    type: str
+    body: SerializableValue
+    metadata: TopicMetadata = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class TopicMessage:
+    id: str
+    topic: str
+    type: str
+    metadata: TopicMetadata
+    body: SerializableValue
+    published_at: str
+
+    @property
+    def publishedAt(self) -> str:
+        return self.published_at
+
+    def as_dict(self) -> dict[str, SerializableValue]:
+        return {
+            "id": self.id,
+            "topic": self.topic,
+            "type": self.type,
+            "metadata": self.metadata,
+            "body": self.body,
+            "publishedAt": self.published_at,
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class PostgresTopicBusOptions:
+    channel: object = _DEFAULT_CHANNEL
+    metadata: TopicMetadata | TopicMetadataProvider | None = None
+    on_error: Callable[[BaseException], None] | None = None
+
+
+class PostgresTopicBus:
+    def __init__(
+        self,
+        engine: AsyncEngine | AsyncEngineLike,
+        options: PostgresTopicBusOptions | None = None,
+        *,
+        channel: object | None = None,
+        metadata: TopicMetadata | TopicMetadataProvider | None = None,
+        on_error: Callable[[BaseException], None] | None = None,
+    ) -> None:
+        configured = options or PostgresTopicBusOptions()
+        channel_value = configured.channel if channel is None else channel
+        self.engine = engine
+        self.channel_name = channel_name(channel_value)
+        self.metadata = configured.metadata if metadata is None else metadata
+        self.on_error = on_error or configured.on_error or (lambda error: None)
+        self._listeners: dict[str, set[TopicListener]] = {}
+        self._raw_connection: Any | None = None
+        self._driver_connection: Any | None = None
+        self._start_lock = asyncio.Lock()
+        self._reconnect_task: asyncio.Task[None] | None = None
+        self._closed = False
+
+    @property
+    def channelName(self) -> str:
+        return self.channel_name
+
+    async def start(self) -> None:
+        if self._driver_connection is not None:
+            return
+        if self._closed:
+            raise RuntimeError("Postgres topic bus is closed")
+        async with self._start_lock:
+            if self._driver_connection is not None:
+                return
+            raw_connection = await self.engine.raw_connection()
+            driver_connection = raw_connection.driver_connection
+            try:
+                await driver_connection.add_listener(self.channel_name, self._handle_notification)
+                add_termination_listener = getattr(
+                    driver_connection, "add_termination_listener", None
+                )
+                if add_termination_listener:
+                    add_termination_listener(self._handle_termination)
+            except BaseException:
+                await _maybe_await(raw_connection.close())
+                raise
+            if self._closed:
+                await driver_connection.remove_listener(
+                    self.channel_name,
+                    self._handle_notification,
+                )
+                await _maybe_await(raw_connection.close())
+                raise RuntimeError("Postgres topic bus is closed")
+            self._raw_connection = raw_connection
+            self._driver_connection = driver_connection
+
+    async def broadcast(
+        self,
+        topic: str,
+        message_input: TopicPublishInput | Mapping[str, Any],
+    ) -> TopicMessage:
+        if not topic.strip():
+            raise TypeError("Topic must not be empty")
+        if self._closed:
+            raise RuntimeError("Postgres topic bus is closed")
+        publish = _publish_input(message_input)
+        if not publish.type.strip():
+            raise TypeError("Message type must not be empty")
+        if not _is_serializable(publish.metadata) or not _is_serializable(publish.body):
+            raise TypeError("Message metadata and body must be JSON serializable without coercion")
+        automatic = await self._resolve_metadata()
+        message = TopicMessage(
+            id=str(uuid.uuid4()),
+            topic=topic,
+            type=publish.type,
+            metadata={**automatic, **publish.metadata},
+            body=publish.body,
+            published_at=datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
+        )
+        encoded = json.dumps(message.as_dict(), separators=(",", ":"), allow_nan=False)
+        if len(encoded.encode("utf-8")) > _MAX_NOTIFY_BYTES:
+            raise ValueError(f"Postgres notification exceeds {_MAX_NOTIFY_BYTES} bytes")
+        async with self.engine.begin() as connection:
+            await connection.execute(
+                text("SELECT pg_notify(:channel, :payload)"),
+                {"channel": self.channel_name, "payload": encoded},
+            )
+        return message
+
+    async def listen(self, topic: str, listener: TopicListener) -> Callable[[], Awaitable[None]]:
+        if not topic.strip():
+            raise TypeError("Topic must not be empty")
+        await self.start()
+        listeners = self._listeners.setdefault(topic, set())
+        listeners.add(listener)
+
+        async def unsubscribe() -> None:
+            listeners.discard(listener)
+            if not listeners:
+                self._listeners.pop(topic, None)
+
+        return unsubscribe
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        if self._reconnect_task:
+            self._reconnect_task.cancel()
+            await asyncio.gather(self._reconnect_task, return_exceptions=True)
+            self._reconnect_task = None
+        raw_connection = self._raw_connection
+        driver_connection = self._driver_connection
+        self._raw_connection = None
+        self._driver_connection = None
+        self._listeners.clear()
+        if driver_connection is not None:
+            remove_termination_listener = getattr(
+                driver_connection,
+                "remove_termination_listener",
+                None,
+            )
+            if remove_termination_listener:
+                remove_termination_listener(self._handle_termination)
+            try:
+                await driver_connection.remove_listener(
+                    self.channel_name,
+                    self._handle_notification,
+                )
+            except Exception as error:
+                self.on_error(error)
+        if raw_connection is not None:
+            await _maybe_await(raw_connection.close())
+
+    async def _resolve_metadata(self) -> TopicMetadata:
+        configured = self.metadata() if callable(self.metadata) else (self.metadata or {})
+        if inspect.isawaitable(configured):
+            configured = await configured
+        if not _is_serializable(configured) or not isinstance(configured, dict):
+            raise TypeError("Bus metadata must be JSON serializable without coercion")
+        return {**_machine_metadata(), **configured}
+
+    def _handle_notification(
+        self,
+        connection: object,
+        process_id: int,
+        channel: str,
+        payload: str,
+    ) -> None:
+        del connection, process_id
+        if channel != self.channel_name:
+            return
+        message = _decode(payload)
+        if message is None:
+            return
+        for listener in tuple(self._listeners.get(message.topic, ())):
+            asyncio.create_task(self._deliver(listener, message))
+
+    async def _deliver(self, listener: TopicListener, message: TopicMessage) -> None:
+        try:
+            result = listener(message)
+            if inspect.isawaitable(result):
+                await result
+        except Exception as error:
+            self.on_error(error)
+
+    def _handle_termination(self, connection: object) -> None:
+        del connection
+        if self._closed or not self._listeners or self._reconnect_task is not None:
+            return
+        raw_connection = self._raw_connection
+        self._raw_connection = None
+        self._driver_connection = None
+        self._reconnect_task = asyncio.create_task(self._reconnect(raw_connection))
+
+    async def _reconnect(self, raw_connection: object | None = None) -> None:
+        delay = 0.0
+        try:
+            if raw_connection is not None:
+                try:
+                    await _maybe_await(raw_connection.close())
+                except Exception as error:
+                    self.on_error(error)
+            while not self._closed and self._listeners:
+                if delay:
+                    await asyncio.sleep(delay)
+                try:
+                    await self.start()
+                    return
+                except Exception as error:
+                    if self._closed:
+                        return
+                    self.on_error(error)
+                    delay = (
+                        _MIN_RECONNECT_DELAY if delay == 0 else min(delay * 2, _MAX_RECONNECT_DELAY)
+                    )
+        finally:
+            self._reconnect_task = None
+
+
+def channel_name(channel: object = _DEFAULT_CHANNEL) -> str:
+    parts = list(channel) if isinstance(channel, (list, tuple)) else [channel]
+    stable = "\0".join(to_stable_key(part) for part in parts)
+    suffix = fnv_hash(stable, length=_CHANNEL_HASH_LENGTH)
+    labels = [part for part in parts if isinstance(part, (str, int, float, bool))]
+    body = to_identifier(*labels, delimiter="_")[: _MAX_CHANNEL_LENGTH - len(suffix) - 1].rstrip(
+        "_"
+    )
+    prefix = body if re.match(r"^[A-Za-z_]", body) else f"{_CHANNEL_FALLBACK}_{body}"
+    return re.sub(r"_+", "_", f"{prefix}_{suffix}")
+
+
+def _publish_input(value: TopicPublishInput | Mapping[str, Any]) -> TopicPublishInput:
+    if isinstance(value, TopicPublishInput):
+        return value
+    return TopicPublishInput(
+        type=value.get("type", ""),
+        metadata=dict(value.get("metadata") or {}),
+        body=value.get("body"),
+    )
+
+
+def _decode(payload: str) -> TopicMessage | None:
+    try:
+        value = json.loads(payload)
+    except (TypeError, json.JSONDecodeError):
+        return None
+    if not isinstance(value, dict) or not _is_serializable(value):
+        return None
+    required = ("id", "topic", "type", "metadata", "body", "publishedAt")
+    if not all(key in value for key in required) or not isinstance(value["metadata"], dict):
+        return None
+    if not all(isinstance(value[key], str) for key in ("id", "topic", "type", "publishedAt")):
+        return None
+    return TopicMessage(
+        id=value["id"],
+        topic=value["topic"],
+        type=value["type"],
+        metadata=value["metadata"],
+        body=value["body"],
+        published_at=value["publishedAt"],
+    )
+
+
+def _machine_metadata() -> TopicMetadata:
+    values: dict[str, SerializableValue | None] = {
+        "project": _first_env(
+            "DATABRICKS_APP_NAME",
+            "DATABRICKS_BUNDLE_NAME",
+            "PROJECT_NAME",
+        ),
+        "hostname": socket.gethostname(),
+        "platform": platform.system().lower(),
+        "environment": _env("PYTHON_ENV") or _env("NODE_ENV"),
+        "appName": _env("DATABRICKS_APP_NAME"),
+        "deploymentId": _env("DATABRICKS_APP_DEPLOYMENT_ID"),
+        "databricksHost": _env("DATABRICKS_HOST"),
+    }
+    return {key: value for key, value in values.items() if value is not None}
+
+
+def _env(name: str) -> str | None:
+    value = os.environ.get(name)
+    return value.strip() if value and value.strip() else None
+
+
+def _first_env(*names: str) -> str | None:
+    return next((value for name in names if (value := _env(name))), None)
+
+
+def _is_serializable(value: object, seen: set[int] | None = None) -> bool:
+    if value is None or isinstance(value, (str, bool, int)):
+        return True
+    if isinstance(value, float):
+        return math.isfinite(value)
+    seen = seen or set()
+    identity = id(value)
+    if identity in seen:
+        return False
+    if isinstance(value, list):
+        seen.add(identity)
+        try:
+            return all(_is_serializable(item, seen) for item in value)
+        finally:
+            seen.remove(identity)
+    if isinstance(value, dict) and all(isinstance(key, str) for key in value):
+        seen.add(identity)
+        try:
+            return all(_is_serializable(item, seen) for item in value.values())
+        finally:
+            seen.remove(identity)
+    return False
+
+
+async def _maybe_await(value: object) -> object:
+    return await value if inspect.isawaitable(value) else value
+
+
+channelName = channel_name

dbx_tools_postgres-0.6.78.dist-info/METADATA

@@ -0,0 +1,188 @@
+Metadata-Version: 2.3
+Name: dbx-tools-postgres
+Version: 0.6.78
+Summary: WorkspaceClient-backed Lakebase Postgres resolution, SQLAlchemy engines, advisory locks, and LISTEN/NOTIFY topic bus
+Requires-Dist: asyncpg>=0.30,<1
+Requires-Dist: databricks-sdk>=0.63.0,<1
+Requires-Dist: greenlet>=3.2,<4
+Requires-Dist: dbx-tools-core==0.6.78
+Requires-Dist: psycopg[binary]>=3.2.9,<4
+Requires-Dist: sqlalchemy>=2.0.41,<3
+Requires-Python: >=3.10
+Project-URL: Source, https://github.com/reggie-db/dbx-tools/tree/main/packages/py/postgres
+Description-Content-Type: text/markdown
+
+# `dbx-tools-postgres`
+
+Python Lakebase/Postgres connection setup, advisory locks, and topic fan-out for
+services that already hold a Databricks `WorkspaceClient`. This package is
+currently an unpublished workspace package, and is the Python counterpart to
+`@dbx-tools/postgres` plus `@dbx-tools/appkit`'s address parsing.
+
+Install directly from this monorepo:
+
+```bash
+pip install "dbx-tools-postgres @ git+https://github.com/reggie-db/dbx-tools.git@main#subdirectory=packages/py/postgres"
+```
+
+Key features:
+
+- accepts the same Postgres URI, Lakebase resource path, hostname, and project-id
+  address shapes as `@dbx-tools/appkit`;
+- resolves missing autoscaling endpoint fields through
+  `WorkspaceClient.api_client`;
+- resolves provisioned Lakebase instance DNS through
+  `WorkspaceClient.database`;
+- injects a fresh database credential on SQLAlchemy's `do_connect` event rather
+  than storing an expiring password in the engine URL, using the SDK's
+  provisioned-instance API or the Autoscaling `/postgres/credentials` endpoint;
+- supports sync psycopg and asyncpg SQLAlchemy engines;
+- derives advisory-lock ids from the same stable structured keys as the Node
+  package and holds one checked-out connection for the full critical section;
+- provides blocking and try-lock context managers for session and transaction
+  locks, with sync and async SQLAlchemy variants;
+- fans messages out to every process on a channel with `PostgresTopicBus`, using
+  the same lifecycle and wire envelope as the Node package.
+
+```python
+from databricks.sdk import WorkspaceClient
+from dbx_tools.postgres import PostgresEngineConfig, create_async_engine
+
+engine = create_async_engine(
+    WorkspaceClient(),
+    PostgresEngineConfig(instance_name="my-lakebase", database="databricks_postgres"),
+    pool_pre_ping=True,
+    pool_recycle=1800,
+)
+```
+
+Pass `credential_provider=` to either engine factory to inject credentials from
+another source while retaining the same connect-time rotation behavior.
+
+```python
+from dbx_tools.postgres import advisory_transaction_lock
+
+with advisory_transaction_lock(engine, ["schema-install", "v2"]) as connection:
+    connection.exec_driver_sql("CREATE TABLE IF NOT EXISTS ...")
+```
+
+## Topic bus
+
+`PostgresTopicBus` is async Postgres topic fan-out built on `LISTEN`/`NOTIFY`. Its
+public lifecycle and wire shape match `@dbx-tools/postgres`'s `PostgresTopicBus`,
+so Node and Python services can share a channel:
+
+- `PostgresTopicBus(engine, options)`;
+- `channelName`;
+- `await start()`;
+- `await broadcast(topic, TopicPublishInput(...))`;
+- `await listen(topic, listener)` returning an async unsubscribe function;
+- `await close()`;
+- envelope fields `id`, `topic`, `type`, `metadata`, `body`, and `publishedAt`.
+
+Channel derivation ports the Node stable-key and FNV rules, so equivalent channel
+parts resolve to the same PostgreSQL identifier in Python and Node.
+
+```python
+from dbx_tools.postgres import PostgresTopicBus, TopicPublishInput
+
+bus = PostgresTopicBus(engine, channel=["billing", "production"])
+
+unsubscribe = await bus.listen("invoice.updated", handle_invoice)
+await bus.broadcast(
+    "invoice.updated",
+    TopicPublishInput(type="invoice.updated", body={"invoice_id": "inv-7"}),
+)
+```
+
+Delivery is live and unstored, like PostgreSQL `LISTEN`/`NOTIFY` itself. Use a
+table or queue when consumers need replay or acknowledgements.
+
+## Databricks notebooks and Spark
+
+Verified end to end against a Lakebase endpoint on serverless notebook compute.
+[`packages/example/notebooks/bus-lakebase.py`](../../example/notebooks/bus-lakebase.py)
+is the runnable version of everything below.
+
+Two things about the Databricks Python runtime change how the bus is called, and
+neither is a limitation of the bus itself:
+
+- **A notebook kernel already runs an event loop**, so `asyncio.run` in a cell
+  raises `RuntimeError: asyncio.run() cannot be called from a running event
+loop`. Drive the coroutine on a short-lived thread with its own loop rather
+  than reaching for `nest_asyncio` — the bus holds a dedicated `LISTEN`
+  connection bound to whichever loop started it, so one loop per bus lifetime is
+  the invariant to preserve.
+- **Install with `%pip install`, not `pip install --target`.** A `--target`
+  install leaves the runtime's preloaded `typing_extensions` ahead of the new
+  one on `sys.path`, and importing `dbx_tools.postgres` then fails with
+  `ImportError: cannot import name 'TypeAliasType'`. `%pip` restarts the Python
+  process, which resolves it.
+
+### Publishing from a Spark UDF
+
+Publishing from executors works. Listening from them does not, and should not be
+attempted: a UDF invocation is short-lived, while `listen` keeps a connection
+open until `close`.
+
+Executors have no Databricks credentials, so they cannot build a
+`WorkspaceClient`. This is where connect-time credential injection pays off — the
+driver mints the Lakebase token once and the UDF closes over it, so the executor
+builds a plain SQLAlchemy engine and installs the token as its provider:
+
+```python
+from sqlalchemy import URL
+from sqlalchemy.ext.asyncio import create_async_engine
+from dbx_tools.postgres import (
+    PostgresTopicBus,
+    TopicPublishInput,
+    install_credential_injection,
+)
+
+# driver: resolve once, capture in the closure
+resolved = resolve_postgres_connection(workspace_client, config)
+token = workspace_client.api_client.do(
+    "POST",
+    "/api/2.0/postgres/credentials",
+    body={"endpoint": resolved.endpoint},
+)["token"]
+
+
+@udf(returnType=StringType())
+def publish(key: str) -> str:
+    async def run() -> str:
+        engine = create_async_engine(
+            URL.create(
+                "postgresql+asyncpg",
+                username=resolved.user,
+                host=resolved.host,
+                port=resolved.port,
+                database=resolved.database,
+                query={"ssl": resolved.ssl_mode},
+            )
+        )
+        install_credential_injection(engine.sync_engine, lambda: token)
+        bus = PostgresTopicBus(engine, channel="app-events")
+        try:
+            message = await bus.broadcast(
+                "row.processed", TopicPublishInput(type="row.processed", body={"key": key})
+            )
+            return message.id
+        finally:
+            await bus.close()
+            await engine.dispose()
+
+    return asyncio.run(run())
+```
+
+Constraints worth knowing before this reaches production:
+
+- A captured token EXPIRES (about an hour). Re-mint per job run; a long-running
+  streaming query needs a provider that refreshes instead of a captured string.
+- Each UDF call opens and closes its own connection, so batch the publish at
+  partition scope (`mapInPandas`, `foreachPartition`) rather than per row.
+- `spark.sparkContext.broadcast` is unavailable on serverless (Spark Connect).
+  A plain closure over driver-side values serializes with the UDF and is enough.
+- Delivery stays live and unstored: if no listener is connected when the UDF
+  publishes, the message is gone. Write to a table when executors produce
+  results a consumer must not miss.

dbx_tools_postgres-0.6.78.dist-info/RECORD

@@ -0,0 +1,8 @@
+dbx_tools/postgres/__init__.py,sha256=Qb2fH8vxKBpEfuLL2wlgh6_5qS-LMrK2hYByjDczrf4,2780
+dbx_tools/postgres/address.py,sha256=ulOuhMUzYeb5ZslYqe0rTNXqDfovk8QR5s8U6Aej20k,3729
+dbx_tools/postgres/advisory_lock.py,sha256=1Lfp2UOoQ4vJvLOfg47hsKVJpIrfeRByGoAbyw2WMYA,7655
+dbx_tools/postgres/engine.py,sha256=QqR49JVTTfvW-hoxiW08vSKfptClsyD4WsSkM08hcaA,15290
+dbx_tools/postgres/topic_bus.py,sha256=Zr0mtHXr1Gv-cdOW-j7mVwQMZsqjMcVyM-SysjpyNIk,13731
+dbx_tools_postgres-0.6.78.dist-info/WHEEL,sha256=lrO5MD1WVAWzcbNy_L2BwtfPrcM3KfUFpbKYXLkJX4A,80
+dbx_tools_postgres-0.6.78.dist-info/METADATA,sha256=5IVKuaMhXQas5iHTHvzpktxmRNnlpApXIJqLazALD44,7623
+dbx_tools_postgres-0.6.78.dist-info/RECORD,,

dbx_tools_postgres-0.6.78.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.1
+Root-Is-Purelib: true
+Tag: py3-none-any