From 1d6d8638601f36c44d566695d61ce6883407d2b2 Mon Sep 17 00:00:00 2001 From: Matt Adams Date: Mon, 8 Jun 2026 14:33:36 -0700 Subject: [PATCH] Add fastapi-postgres-app template MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Minimal FastAPI service 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 apps-cookbook.dev FastAPI + Lakebase recipe, adapted for Autoscaling — uses w.postgres.generate_database_credential(endpoint=...) instead of the Provisioned variant. 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 uvicorn + PGENDPOINT valueFrom ├── requirements.txt └── .env.example Local dev only Customer flow on Databricks Apps: 1. Create the app 2. Settings → Resources → Add resource: Lakebase Autoscaling project, Permission 'Can connect and create', Resource key 'postgres' 3. Deploy The resource binding auto-injects PG* env vars and app.yaml's `valueFrom: postgres` populates PGENDPOINT. No manual env config needed for deploy. Naming and conventions match the existing peer templates (flask-postgres-app, dash-postgres-app, streamlit-postgres-app): `postgres_spec` in manifest, resource key `postgres`, env var `PGENDPOINT`, requirements.txt for deps. Fills a gap: the repo had flask/dash/streamlit-postgres-app and appkit-lakebase (Node.js) but no FastAPI equivalent for Lakebase Autoscaling. README also includes a Postman service-principal recipe for exercising the deployed API. --- fastapi-postgres-app/.env.example | 46 ++++++ fastapi-postgres-app/README.md | 104 ++++++++++++++ fastapi-postgres-app/app.yaml | 15 ++ fastapi-postgres-app/config/database.py | 183 ++++++++++++++++++++++++ fastapi-postgres-app/main.py | 40 ++++++ fastapi-postgres-app/manifest.yaml | 9 ++ fastapi-postgres-app/requirements.txt | 6 + fastapi-postgres-app/routes/__init__.py | 0 fastapi-postgres-app/routes/health.py | 59 ++++++++ 9 files changed, 462 insertions(+) create mode 100644 fastapi-postgres-app/.env.example create mode 100644 fastapi-postgres-app/README.md create mode 100644 fastapi-postgres-app/app.yaml create mode 100644 fastapi-postgres-app/config/database.py create mode 100644 fastapi-postgres-app/main.py create mode 100644 fastapi-postgres-app/manifest.yaml create mode 100644 fastapi-postgres-app/requirements.txt create mode 100644 fastapi-postgres-app/routes/__init__.py create mode 100644 fastapi-postgres-app/routes/health.py diff --git a/fastapi-postgres-app/.env.example b/fastapi-postgres-app/.env.example new file mode 100644 index 00000000..5adf5369 --- /dev/null +++ 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//branches//endpoints/ +# Find via: +# databricks api get /api/2.0/postgres/projects//branches//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 new file mode 100644 index 00000000..db7a8a12 --- /dev/null +++ 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:///oidc/v1/token` | +| Client ID | `` | +| 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://-.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 new file mode 100644 index 00000000..52f68ba0 --- /dev/null +++ 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 new file mode 100644 index 00000000..818fcd63 --- /dev/null +++ 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: `: + + PGENDPOINT projects//branches//endpoints/ + +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: ` in app.yaml. " + "Format: projects//branches//endpoints/." + ) + 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 new file mode 100644 index 00000000..c3c6f825 --- /dev/null +++ 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."} diff --git a/fastapi-postgres-app/manifest.yaml b/fastapi-postgres-app/manifest.yaml new file mode 100644 index 00000000..9dbcc8d1 --- /dev/null +++ b/fastapi-postgres-app/manifest.yaml @@ -0,0 +1,9 @@ +version: 1 +name: "Lakebase Autoscaling FastAPI app" +description: "A FastAPI service that reads from a Lakebase Autoscaling (Postgres) database. Async SQLAlchemy + asyncpg, OAuth token refresh every 50 minutes, no front-end. Pairs with the cookbook recipe at apps-cookbook.dev/docs/fastapi/getting_started/lakebase_connection (adapted for Autoscale)." + +resource_specs: + - name: "postgres" + description: "The Lakebase Autoscaling endpoint the app reads from." + postgres_spec: + permission: "CAN_CONNECT_AND_CREATE" diff --git a/fastapi-postgres-app/requirements.txt b/fastapi-postgres-app/requirements.txt new file mode 100644 index 00000000..64225159 --- /dev/null +++ b/fastapi-postgres-app/requirements.txt @@ -0,0 +1,6 @@ +databricks-sdk>=0.81.0 +sqlalchemy>=2.0.0 +asyncpg>=0.29.0 +python-dotenv>=1.0.0 +fastapi>=0.115.0 +uvicorn>=0.32.0 diff --git a/fastapi-postgres-app/routes/__init__.py b/fastapi-postgres-app/routes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/fastapi-postgres-app/routes/health.py b/fastapi-postgres-app/routes/health.py new file mode 100644 index 00000000..1b9e3f27 --- /dev/null +++ b/fastapi-postgres-app/routes/health.py @@ -0,0 +1,59 @@ +"""Liveness and readiness endpoints. + +`/healthz` is a cheap liveness probe that does NOT touch Postgres. +`/api/health` is a readiness probe that round-trips to Lakebase. +""" +from typing import Annotated + +from fastapi import APIRouter, Depends +from pydantic import BaseModel, Field +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +from config.database import get_async_db + +router = APIRouter(tags=["health"]) + + +class LivenessResponse(BaseModel): + status: str = Field(..., examples=["alive"]) + + +class ReadinessResponse(BaseModel): + status: str = Field(..., examples=["ok", "degraded"]) + database: bool = Field(..., description="True if a SELECT 1 against Lakebase succeeded.") + + +@router.get( + "/healthz", + response_model=LivenessResponse, + summary="Liveness probe (no DB)", + description=( + "Cheap liveness check. **Does NOT touch Postgres.** Returns immediately. " + "Use this for orchestrator liveness probes where you only want to know " + "whether the process is up." + ), +) +async def healthz() -> LivenessResponse: + return LivenessResponse(status="alive") + + +@router.get( + "/api/health", + response_model=ReadinessResponse, + summary="Readiness probe (round-trips Postgres)", + description=( + "Readiness check that **does** round-trip to Lakebase via `SELECT 1`. " + "Returns `database: true` on success, `database: false` (and " + "`status: degraded`) on any DB error. Use this to confirm the app can " + "actually reach the database, not just that the process is alive." + ), +) +async def api_health( + db: Annotated[AsyncSession, Depends(get_async_db)], +) -> ReadinessResponse: + try: + await db.execute(text("SELECT 1")) + return ReadinessResponse(status="ok", database=True) + except Exception: + return ReadinessResponse(status="degraded", database=False)