From c27248b0515fd37bd135c62f29ad5cfee1b885ae Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:00:43 +0200 Subject: [PATCH 1/7] Add Databricks Lakebase Entra credential provider --- pgdevkit/lakebase.py | 86 +++++++++++++++++++++++++++++++++ tests/test_lakebase.py | 106 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+) create mode 100644 pgdevkit/lakebase.py create mode 100644 tests/test_lakebase.py diff --git a/pgdevkit/lakebase.py b/pgdevkit/lakebase.py new file mode 100644 index 0000000..3fb3d14 --- /dev/null +++ b/pgdevkit/lakebase.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import datetime +import json +import time +import uuid +from urllib import request as urllib_request +from urllib.error import HTTPError + +_DATABRICKS_RESOURCE_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default" +_REFRESH_MARGIN_SECONDS = 300 + +_credential: object | None = None +_credential_cache: dict[tuple[str, str], tuple[str, float]] = {} +_dns_cache: dict[tuple[str, str], str] = {} + + +def _databricks_resource_token() -> str: + try: + from azure.identity import DefaultAzureCredential + except ImportError: + raise ImportError("Install azure-identity extra: pip install pgdevkit[azure]") + global _credential + if _credential is None: + _credential = DefaultAzureCredential() + return _credential.get_token(_DATABRICKS_RESOURCE_SCOPE).token + + +def _request_json(url: str, token: str, *, body: dict | None = None) -> dict: + data = json.dumps(body).encode("utf-8") if body is not None else None + req = urllib_request.Request( + url, + data=data, + method="POST" if body is not None else "GET", + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + try: + with urllib_request.urlopen(req, timeout=10) as resp: + return json.loads(resp.read().decode("utf-8")) + except HTTPError as e: + detail = e.read().decode("utf-8", errors="replace") + raise RuntimeError(f"Databricks API request failed [{e.code}]: {detail}") from e + + +def _parse_expiration(expiration_time: str) -> float: + return datetime.datetime.fromisoformat(expiration_time).timestamp() + + +def get_lakebase_password(workspace_host: str, instance_name: str) -> str: + """Return a valid short-lived Postgres password for the given Lakebase + instance, fetching/refreshing the credential as needed.""" + key = (workspace_host, instance_name) + cached = _credential_cache.get(key) + if cached is not None: + token, expires_at = cached + if time.time() < expires_at - _REFRESH_MARGIN_SECONDS: + return token + + entra_token = _databricks_resource_token() + response = _request_json( + f"{workspace_host.rstrip('/')}/api/2.0/database/credentials", + entra_token, + body={"instance_names": [instance_name], "request_id": str(uuid.uuid4())}, + ) + token = response["token"] + expires_at = _parse_expiration(response["expiration_time"]) + _credential_cache[key] = (token, expires_at) + return token + + +def resolve_read_write_dns(workspace_host: str, instance_name: str) -> str: + """Resolve and cache the Postgres read/write DNS endpoint for a Lakebase instance.""" + key = (workspace_host, instance_name) + if key in _dns_cache: + return _dns_cache[key] + entra_token = _databricks_resource_token() + response = _request_json( + f"{workspace_host.rstrip('/')}/api/2.0/database/instances/{instance_name}", + entra_token, + ) + dns = response["read_write_dns"] + _dns_cache[key] = dns + return dns diff --git a/tests/test_lakebase.py b/tests/test_lakebase.py new file mode 100644 index 0000000..0830591 --- /dev/null +++ b/tests/test_lakebase.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import json +from io import BytesIO +from urllib.error import HTTPError + +import pytest + +import pgdevkit.lakebase as lakebase + + +@pytest.fixture(autouse=True) +def _clear_caches(): + lakebase._credential_cache.clear() + lakebase._dns_cache.clear() + yield + lakebase._credential_cache.clear() + lakebase._dns_cache.clear() + + +class _FakeResponse: + def __init__(self, payload: dict) -> None: + self._payload = payload + + def read(self) -> bytes: + return json.dumps(self._payload).encode("utf-8") + + def __enter__(self) -> "_FakeResponse": + return self + + def __exit__(self, *exc) -> None: + return None + + +def test_get_lakebase_password_fetches_and_caches(monkeypatch): + monkeypatch.setattr(lakebase, "_databricks_resource_token", lambda: "ENTRA_TOKEN") + calls = [] + + def fake_urlopen(req, timeout=10): + calls.append(req) + return _FakeResponse({"token": "PGTOKEN", "expiration_time": "2099-01-01T00:00:00Z"}) + + monkeypatch.setattr(lakebase.urllib_request, "urlopen", fake_urlopen) + + token = lakebase.get_lakebase_password("https://adb-123.azuredatabricks.net", "myinstance") + assert token == "PGTOKEN" + assert len(calls) == 1 + assert calls[0].full_url == "https://adb-123.azuredatabricks.net/api/2.0/database/credentials" + assert calls[0].get_header("Authorization") == "Bearer ENTRA_TOKEN" + body = json.loads(calls[0].data.decode("utf-8")) + assert body["instance_names"] == ["myinstance"] + + # Second call within validity window must use the cache, not call urlopen again. + token2 = lakebase.get_lakebase_password("https://adb-123.azuredatabricks.net", "myinstance") + assert token2 == "PGTOKEN" + assert len(calls) == 1 + + +def test_get_lakebase_password_refreshes_when_near_expiry(monkeypatch): + monkeypatch.setattr(lakebase, "_databricks_resource_token", lambda: "ENTRA_TOKEN") + responses = iter( + [ + {"token": "FIRST", "expiration_time": "2020-01-01T00:00:01Z"}, + {"token": "SECOND", "expiration_time": "2099-01-01T00:00:00Z"}, + ] + ) + + def fake_urlopen(req, timeout=10): + return _FakeResponse(next(responses)) + + monkeypatch.setattr(lakebase.urllib_request, "urlopen", fake_urlopen) + + first = lakebase.get_lakebase_password("https://adb-123.azuredatabricks.net", "myinstance") + assert first == "FIRST" + second = lakebase.get_lakebase_password("https://adb-123.azuredatabricks.net", "myinstance") + assert second == "SECOND" + + +def test_get_lakebase_password_raises_on_http_error(monkeypatch): + monkeypatch.setattr(lakebase, "_databricks_resource_token", lambda: "ENTRA_TOKEN") + + def fake_urlopen(req, timeout=10): + raise HTTPError(req.full_url, 403, "Forbidden", {}, BytesIO(b'{"message": "no access"}')) + + monkeypatch.setattr(lakebase.urllib_request, "urlopen", fake_urlopen) + + with pytest.raises(RuntimeError, match="403"): + lakebase.get_lakebase_password("https://adb-123.azuredatabricks.net", "myinstance") + + +def test_resolve_read_write_dns_caches(monkeypatch): + monkeypatch.setattr(lakebase, "_databricks_resource_token", lambda: "ENTRA_TOKEN") + calls = [] + + def fake_urlopen(req, timeout=10): + calls.append(req) + return _FakeResponse({"read_write_dns": "instance-abc.database.azuredatabricks.net"}) + + monkeypatch.setattr(lakebase.urllib_request, "urlopen", fake_urlopen) + + dns = lakebase.resolve_read_write_dns("https://adb-123.azuredatabricks.net", "myinstance") + assert dns == "instance-abc.database.azuredatabricks.net" + assert calls[0].full_url == "https://adb-123.azuredatabricks.net/api/2.0/database/instances/myinstance" + + lakebase.resolve_read_write_dns("https://adb-123.azuredatabricks.net", "myinstance") + assert len(calls) == 1 # cached, no second request From e5ae592e25660a439be901aa8550cd49f9a7f2c2 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:04:16 +0200 Subject: [PATCH 2/7] Detect Lakebase vs Azure Postgres hosts, dispatch build_conninfo accordingly --- pgdevkit/connection.py | 43 +++++++++++++++++++++++++---- tests/test_connection.py | 58 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 5 deletions(-) create mode 100644 tests/test_connection.py diff --git a/pgdevkit/connection.py b/pgdevkit/connection.py index 1e03a9c..76574d6 100644 --- a/pgdevkit/connection.py +++ b/pgdevkit/connection.py @@ -1,8 +1,23 @@ from __future__ import annotations -from urllib.parse import urlparse, urlunparse, quote +from typing import Literal +from urllib.parse import quote, urlparse, urlunparse -def _get_entra_token() -> str: +_LAKEBASE_HOST_SUFFIXES = ( + ".database.azuredatabricks.net", + ".database.cloud.databricks.com", +) + + +def detect_provider(host: str) -> Literal["azure_postgres", "databricks_lakebase"]: + """Classify a Postgres hostname: Databricks Lakebase (needs credential + exchange) or the default Azure Postgres AAD token flow.""" + if any(host.endswith(suffix) for suffix in _LAKEBASE_HOST_SUFFIXES): + return "databricks_lakebase" + return "azure_postgres" + + +def get_azure_postgres_password() -> str: try: from azure.identity import DefaultAzureCredential except ImportError: @@ -12,12 +27,30 @@ def _get_entra_token() -> str: return token.token -def build_conninfo(url: str, entra_user: str | None = None) -> str: +def build_conninfo( + url: str, + entra_user: str | None = None, + *, + databricks_workspace_host: str | None = None, + databricks_instance: str | None = None, +) -> str: if entra_user is None: return url - token = _get_entra_token() + parsed = urlparse(url) host = parsed.hostname or "" port = f":{parsed.port}" if parsed.port else "" - netloc = f"{quote(entra_user, safe='')}:{quote(token, safe='')}@{host}{port}" + + if detect_provider(host) == "databricks_lakebase": + if not databricks_workspace_host or not databricks_instance: + raise ValueError( + "Lakebase host detected — pass --databricks-workspace-host and --databricks-instance" + ) + from .lakebase import get_lakebase_password + + password = get_lakebase_password(databricks_workspace_host, databricks_instance) + else: + password = get_azure_postgres_password() + + netloc = f"{quote(entra_user, safe='')}:{quote(password, safe='')}@{host}{port}" return urlunparse(parsed._replace(netloc=netloc)) diff --git a/tests/test_connection.py b/tests/test_connection.py new file mode 100644 index 0000000..0e6b5bb --- /dev/null +++ b/tests/test_connection.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import pytest + +from pgdevkit.connection import build_conninfo, detect_provider + + +@pytest.mark.parametrize( + "host,expected", + [ + ("myserver.postgres.database.azure.com", "azure_postgres"), + ("myserver.postgres.cosmos.azure.com", "azure_postgres"), + ("instance-abc123.database.azuredatabricks.net", "databricks_lakebase"), + ("instance-abc123.database.cloud.databricks.com", "databricks_lakebase"), + ("localhost", "azure_postgres"), + ("db.internal.example.com", "azure_postgres"), + ], +) +def test_detect_provider(host, expected): + assert detect_provider(host) == expected + + +def test_build_conninfo_without_entra_user_returns_url_unchanged(): + url = "postgresql://user:pass@host:5432/db" + assert build_conninfo(url) == url + + +def test_build_conninfo_azure_postgres_uses_token_as_password(monkeypatch): + monkeypatch.setattr("pgdevkit.connection.get_azure_postgres_password", lambda: "TOKEN123") + conninfo = build_conninfo( + "postgresql://myserver.postgres.database.azure.com:5432/db", "alice@example.com" + ) + assert conninfo == "postgresql://alice%40example.com:TOKEN123@myserver.postgres.database.azure.com:5432/db" + + +def test_build_conninfo_lakebase_uses_credential_exchange(monkeypatch): + monkeypatch.setattr( + "pgdevkit.lakebase.get_lakebase_password", + lambda workspace_host, instance_name: "LAKEBASE_TOKEN", + ) + conninfo = build_conninfo( + "postgresql://instance-abc.database.azuredatabricks.net:5432/databricks_postgres", + "alice@example.com", + databricks_workspace_host="https://adb-123.azuredatabricks.net", + databricks_instance="myinstance", + ) + assert conninfo == ( + "postgresql://alice%40example.com:LAKEBASE_TOKEN" + "@instance-abc.database.azuredatabricks.net:5432/databricks_postgres" + ) + + +def test_build_conninfo_lakebase_missing_flags_raises(): + with pytest.raises(ValueError, match="databricks-workspace-host"): + build_conninfo( + "postgresql://instance-abc.database.azuredatabricks.net:5432/db", + "alice@example.com", + ) From 1af175955de91ac13b054a67b3b40f291c13039b Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:06:42 +0200 Subject: [PATCH 3/7] Add --databricks-workspace-host/--databricks-instance to pgdb compare --- pgdevkit/cli.py | 19 +++++++++++- tests/test_cli_compare.py | 64 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+), 1 deletion(-) create mode 100644 tests/test_cli_compare.py diff --git a/pgdevkit/cli.py b/pgdevkit/cli.py index d5baa8e..4d60856 100644 --- a/pgdevkit/cli.py +++ b/pgdevkit/cli.py @@ -28,6 +28,14 @@ def compare( url: str = typer.Option(..., "--url", help="PostgreSQL DSN (postgresql://user:pass@host:port/db)"), entra_user: str | None = typer.Option(None, "--entra-user", help="Azure Entra user (triggers token auth)"), + databricks_workspace_host: str | None = typer.Option( + None, + "--databricks-workspace-host", + help="Databricks workspace URL, e.g. https://adb-....azuredatabricks.net (required for Lakebase hosts)", + ), + databricks_instance: str | None = typer.Option( + None, "--databricks-instance", help="Lakebase instance name (required for Lakebase hosts)" + ), report_extra_db: bool = typer.Option(False, "--report-extra-db", help="Report objects in DB but not in scripts"), scripts_dir: Path = typer.Argument(..., help="Directory containing SQL scripts"), ) -> None: @@ -36,7 +44,16 @@ def compare( err_console.print(f"[red]Error:[/red] {scripts_dir} is not a directory") raise typer.Exit(2) - conninfo = build_conninfo(url, entra_user) + try: + conninfo = build_conninfo( + url, + entra_user, + databricks_workspace_host=databricks_workspace_host, + databricks_instance=databricks_instance, + ) + except ValueError as e: + err_console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(2) with console.status("Parsing SQL scripts..."): scripts_schema = parse_directory(scripts_dir) diff --git a/tests/test_cli_compare.py b/tests/test_cli_compare.py new file mode 100644 index 0000000..46d4b8c --- /dev/null +++ b/tests/test_cli_compare.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from pathlib import Path + +from typer.testing import CliRunner + +import pgdevkit.connection as connection +from pgdevkit.cli import app + +runner = CliRunner() + + +def test_compare_lakebase_host_without_databricks_flags_fails_fast(tmp_path: Path): + (tmp_path / "empty").mkdir() + result = runner.invoke( + app, + [ + "compare", + "--url", + "postgresql://user@instance-abc.database.azuredatabricks.net:5432/db", + "--entra-user", + "alice@example.com", + str(tmp_path / "empty"), + ], + ) + assert result.exit_code == 2 + assert "--databricks-workspace-host" in result.output + + +def test_compare_lakebase_host_with_flags_reaches_introspection(tmp_path: Path, monkeypatch): + (tmp_path / "empty").mkdir() + monkeypatch.setattr(connection, "get_azure_postgres_password", lambda: "unused") + + calls = [] + + def fake_get_lakebase_password(workspace_host, instance_name): + calls.append((workspace_host, instance_name)) + return "LAKEBASE_TOKEN" + + monkeypatch.setattr("pgdevkit.lakebase.get_lakebase_password", fake_get_lakebase_password) + + def fake_introspect_db(conninfo): + assert "LAKEBASE_TOKEN" in conninfo + raise RuntimeError("stop after conninfo built — introspection itself isn't under test here") + + monkeypatch.setattr("pgdevkit.cli.introspect_db", fake_introspect_db) + + result = runner.invoke( + app, + [ + "compare", + "--url", + "postgresql://user@instance-abc.database.azuredatabricks.net:5432/db", + "--entra-user", + "alice@example.com", + "--databricks-workspace-host", + "https://adb-123.azuredatabricks.net", + "--databricks-instance", + "myinstance", + str(tmp_path / "empty"), + ], + ) + assert calls == [("https://adb-123.azuredatabricks.net", "myinstance")] + assert result.exit_code == 1 # CliRunner captures the exception as a non-zero exit From ce2249140562065f818393c9e59a73d749e77b7b Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:09:22 +0200 Subject: [PATCH 4/7] Add entra_user support to PgPool for Azure Postgres and Lakebase --- pgdevkit/db/connection.py | 48 +++++++++++++++++++++++------- tests/db/test_connection.py | 58 +++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 10 deletions(-) create mode 100644 tests/db/test_connection.py diff --git a/pgdevkit/db/connection.py b/pgdevkit/db/connection.py index c0d63a6..a2b104b 100644 --- a/pgdevkit/db/connection.py +++ b/pgdevkit/db/connection.py @@ -1,28 +1,56 @@ from __future__ import annotations +import asyncio import os from psycopg_pool import AsyncConnectionPool +from ..connection import detect_provider, get_azure_postgres_password +from ..lakebase import get_lakebase_password + class PgPool: """A connection pool keyed off `{env_prefix}HOST/PORT/DB/USER/PASSWORD` - environment variables (e.g. env_prefix="MDM_POSTGRES_").""" + environment variables (e.g. env_prefix="MDM_POSTGRES_"). + + Pass `entra_user` to authenticate via Entra ID instead of a static + password. The Postgres host is inspected to pick the Azure Postgres AAD + token flow or Databricks Lakebase credential exchange; for the latter + also set `{env_prefix}DATABRICKS_WORKSPACE_HOST` and + `{env_prefix}DATABRICKS_INSTANCE`. + """ - def __init__(self, env_prefix: str = "POSTGRES_", max_size: int = 40) -> None: + def __init__( + self, + env_prefix: str = "POSTGRES_", + max_size: int = 40, + *, + entra_user: str | None = None, + ) -> None: self._env_prefix = env_prefix self._max_size = max_size + self._entra_user = entra_user self._pool: AsyncConnectionPool | None = None - def _dsn(self) -> str: + async def _dsn(self) -> str: p = self._env_prefix - return ( - f"host={os.environ[p + 'HOST']} " - f"port={os.environ[p + 'PORT']} " - f"dbname={os.environ[p + 'DB']} " - f"user={os.environ[p + 'USER']} " - f"password={os.environ[p + 'PASSWORD']}" - ) + host = os.environ[p + "HOST"] + port = os.environ[p + "PORT"] + dbname = os.environ[p + "DB"] + + if self._entra_user is None: + user = os.environ[p + "USER"] + password = os.environ[p + "PASSWORD"] + else: + user = self._entra_user + if detect_provider(host) == "databricks_lakebase": + workspace_host = os.environ[p + "DATABRICKS_WORKSPACE_HOST"] + instance_name = os.environ[p + "DATABRICKS_INSTANCE"] + password = await asyncio.to_thread(get_lakebase_password, workspace_host, instance_name) + else: + password = await asyncio.to_thread(get_azure_postgres_password) + + return f"host={host} port={port} dbname={dbname} user={user} password={password}" async def open(self) -> None: if self._pool is None: diff --git a/tests/db/test_connection.py b/tests/db/test_connection.py new file mode 100644 index 0000000..2509eea --- /dev/null +++ b/tests/db/test_connection.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import pytest + +from pgdevkit.db.connection import PgPool + + +ENV_PREFIX = "PGDEVKIT_CONNTEST_" + + +async def test_dsn_static_password_when_entra_user_unset(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "localhost") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setenv(f"{ENV_PREFIX}USER", "myuser") + monkeypatch.setenv(f"{ENV_PREFIX}PASSWORD", "mypassword") + + pool = PgPool(env_prefix=ENV_PREFIX) + dsn = await pool._dsn() + assert dsn == "host=localhost port=5432 dbname=mydb user=myuser password=mypassword" + + +async def test_dsn_azure_postgres_entra(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "myserver.postgres.database.azure.com") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "mydb") + monkeypatch.setattr("pgdevkit.db.connection.get_azure_postgres_password", lambda: "AADTOKEN") + + pool = PgPool(env_prefix=ENV_PREFIX, entra_user="alice@example.com") + dsn = await pool._dsn() + assert dsn == ( + "host=myserver.postgres.database.azure.com port=5432 dbname=mydb " + "user=alice@example.com password=AADTOKEN" + ) + + +async def test_dsn_databricks_lakebase_entra(monkeypatch): + monkeypatch.setenv(f"{ENV_PREFIX}HOST", "instance-abc.database.azuredatabricks.net") + monkeypatch.setenv(f"{ENV_PREFIX}PORT", "5432") + monkeypatch.setenv(f"{ENV_PREFIX}DB", "databricks_postgres") + monkeypatch.setenv(f"{ENV_PREFIX}DATABRICKS_WORKSPACE_HOST", "https://adb-123.azuredatabricks.net") + monkeypatch.setenv(f"{ENV_PREFIX}DATABRICKS_INSTANCE", "myinstance") + + calls = [] + + def fake_get_lakebase_password(workspace_host, instance_name): + calls.append((workspace_host, instance_name)) + return "LAKEBASE_TOKEN" + + monkeypatch.setattr("pgdevkit.db.connection.get_lakebase_password", fake_get_lakebase_password) + + pool = PgPool(env_prefix=ENV_PREFIX, entra_user="alice@example.com") + dsn = await pool._dsn() + assert dsn == ( + "host=instance-abc.database.azuredatabricks.net port=5432 dbname=databricks_postgres " + "user=alice@example.com password=LAKEBASE_TOKEN" + ) + assert calls == [("https://adb-123.azuredatabricks.net", "myinstance")] From 6db3e4473b4dbc8cb2d7ce317d13407beab53dce Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:12:19 +0200 Subject: [PATCH 5/7] Document Entra auth for Azure Postgres and Databricks Lakebase --- README.md | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/README.md b/README.md index 3c62aef..1bec690 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,33 @@ convention) against a live database and report differences: pgdb compare --url postgresql://user:pass@host:port/db path/to/database/ ``` +### Entra ID auth (Azure Postgres / Databricks Lakebase) + +Pass `--entra-user ` to `pgdb compare` to authenticate with an +Entra ID token instead of a static password. Which token flow is used is +auto-detected from the database hostname: + +- **Azure Database for PostgreSQL** (`*.postgres.database.azure.com`, + `*.postgres.cosmos.azure.com`) — the default: fetches a token via + `DefaultAzureCredential` and uses it directly as the password. Requires + the `azure` extra: `pip install pgdevkit[azure]`. +- **Databricks Lakebase** (`*.database.azuredatabricks.net`, + `*.database.cloud.databricks.com`) — fetches a Databricks-scoped Entra + token, then exchanges it for a short-lived Postgres credential via the + Databricks workspace API. Also requires `--databricks-workspace-host` + and `--databricks-instance`: + +```bash +pgdb compare --url postgresql://instance-abc.database.azuredatabricks.net:5432/databricks_postgres \ + --entra-user alice@example.com \ + --databricks-workspace-host https://adb-123456789.azuredatabricks.net \ + --databricks-instance myinstance \ + path/to/database/ +``` + +(`--url`'s own user/password, if any, are discarded and replaced — `--entra-user` +plus the fetched token become the connection's actual credentials.) + ## `pgdb testdb` Manages a single shared, Podman-backed Postgres container for local tests @@ -59,6 +86,11 @@ Install with the `db` extra: `pip install pgdevkit[db]`. - **`PgPool`** — an async connection pool keyed off `{env_prefix}HOST/PORT/DB/USER/PASSWORD` env vars. Call `await pool.open()` once at startup, then use `async with pool.connection() as con:`. + Pass `entra_user` to authenticate via Entra ID instead of a static + password — same host-based auto-detection as `pgdb compare`'s + `--entra-user`. For Lakebase hosts, also set the + `{env_prefix}DATABRICKS_WORKSPACE_HOST` and `{env_prefix}DATABRICKS_INSTANCE` + env vars. - **CRUD functions** — `pg_retrieve`, `pg_retrieve_many`, `pg_insert`, `pg_insert_many`, `pg_update`, `pg_update_dict`, `pg_upsert`, `pg_upsert_dict`, `pg_upsert_many`, `pg_upsert_many_dict`, `pg_delete`, From ac976ba6aa4cc168818d95cee60f3b13dcfb1186 Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:23:52 +0200 Subject: [PATCH 6/7] Bump psycopg-pool floor to 3.3.0 for callable conninfo support Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01GZmfryNuuAYSy38WUfva7N --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 381b84f..20219e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ cli = [ "typer>=0.26.7", ] db = [ - "psycopg-pool>=3.2.0", + "psycopg-pool>=3.3.0", "pydantic>=2.0", ] diff --git a/uv.lock b/uv.lock index 4e240dc..27a8012 100644 --- a/uv.lock +++ b/uv.lock @@ -302,7 +302,7 @@ test = [ requires-dist = [ { name = "azure-identity", marker = "extra == 'azure'", specifier = ">=1.19.0" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.2.0" }, - { name = "psycopg-pool", marker = "extra == 'db'", specifier = ">=3.2.0" }, + { name = "psycopg-pool", marker = "extra == 'db'", specifier = ">=3.3.0" }, { name = "pydantic", marker = "extra == 'db'", specifier = ">=2.0" }, { name = "rich", marker = "extra == 'cli'", specifier = ">=13.0.0" }, { name = "sqlglot", extras = ["c"], specifier = ">=30.11.0" }, From bc4d3e5f2e82e73c25885c9630c1392ca6572dbc Mon Sep 17 00:00:00 2001 From: Adrian Ehrsam Date: Thu, 16 Jul 2026 11:31:29 +0200 Subject: [PATCH 7/7] Fix ty check: type _credential as DefaultAzureCredential | None ty couldn't see .get_token() on a lazily-created object typed as `object | None`. Use a TYPE_CHECKING-guarded import so the annotation is accurate without making azure-identity an unconditional import. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01GZmfryNuuAYSy38WUfva7N --- pgdevkit/lakebase.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pgdevkit/lakebase.py b/pgdevkit/lakebase.py index 3fb3d14..f5d4eb6 100644 --- a/pgdevkit/lakebase.py +++ b/pgdevkit/lakebase.py @@ -4,13 +4,17 @@ import json import time import uuid +from typing import TYPE_CHECKING from urllib import request as urllib_request from urllib.error import HTTPError +if TYPE_CHECKING: + from azure.identity import DefaultAzureCredential + _DATABRICKS_RESOURCE_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default" _REFRESH_MARGIN_SECONDS = 300 -_credential: object | None = None +_credential: DefaultAzureCredential | None = None _credential_cache: dict[tuple[str, str], tuple[str, float]] = {} _dns_cache: dict[tuple[str, str], str] = {}