SuperAPI AI Directory
Partner with us中文AI newsFavorites
Skillslakebase-setup

lakebase-setup

Agents & models

View repository

Updated 10-16 05:05

About this skill

Lakebase Setup for Agent Persistence

Profile reminder: All databricks CLI commands must include the profile from .env: databricks <command> --profile <profile> or DATABRICKS_CONFIG_PROFILE=<profile> databricks <command>

Lakebase (autoscaling): This skill uses autoscaling Lakebase — the project/branch/endpoint model. Make sure you have the autoscaling endpoint (a short endpoint name or the full resource path projects/<p>/branches/<b>/endpoints/<e>), or the project + branch.

Use Cases

Lakebase is used for three distinct purposes across the agent templates:

Use case Templates Description
Chat UI conversation history All templates The built-in chat UI (e2e-chatbot-app-next) can persist conversations across page refreshes and browser sessions. This is purely UI-side persistence — the agent itself is stateless.
Agent short-term memory agent-langgraph-advanced, agent-openai-advanced Conversation threads within a session via AsyncCheckpointSaver (LangGraph) or AsyncDatabricksSession (OpenAI SDK). The agent remembers what was said earlier in the same conversation.
Agent long-term memory agent-langgraph-advanced User facts across sessions via AsyncDatabricksStore. The agent remembers things about a user from previous conversations.

Note: When the quickstart prompts for Lakebase on a non-memory template, it's for chat UI history only — not for the agent. Memory templates always require Lakebase.

Overview

Lakebase provides persistent PostgreSQL storage for agents:

  • Short-term memory (LangGraph): Conversation history within a thread (AsyncCheckpointSaver)
  • Long-term memory (LangGraph): User facts across sessions (AsyncDatabricksStore)
  • Short-term memory (OpenAI SDK): Conversation history via AsyncDatabricksSession
  • Long-running agent persistence (OpenAI SDK): Background task state via custom SQLAlchemy tables (agent_server schema)

Note: For pre-configured memory templates, see:

  • agent-langgraph-advanced - Short-term memory, long-term memory, and long-running background tasks (LangGraph)
  • agent-openai-advanced - Short-term memory and long-running background tasks (OpenAI SDK)

Complete Setup Workflow

┌───────────────────────────────────────────────────────────────────────────┐
│  1. Add dependency  →  2. Get instance  →  3. Configure DAB              │
│  4. Configure .env  →  5. Deploy  →  6. Grant SP permissions  →  7. Run  │
└───────────────────────────────────────────────────────────────────────────┘

Shortcut: If using a pre-configured memory template, uv run quickstart with Lakebase flags handles steps 2-4 automatically. You still need to do steps 5-7 manually.

Step 1: Add Memory Dependency

Add the memory extra to your pyproject.toml:

dependencies = [
    "databricks-langchain[memory]",
    # ... other dependencies
]

Then sync dependencies:

uv sync

Step 2: Create or Get Lakebase Instance

Autoscaling uses a project/branch model. You need three values:

  • Project name (e.g., my-project)
  • Branch name (e.g., my-branch)
  • Database ID (e.g., db-xxxx-xxxxxxxxxx)

Find these via the postgres API:

# List projects
databricks api get /api/2.0/postgres/projects --profile <profile>

# List branches for a project
databricks api get /api/2.0/postgres/projects/<project-name>/branches --profile <profile>

# List databases for a branch
databricks api get /api/2.0/postgres/projects/<project-name>/branches/<branch-name>/databases --profile <profile>

Important: The database ID is the internal ID (e.g., db-xxxx-xxxxxxxxxx), NOT databricks_postgres.

Step 3: Configure databricks.yml (Lakebase Resource)

Note: If you ran uv run quickstart with the Lakebase flag (--lakebase-autoscaling-endpoint), the quickstart already configured databricks.yml for you — including fetching the branch/database for autoscaling. Manual configuration is only needed if you didn't use quickstart or need to change values.

Add the postgres resource to your app in databricks.yml:

