Lakebase entra auth

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 <identity>` 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`,

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)

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))

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:

pgdevkit/lakebase.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import datetime
+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: DefaultAzureCredential | 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

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",
 ]
 

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")]