dbx-tools-postgres 0.6.89__tar.gz → 0.6.91__tar.gz

{dbx_tools_postgres-0.6.89 → dbx_tools_postgres-0.6.91}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: dbx-tools-postgres
-Version: 0.6.89
+Version: 0.6.91
 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.89
-Requires-Dist: psycopg[binary]>=3.2.9,<4
-Requires-Dist: sqlalchemy>=2.0.41,<3
+Requires-Dist: asyncpg>=0.30
+Requires-Dist: databricks-sdk>=0.63.0
+Requires-Dist: greenlet>=3.2
+Requires-Dist: dbx-tools-core==0.6.91
+Requires-Dist: psycopg[binary]>=3.2.9
+Requires-Dist: sqlalchemy>=2.0.41
 Requires-Python: >=3.10
 Project-URL: Source, https://github.com/reggie-db/dbx-tools/tree/main/packages/py/postgres
 Description-Content-Type: text/markdown
@@ -39,9 +39,11 @@ Key features:
   `WorkspaceClient.api_client`;
 - resolves provisioned Lakebase instance DNS through
   `WorkspaceClient.database`;
-- injects a fresh database credential on SQLAlchemy's `do_connect` event rather
+- injects a cached 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;
+- refreshes built-in credential providers ahead of expiry with a process-local
+  check-lock-check load, so concurrent pool connections share one mint;
 - 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;
@@ -63,7 +65,8 @@ engine = create_async_engine(
 ```
 
 Pass `credential_provider=` to either engine factory to inject credentials from
-another source while retaining the same connect-time rotation behavior.
+another source. A custom provider owns its own cache, expiry, and refresh
+serialization policy.
 
 ```python
 from dbx_tools.postgres import advisory_transaction_lock

{dbx_tools_postgres-0.6.89 → dbx_tools_postgres-0.6.91}/README.md

@@ -25,9 +25,11 @@ Key features:
   `WorkspaceClient.api_client`;
 - resolves provisioned Lakebase instance DNS through
   `WorkspaceClient.database`;
-- injects a fresh database credential on SQLAlchemy's `do_connect` event rather
+- injects a cached 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;
+- refreshes built-in credential providers ahead of expiry with a process-local
+  check-lock-check load, so concurrent pool connections share one mint;
 - 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;
@@ -49,7 +51,8 @@ engine = create_async_engine(
 ```
 
 Pass `credential_provider=` to either engine factory to inject credentials from
-another source while retaining the same connect-time rotation behavior.
+another source. A custom provider owns its own cache, expiry, and refresh
+serialization policy.
 
 ```python
 from dbx_tools.postgres import advisory_transaction_lock

{dbx_tools_postgres-0.6.89 → dbx_tools_postgres-0.6.91}/pyproject.toml

@@ -1,16 +1,16 @@
 [project]
 name = "dbx-tools-postgres"
-version = "0.6.89"
+version = "0.6.91"
 description = "WorkspaceClient-backed Lakebase Postgres resolution, SQLAlchemy engines, advisory locks, and LISTEN/NOTIFY topic bus"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-    "asyncpg>=0.30,<1",
-    "databricks-sdk>=0.63.0,<1",
-    "greenlet>=3.2,<4",
-    "dbx-tools-core==0.6.89",
-    "psycopg[binary]>=3.2.9,<4",
-    "sqlalchemy>=2.0.41,<3",
+    "asyncpg>=0.30",
+    "databricks-sdk>=0.63.0",
+    "greenlet>=3.2",
+    "dbx-tools-core==0.6.91",
+    "psycopg[binary]>=3.2.9",
+    "sqlalchemy>=2.0.41",
 ]
 
 [project.urls]

{dbx_tools_postgres-0.6.89 → dbx_tools_postgres-0.6.91}/pyproject.toml.orig

@@ -2,17 +2,17 @@
 
 [project]
 name = "dbx-tools-postgres"
-version = "0.6.89"
+version = "0.6.91"
 description = "WorkspaceClient-backed Lakebase Postgres resolution, SQLAlchemy engines, advisory locks, and LISTEN/NOTIFY topic bus"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-  "asyncpg>=0.30,<1",
-  "databricks-sdk>=0.63.0,<1",
-  "greenlet>=3.2,<4",
-  "dbx-tools-core==0.6.89",
-  "psycopg[binary]>=3.2.9,<4",
-  "sqlalchemy>=2.0.41,<3"
+  "asyncpg>=0.30",
+  "databricks-sdk>=0.63.0",
+  "greenlet>=3.2",
+  "dbx-tools-core==0.6.91",
+  "psycopg[binary]>=3.2.9",
+  "sqlalchemy>=2.0.41"
 ]
 
   [project.urls]

{dbx_tools_postgres-0.6.89 → dbx_tools_postgres-0.6.91}/src/dbx_tools/postgres/engine.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import datetime as dt
 import os
+import threading
 import uuid
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
@@ -13,9 +15,15 @@ from sqlalchemy.ext.asyncio import create_async_engine as sqlalchemy_create_asyn
 
 from .address import SSL_MODES, ParsedAddress, SslMode, parse_address, parse_resource_path
 