resources:
  apps:
    your_app:
      name: "your-app-name"
      source_code_path: ./
      resources:
        # ... other resources (experiment, UC functions, etc.) ...

        # Autoscaling Lakebase instance for long-term memory
        - name: 'postgres'
          postgres:
            branch: "projects/<project-name>/branches/<branch-name>"
            database: "projects/<project-name>/branches/<branch-name>/databases/<database-id>"
            permission: 'CAN_CONNECT_AND_CREATE'

Important: The branch and database fields use full resource path format.

See .claude/skills/add-tools/examples/lakebase-autoscaling.yaml for the YAML snippet.

Add Environment Variables to databricks.yml config block

      config:
        env:
          # Autoscaling Lakebase endpoint - resolved from postgres resource at deploy time
          - name: LAKEBASE_AUTOSCALING_ENDPOINT
            value_from: "postgres"
          # Static values for embedding configuration
          - name: EMBEDDING_ENDPOINT
            value: "databricks-gte-large-en"
          - name: EMBEDDING_DIMS
            value: "1024"

Step 4: Configure .env (Local Development)

For local development, add to .env:

LAKEBASE_AUTOSCALING_ENDPOINT=<your-endpoint>
EMBEDDING_ENDPOINT=databricks-gte-large-en
EMBEDDING_DIMS=1024

Important: embedding_dims must match the embedding endpoint:

Endpoint Dimensions
databricks-gte-large-en 1024
databricks-bge-large-en 1024

Note: .env is only for local development. When deployed, the app gets values from databricks.yml config env.

Step 5: Initialize Tables

Step 5: Deploy

Deploy the app so the service principal and resources are created:

DATABRICKS_CONFIG_PROFILE=<profile> databricks bundle deploy

Step 6: Grant SP Permissions (CRITICAL)

WARNING: You MUST complete this step before running the app. Without it, the app will fail with database migration errors like CREATE TABLE IF NOT EXISTS "drizzle"."__drizzle_migrations" — permission denied.

After deploying, the app's service principal needs Postgres roles to access Lakebase tables. The DAB resource grants basic connectivity, but you must also grant Postgres-level schema and table permissions.

Step 1: Get the app's service principal client ID:

DATABRICKS_CONFIG_PROFILE=<profile> databricks apps get <app-name> --output json | jq -r '.service_principal_client_id'

Step 2: Grant permissions using the grant script:

# Autoscaling (endpoint — reads LAKEBASE_AUTOSCALING_ENDPOINT from .env by default):
DATABRICKS_CONFIG_PROFILE=<profile> uv run python scripts/grant_lakebase_permissions.py <sp-client-id> \
  --memory-type <type> --autoscaling-endpoint <endpoint>

# Autoscaling (project + branch):
DATABRICKS_CONFIG_PROFILE=<profile> uv run python scripts/grant_lakebase_permissions.py <sp-client-id> \
  --memory-type <type> --project <project> --branch <branch>

Memory type by template:

Template --memory-type value
agent-langgraph-advanced langgraph
agent-openai-advanced openai

