databricks · sysopmatt · Jun 8, 2026
diff --git a/fastapi-postgres-app/.env.example b/fastapi-postgres-app/.env.example
@@ -0,0 +1,46 @@
+# ============================================
+# This file is for LOCAL DEV ONLY
+# ============================================
+# Deployed apps don't need a .env file — the resource binding auto-injects
+# PG* and app.yaml's `valueFrom: postgres` populates PGENDPOINT.
+#
+# To run locally: copy to `.env` and fill in the TODOs below. All TODOs
+# are required for local.
+
+# ============================================
+# Lakebase endpoint resource name (local: required)
+# ============================================
+# Format: projects/<project-id>/branches/<branch-id>/endpoints/<endpoint-id>
+# Find via:
+#   databricks api get /api/2.0/postgres/projects/<project>/branches/<branch>/endpoints
+#
+# TODO: set PGENDPOINT
+# PGENDPOINT=projects/my-project/branches/main/endpoints/primary
+
+# ============================================
+# Postgres connection (local: required)
+# ============================================
+# Find values in the Lakebase Connect modal → Parameters only.
+#
+# TODO: set PGHOST
+# PGHOST=ep-xxx-yyy.database.eastus.azuredatabricks.net
+#
+# TODO: set PGUSER — your identity email (must be registered as a Lakebase
+# role on the target project)
+# PGUSER=you@domain.com
+#
+# TODO: set PGDATABASE
+# PGDATABASE=my_database
+#
+# Defaults — usually no change needed:
+# PGPORT=5432
+# PGSSLMODE=require
+
+# ============================================
+# Connection pool tuning (always optional)
+# ============================================
+# DB_POOL_SIZE=5
+# DB_MAX_OVERFLOW=10
+# DB_COMMAND_TIMEOUT=30
+# DB_POOL_TIMEOUT=10
+# DB_POOL_RECYCLE_INTERVAL=3600
diff --git a/fastapi-postgres-app/README.md b/fastapi-postgres-app/README.md
@@ -0,0 +1,104 @@
+# Lakebase Autoscaling FastAPI app
+
+A minimal FastAPI app that connects to a Lakebase Autoscaling (Postgres) database. Async stack (SQLAlchemy `AsyncEngine` + `asyncpg`), background OAuth token refresh every 50 minutes, no front-end. Pairs with the [Databricks Apps Cookbook FastAPI + Lakebase recipe](https://apps-cookbook.dev/docs/fastapi/getting_started/lakebase_connection), adapted for Autoscaling.
+
+## Layout
+
+```
+fastapi-postgres-app/
+├── main.py                   FastAPI entrypoint + lifespan
+├── config/database.py        Lakebase connection + token refresh
+├── routes/health.py          /healthz, /api/health
+├── manifest.yaml             Resource binding (postgres_spec)
+├── app.yaml                  Apps deploy config (uvicorn + PGENDPOINT valueFrom)
+├── requirements.txt
+└── .env.example              Local dev only
+```
+
+## Routes
+
+| Method | Path | Description |
+|---|---|---|
+| GET | `/` | Hello world |
+| GET | `/docs` | OpenAPI / Swagger UI |
+| GET | `/healthz` | Cheap liveness probe (no DB) |
+| GET | `/api/health` | Round-trips to Lakebase |
+
+## Deploy to Databricks Apps
+
+1. **Create the app** in your workspace.
+2. **Bind the Lakebase resource.** App → **Settings → Resources → Add resource**:
+   - **Database**: pick your Lakebase Autoscaling project, branch, and database
+   - **Permission**: `Can connect and create`
+   - **Resource key**: `postgres` (must match `valueFrom: postgres` in `app.yaml` — change either side if you prefer a different key)
+3. **Deploy** (from a connected Git repo or `databricks apps deploy`).
+
+The resource binding auto-injects `PGHOST` / `PGPORT` / `PGDATABASE` / `PGUSER` / `PGSSLMODE` / `PGAPPNAME` and `app.yaml`'s `valueFrom` populates `PGENDPOINT`. Nothing to set manually.
+
+After the app starts, hit `/api/health` to confirm: `{"status": "ok", "database": true}`.
+
+## Testing the API with Postman (service principal auth)
+
+Databricks Apps require an OAuth Bearer token on every request. Easiest way to exercise the API from Postman is with a service principal and Client Credentials grant — Postman fetches and refreshes the token automatically.
+
+### One-time setup
+
+You'll need a workspace admin (or anyone with permission to create service principals) to do this once.
+
+1. **Create a service principal.**
+   Workspace → **Settings → Identity and access → Service principals → Add service principal**.
+
+2. **Generate an OAuth client secret.** Open the SP → **Secrets** tab → **Generate secret**. Save both the **client ID** (the SP's application ID) and the **client secret** — the secret is shown only once.
+
+3. **Grant the SP `Can Use` on the app.** Open the app in the workspace → **Permissions** tab → add the service principal with `Can Use`.
+
+### Postman config
+
+In your request, open the **Authorization** tab and configure:
+
+| Field | Value |
+|---|---|
+| Type | OAuth 2.0 |
+| Grant Type | Client Credentials |
+| Access Token URL | `https://<your-workspace-host>/oidc/v1/token` |
+| Client ID | `<sp-client-id>` |
+| Client Secret | `<sp-client-secret>` |
+| Scope | `all-apis` |
+| Client Authentication | Send as Basic Auth header |
+
+Your **workspace host** is what's in the browser address bar when you're logged into Databricks (everything between `https://` and the first slash). It's **not** the app URL — the token endpoint lives on the workspace itself.
+
+Click **Get New Access Token**, then the orange **Use Token** button. Postman fetches a JWT (starts with `ey...`) and applies it as a Bearer token. Tokens last about an hour; Postman re-fetches as needed.
+
+Point the request at your app's public URL, e.g. `GET https://<app-name>-<workspace-id>.cloud.databricksapps.com/api/health`.
+
+## Local dev
+
+```bash
+cp .env.example .env
+# fill in the TODOs (see comments in .env.example)
+uv pip install -r requirements.txt
+uvicorn main:app --reload --port 8000
+```
+
+Local dev requires your identity to be registered as a Lakebase role on the target project (one-time; the deployed app's SP gets auto-registered by the resource binding):
+
+```sql
+CREATE EXTENSION IF NOT EXISTS databricks_auth;
+SELECT databricks_create_role('you@domain.com', 'USER');
+```
+
+## Caveats
+
+- **Token lifecycle.** Lakebase OAuth tokens live ~60 min. The background task refreshes every 50.
+- **SSL required.** Hardcoded in `connect_args`.
+- **Resource key must match `valueFrom`.** `app.yaml` references `valueFrom: postgres`. If you use a different key in the UI, change either side to match.
+
+## What this example doesn't show
+
+Kept intentionally minimal. Not included:
+- Front-end UI
+- Routes that query application tables (bring your own)
+- Write operations (INSERT/UPDATE/DELETE)
+- User-authorized queries (connects as the app SP, not the end user)
+- Migrations / Alembic
diff --git a/fastapi-postgres-app/app.yaml b/fastapi-postgres-app/app.yaml
@@ -0,0 +1,15 @@
+command:
+  - uvicorn
+  - main:app
+  - --host
+  - 0.0.0.0
+  - --port
+  - "8000"
+
+env:
+  # Pulls the Lakebase Autoscale endpoint resource name from the bound resource
+  # (set the resource's "Resource key" in the Apps UI to `postgres`).
+  # PGHOST, PGPORT, PGDATABASE, PGUSER, PGSSLMODE, PGAPPNAME are all auto-injected
+  # by the platform when the Lakebase resource is bound — no manual config needed.
+  - name: PGENDPOINT
+    valueFrom: postgres
diff --git a/fastapi-postgres-app/config/database.py b/fastapi-postgres-app/config/database.py
@@ -0,0 +1,183 @@
+"""Lakebase Autoscale connection helper for a FastAPI app.
+
+Based on the Databricks Apps Cookbook recipe at
+https://apps-cookbook.dev/docs/fastapi/getting_started/lakebase_connection,
+adapted for Lakebase Autoscale (endpoint-based credential generation via
+`w.postgres.generate_database_credential(endpoint=...)`).
+
+Connection details come from the platform-injected env vars that Databricks
+Apps populates when a Lakebase resource is bound to the app:
+
+  PGHOST        Lakebase endpoint host
+  PGPORT        Postgres port (5432)
+  PGDATABASE    Database name
+  PGUSER        App service principal's client ID (also the PG role name)
+  PGSSLMODE     `require`
+  PGAPPNAME     The app name
+
+For the OAuth token we also need the endpoint *resource name*, which is NOT
+auto-injected — set it in app.yaml via `valueFrom: <resource-key>`:
+
+  PGENDPOINT   projects/<project>/branches/<branch>/endpoints/<endpoint>
+
+See https://docs.databricks.com/aws/en/dev-tools/databricks-apps/lakebase
+for the auto-injected variable list.
+"""
+import asyncio
+import logging
+import os
+import time
+from typing import AsyncGenerator
+
+from databricks.sdk import WorkspaceClient
+from dotenv import load_dotenv
+from sqlalchemy import URL, event
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+
+load_dotenv()
+logger = logging.getLogger(__name__)
+
+# Module-level state. Initialized at FastAPI startup via init_engine().
+engine: AsyncEngine | None = None
+AsyncSessionLocal: sessionmaker | None = None
+workspace_client: WorkspaceClient | None = None
+
+# Current Lakebase OAuth token. Rotated by the background task every 50 min;
+# the SQLAlchemy event listener reads this on every new connection.
+postgres_password: str | None = None
+last_password_refresh: float = 0.0
+token_refresh_task: asyncio.Task | None = None
+
+
+def _endpoint() -> str:
+    val = os.getenv("PGENDPOINT")
+    if not val:
+        raise RuntimeError(
+            "PGENDPOINT is required. Bind a Lakebase resource to the app "
+            "and reference it via `valueFrom: <resource-key>` in app.yaml. "
+            "Format: projects/<project>/branches/<branch>/endpoints/<endpoint>."
+        )
+    return val
+
+
+def _mint_token() -> str:
+    """Generate a fresh OAuth token for the Lakebase endpoint."""
+    assert workspace_client is not None, "WorkspaceClient not initialized"
+    cred = workspace_client.postgres.generate_database_credential(endpoint=_endpoint())
+    if not cred.token:
+        raise RuntimeError("generate_database_credential returned no token")
+    return cred.token
+
+
+async def refresh_token_background():
+    """Refresh the Lakebase OAuth token every 50 minutes.
+
+    Tokens are valid ~60 min; refreshing at 50 gives a 10-min safety margin.
+    The new token is stored in module-global `postgres_password`; the
+    SQLAlchemy event listener reads it for each new connection.
+    """
+    global postgres_password, last_password_refresh
+    while True:
+        try:
+            await asyncio.sleep(50 * 60)
+            logger.info("Refreshing Lakebase OAuth token")
+            postgres_password = _mint_token()
+            last_password_refresh = time.time()
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            logger.error("Background token refresh failed: %s", e)
+
+
+def init_engine() -> None:
+    """Initialize the async SQLAlchemy engine with Lakebase Autoscale auth."""
+    global engine, AsyncSessionLocal, workspace_client, postgres_password, last_password_refresh
+
+    workspace_client = WorkspaceClient()
+
+    # Auto-injected by the Apps platform when a Lakebase resource is bound.
+    host = os.environ.get("PGHOST")
+    database = os.environ.get("PGDATABASE")
+    port = int(os.environ.get("PGPORT", "5432"))
+    user = os.environ.get("PGUSER") or workspace_client.current_user.me().user_name
+    sslmode = os.environ.get("PGSSLMODE", "require")
+
+    if not host or not database:
+        raise RuntimeError(
+            "PGHOST and PGDATABASE are required. They're auto-injected when a "
+            "Lakebase resource is bound to the app (per "
+            "docs.databricks.com/aws/en/dev-tools/databricks-apps/lakebase). "
+            "For local dev, set them in .env (see .env.example)."
+        )
+
+    # Initial token.
+    postgres_password = _mint_token()
+    last_password_refresh = time.time()
+    logger.info("Initial Lakebase token minted; user=%s host=%s db=%s", user, host, database)
+
+    url = URL.create(
+        drivername="postgresql+asyncpg",
+        username=user,
+        password="",  # injected by event listener below
+        host=host,
+        port=port,
+        database=database,
+    )
+
+    # asyncpg uses `ssl` not `sslmode` in connect_args. Map common PG values.
+    ssl_arg = "require" if sslmode.lower() in ("require", "verify-ca", "verify-full") else False
+
+    engine = create_async_engine(
+        url,
+        pool_pre_ping=False,
+        echo=False,
+        pool_size=int(os.getenv("DB_POOL_SIZE", "5")),
+        max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "10")),
+        pool_timeout=int(os.getenv("DB_POOL_TIMEOUT", "10")),
+        pool_recycle=int(os.getenv("DB_POOL_RECYCLE_INTERVAL", "3600")),
+        connect_args={
+            "command_timeout": int(os.getenv("DB_COMMAND_TIMEOUT", "30")),
+            "server_settings": {
+                "application_name": os.environ.get("PGAPPNAME", "lakebase_fastapi_example"),
+            },
+            "ssl": ssl_arg,
+        },
+    )
+
+    @event.listens_for(engine.sync_engine, "do_connect")
+    def provide_token(dialect, conn_rec, cargs, cparams):
+        cparams["password"] = postgres_password
+
+    AsyncSessionLocal = sessionmaker(bind=engine, class_=AsyncSession, expire_on_commit=False)
+    logger.info("Database engine initialized")
+
+
+async def start_token_refresh() -> None:
+    global token_refresh_task
+    if token_refresh_task is None or token_refresh_task.done():
+        token_refresh_task = asyncio.create_task(refresh_token_background())
+        logger.info("Background token refresh task started")
+
+
+async def stop_token_refresh() -> None:
+    global token_refresh_task
+    if token_refresh_task and not token_refresh_task.done():
+        token_refresh_task.cancel()
+        try:
+            await token_refresh_task
+        except asyncio.CancelledError:
+            pass
+        logger.info("Background token refresh task stopped")
+
+
+async def get_async_db() -> AsyncGenerator[AsyncSession, None]:
+    """FastAPI dependency for getting an async DB session per request.
+
+    Inject into endpoints with:
+        db: Annotated[AsyncSession, Depends(get_async_db)]
+    """
+    if AsyncSessionLocal is None:
+        raise RuntimeError("Engine not initialized; call init_engine() first")
+    async with AsyncSessionLocal() as session:
+        yield session
diff --git a/fastapi-postgres-app/main.py b/fastapi-postgres-app/main.py
@@ -0,0 +1,40 @@
+"""FastAPI app entrypoint.
+
+Pattern from https://apps-cookbook.dev/docs/fastapi/getting_started/lakebase_connection,
+adapted for Lakebase Autoscale. See config/database.py for the auth flow.
+"""
+import logging
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from config.database import init_engine, start_token_refresh, stop_token_refresh
+from routes import health
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    init_engine()
+    await start_token_refresh()
+    logger.info("App started with Lakebase connection")
+    yield
+    await stop_token_refresh()
+    logger.info("App shutdown complete")
+
+
+app = FastAPI(
+    title="Lakebase FastAPI Example",
+    description="Minimal FastAPI app connecting to Lakebase Autoscale with async SQLAlchemy.",
+    version="0.1.0",
+    lifespan=lifespan,
+)
+
+app.include_router(health.router)
+
+
+@app.get("/")
+async def root():
+    return {"message": "Lakebase FastAPI example. See /docs for the OpenAPI UI."}