+"""Lakebase connection resolution and connect-time credential injection."""
+
 CredentialProvider = Callable[[], str]
+CredentialLoader = Callable[[], tuple[str, dt.datetime | None]]
 
 _API_BASE = "/api/2.0/postgres"
+_CREDENTIAL_REFRESH_LEAD = dt.timedelta(minutes=5)
+_DEFAULT_CREDENTIAL_LIFETIME = dt.timedelta(minutes=50)
+_MINIMUM_CREDENTIAL_LIFETIME = dt.timedelta(minutes=1)
 _DEFAULT_DATABASE = "databricks_postgres"
 _DEFAULT_PORT = 5432
 _DEFAULT_SSL_MODE: SslMode = "require"
@@ -72,6 +80,28 @@ class ResolvedPostgresConnection:
         )
 
 
+class _CachedCredentialProvider:
+    """Cache one credential and serialize refreshes within the Python process."""
+
+    def __init__(self, load: CredentialLoader) -> None:
+        self._load = load
+        self._lock = threading.RLock()
+        self._token: str | None = None
+        self._renew_at = dt.datetime.min.replace(tzinfo=dt.timezone.utc)
+
+    def __call__(self) -> str:
+        token = self._token
+        if token is not None and _utcnow() < self._renew_at:
+            return token
+        with self._lock:
+            if self._token is None or _utcnow() >= self._renew_at:
+                token, expiration = self._load()
+                self._token = token
+                self._renew_at = _credential_renewal(expiration)
+            assert self._token is not None
+            return self._token
+
+
 def workspace_credential_provider(
     workspace_client: WorkspaceClientLike,
     instance_name: str,
@@ -79,7 +109,7 @@ def workspace_credential_provider(
     if not instance_name.strip():
         raise ValueError("instance_name must not be empty")
 
-    def provide() -> str:
+    def load() -> tuple[str, dt.datetime | None]:
         credential = workspace_client.database.generate_database_credential(
             request_id=str(uuid.uuid4()),
             instance_names=[instance_name],
@@ -87,9 +117,9 @@ def workspace_credential_provider(
         token = getattr(credential, "token", None)
         if not isinstance(token, str) or not token:
             raise RuntimeError("WorkspaceClient returned no Lakebase database credential token")
-        return token
+        return token, _credential_expiration(getattr(credential, "expiration_time", None))
 
-    return provide
+    return _CachedCredentialProvider(load)
 
 
 def autoscaling_credential_provider(
@@ -99,7 +129,7 @@ def autoscaling_credential_provider(
     if not endpoint.strip():
         raise ValueError("endpoint must not be empty")
 
-    def provide() -> str:
+    def load() -> tuple[str, dt.datetime | None]:
         credential = workspace_client.api_client.do(
             "POST",
             f"{_API_BASE}/credentials",
@@ -109,9 +139,9 @@ def autoscaling_credential_provider(
         token = credential.get("token") if isinstance(credential, dict) else None
         if not isinstance(token, str) or not token:
             raise RuntimeError("WorkspaceClient returned no Lakebase database credential token")
-        return token
+        return token, _mapping_credential_expiration(credential)
 
-    return provide
+    return _CachedCredentialProvider(load)
 
 
 def install_credential_injection(engine: Engine, provider: CredentialProvider) -> None:
@@ -227,6 +257,47 @@ def resolve_postgres_connection(
     )
 
 
+def _credential_renewal(expiration: dt.datetime | None) -> dt.datetime:
+    now = _utcnow()
+    if expiration is None:
+        return now + _DEFAULT_CREDENTIAL_LIFETIME
+    normalized = (
+        expiration.astimezone(dt.timezone.utc)
+        if expiration.tzinfo is not None
+        else expiration.replace(tzinfo=dt.timezone.utc)
+    )
+    return max(normalized - _CREDENTIAL_REFRESH_LEAD, now + _MINIMUM_CREDENTIAL_LIFETIME)
+
+
+def _credential_expiration(value: object) -> dt.datetime | None:
+    if isinstance(value, dt.datetime):
+        return value
+    if isinstance(value, (int, float)):
+        try:
+            return dt.datetime.fromtimestamp(value, tz=dt.timezone.utc)
+        except (OSError, OverflowError, ValueError):
+            return None
+    if not isinstance(value, str) or not value.strip():
+        return None
+    normalized = value.strip().replace("Z", "+00:00")
+    try:
+        return dt.datetime.fromisoformat(normalized)
+    except ValueError:
+        return None
+
+
+def _mapping_credential_expiration(credential: Mapping[str, object]) -> dt.datetime | None:
+    for key in ("expiration_time", "expirationTime", "expires_at", "expiresAt"):
+        expiration = _credential_expiration(credential.get(key))
+        if expiration is not None:
+            return expiration
+    return None
+
+
+def _utcnow() -> dt.datetime:
+    return dt.datetime.now(tz=dt.timezone.utc)
+
+
 def _default_provider(
     workspace_client: WorkspaceClientLike,
     resolved: ResolvedPostgresConnection,