The script handles fresh branches gracefully (warns but doesn't fail if tables don't exist yet — they'll be created on first app startup).

Step 7: Run Your App

DATABRICKS_CONFIG_PROFILE=<profile> databricks bundle run {{BUNDLE_NAME}}

Note: bundle deploy only uploads files and configures resources. bundle run is required to actually start the app with the new code.

Complete Example: databricks.yml with Autoscaling Lakebase

bundle:
  name: agent_langgraph

resources:
  apps:
    agent_langgraph:
      name: "my-agent-app"
      description: "Agent with long-term memory"
      source_code_path: ./
      config:
        command: ["uv", "run", "start-app"]
        env:
          - name: MLFLOW_TRACKING_URI
            value: "databricks"
          - name: MLFLOW_REGISTRY_URI
            value: "databricks-uc"
          - name: API_PROXY
            value: "http://localhost:8000/invocations"
          - name: CHAT_APP_PORT
            value: "3000"
          - name: CHAT_PROXY_TIMEOUT_SECONDS
            value: "300"
          - name: MLFLOW_EXPERIMENT_ID
            value_from: "experiment"
          # Autoscaling Lakebase config
          - name: LAKEBASE_AUTOSCALING_PROJECT
            value: "<your-project-name>"
          - name: LAKEBASE_AUTOSCALING_BRANCH
            value: "<your-branch-name>"
          # Static values for embedding configuration
          - name: EMBEDDING_ENDPOINT
            value: "databricks-gte-large-en"
          - name: EMBEDDING_DIMS
            value: "1024"

      resources:
        - name: 'experiment'
          experiment:
            experiment_id: ""
            permission: 'CAN_MANAGE'
        - name: 'postgres'
          postgres:
            branch: "projects/<your-project-name>/branches/<your-branch-name>"
            database: "projects/<your-project-name>/branches/<your-branch-name>/databases/<your-database-id>"
            permission: 'CAN_CONNECT_AND_CREATE'

targets:
  dev:
    mode: development
    default: true

Troubleshooting

Issue Cause Solution
"embedding_dims is required when embedding_endpoint is specified" Missing embedding_dims parameter Add embedding_dims=1024 to AsyncDatabricksStore
"relation 'store' does not exist" Tables not initialized The app creates tables on first use; ensure SP has CREATE permission
"Unable to resolve Lakebase endpoint 'None'" Missing env var in deployed app Add LAKEBASE_AUTOSCALING_ENDPOINT to databricks.yml config.env
"permission denied for table store" Missing grants Run uv run python scripts/grant_lakebase_permissions.py <sp-client-id> to grant permissions
"Failed to connect to Lakebase" Wrong instance name or project/branch Verify values in databricks.yml and .env
Connection pool errors on exit Python cleanup race Ignore PythonFinalizationError - it's harmless
App not updated after deploy Forgot to run bundle Run databricks bundle run <app> after deploy
value_from not resolving Resource name mismatch Ensure value_from value matches name in databricks.yml resources
"Invalid postgres resource parameters" Missing database field in postgres resource Add full database path: projects/<project>/branches/<branch>/databases/<db-id>
CREATE TABLE IF NOT EXISTS "drizzle"."__drizzle_migrations" fails Grant step was skipped — SP lacks Postgres permissions Run grant_lakebase_permissions.py with --memory-type, then restart the app

LakebaseClient API (for reference)

from databricks_ai_bridge.lakebase import LakebaseClient, SchemaPrivilege, TablePrivilege

# Autoscaling:
client = LakebaseClient(project="...", branch="...")

# Create role (must do first)
client.create_role(identity_name, "SERVICE_PRINCIPAL")

# Grant schema (note: schemas is a list, grantee not role)
client.grant_schema(
    grantee="...",
    schemas=["public"],
    privileges=[SchemaPrivilege.USAGE, SchemaPrivilege.CREATE],
)

# Grant tables (note: tables includes schema prefix)
client.grant_table(
    grantee="...",
    tables=["public.store"],
    privileges=[TablePrivilege.SELECT, TablePrivilege.INSERT, ...],
)

# Execute raw SQL
client.execute("SELECT * FROM pg_tables WHERE schemaname = 'public'")

Service Principal Identifiers

When granting permissions manually, note that Databricks apps have multiple identifiers:

Field Format Example
service_principal_id Numeric ID 1234567890123456
service_principal_client_id UUID a1b2c3d4-e5f6-7890-abcd-ef1234567890
service_principal_name String name my-app-service-principal

Get all identifiers:

DATABRICKS_CONFIG_PROFILE=<profile> databricks apps get <app-name> --output json | jq '{
  id: .service_principal_id,
  client_id: .service_principal_client_id,
  name: .service_principal_name
}'

Which to use:

  • LakebaseClient.create_role() - Use service_principal_client_id (UUID) or service_principal_name
  • Raw SQL grants - Use service_principal_client_id (UUID)

Next Steps

  • Add memory to agent code: see agent-memory skill
  • Test locally: see run-locally skill
  • Deploy: see deploy skill