lakebase-setup

name	lakebase-setup
description	Configure Lakebase for agent memory storage. Use when: (1) Adding memory capabilities to the agent, (2) 'Failed to connect to Lakebase' errors, (3) Permission errors on checkpoint/store tables, (4) User says 'lakebase', 'memory setup', or 'add memory'.

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
             
            
             
          
            
             
            
             

      
          
          
             
             
          
          
             
             
             


  
     
     

---

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