Lakebase: Postgres Meets the Lakehouse — Hands-On with Databricks' OLTP Play

Table of Contents

Databricks Deep Dive - This article is part of a series.

Part : Agent Bricks: Building Enterprise AI Agents on Databricks

Part : This Article

Part : Lakeflow GA: The Unified Data Engineering Platform You've Been Waiting For

Part : MLflow 3: Rebuilt for the GenAI Era

Part : Mosaic AI Gateway: Unified AI Governance at Scale

Part : Serverless GPU & AI Runtime: Running GPU Workloads Without Cluster Management on Databricks

Part : Unity Catalog Deep Dive: Fine-Grained Access Control

For years, the Databricks lakehouse had a conspicuous gap: operational workloads. You could build world-class analytics pipelines, train ML models, and govern petabytes of data — but the moment you needed a transactional database for your application, you were back to provisioning an external Postgres, managing credentials, and stitching together CDC pipelines to keep things in sync.

Lakebase changes that. It’s a fully managed, Postgres-compatible OLTP database running inside the Databricks platform. Same wire protocol. Same SQL dialect. Same psql you’ve used for a decade. But with lakehouse-native superpowers: autoscaling, scale-to-zero, database branching, point-in-time recovery, and built-in Unity Catalog integration.

In this article, we’ll go beyond the marketing and get hands-on: create a project, connect, build tables, branch your database, register it in Unity Catalog, and wire up pgvector for semantic search — all with real commands you can run today.

Why Lakebase Matters
#

Before Lakebase, the typical architecture for a data + AI application on Databricks looked like this:

1
2
3
4
5
6
7


[Application] --> [External Postgres (RDS/Cloud SQL)]
                        |
                  CDC Pipeline (Debezium, Fivetran, etc.)
                        |
                  [Delta Lake in Lakehouse]
                        |
                  [Analytics / ML / AI]

That’s a lot of moving parts. Every CDC pipeline is a failure point. Every external database is a credential to manage, a network path to secure, and a bill to optimize separately.

With Lakebase, this collapses:

1
2
3
4
5
6
7


[Application] --> [Lakebase (managed Postgres)]
                        |
                  Lakehouse Sync (built-in CDC via wal2delta)
                        |
                  [Delta Lake in Unity Catalog]
                        |
                  [Analytics / ML / AI]

One platform. One governance model. One bill.

Lakebase runs real Postgres 17 — not a Postgres-like API or a compatibility layer. Your existing ORMs, drivers, and psql scripts work unchanged. It also ships with pgvector for vector similarity search out of the box.

Prerequisites
#

Before you start:

A Databricks workspace on AWS (GA) or Azure (public preview)
Workspace admin privileges (to create Lakebase projects)
psql installed locally (or any Postgres client)
Basic familiarity with SQL and the Databricks UI
A Unity Catalog metastore attached to your workspace (for the UC integration sections)

Autoscaling vs. Provisioned: Which One?
#

Since March 2026, all new Lakebase instances default to Autoscaling. Here’s the quick comparison:

Feature	Provisioned	Autoscaling
Scaling	Manual	Automatic
Scale-to-zero	No	Yes (default 5 min idle)
RAM per CU	~16 GB	~2 GB (more granular)
Branching	No	Yes
Instant restore	No	Yes
New instances	Legacy	Default since March 2026

If you’re starting fresh, use Autoscaling. Provisioned instances continue to work but won’t receive the newer capabilities.

Step 1: Create a Lakebase Project
#

In the Databricks sidebar, navigate to Lakebase and click New project.

Give your project a name (e.g., ai-app-backend)
Select Postgres 17
Click Create

Your project spins up with:

A production branch (the main database)
A default databricks_postgres database
Autoscaling compute attached to the branch

Scale-to-zero is enabled by default. If nobody connects for 5 minutes, compute suspends entirely — you pay nothing. The next connection wakes it up in seconds.

Step 2: Connect with psql
#

Grab your connection details from the project’s Connect tab. You’ll get a connection string like:

1

psql 'postgresql://your_role:[email protected]/databricks_postgres?sslmode=require'

Or set environment variables:

1
2
3
4
5
6
7


export PGHOST="ep-your-project-id.databricks.com"
export PGPORT="5432"
export PGDATABASE="databricks_postgres"
export PGUSER="your_role"
export PGPASSWORD="your_password"

psql

You should see the familiar Postgres prompt. Let’s verify:

1

SELECT version();

1
2
3
4


                                                  version
------------------------------------------------------------------------------------------------------------
 PostgreSQL 17.4 on x86_64-pc-linux-gnu, compiled by gcc (GCC) 13.2.0, 64-bit
(1 row)

This is real Postgres. Your existing tools — pgAdmin, DBeaver, SQLAlchemy, Prisma, Django ORM — all connect the same way.

Connect from Python
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


import psycopg2

conn = psycopg2.connect(
    host="ep-your-project-id.databricks.com",
    port=5432,
    dbname="databricks_postgres",
    user="your_role",
    password="your_password",
    sslmode="require"
)

cur = conn.cursor()
cur.execute("SELECT version();")
print(cur.fetchone()[0])
# PostgreSQL 17.4 on x86_64-pc-linux-gnu ...

Or with SQLAlchemy (works with Flask, FastAPI, Django):

1
2
3
4
5
6
7


from sqlalchemy import create_engine

engine = create_engine(
    "postgresql+psycopg2://your_role:your_password"
    "@ep-your-project-id.databricks.com:5432/databricks_postgres",
    connect_args={"sslmode": "require"}
)

No special drivers. No Databricks SDK. Standard Postgres.

Step 3: Build Your Schema
#

Let’s create a practical schema for an AI chatbot application:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


-- Conversation tracking for an AI agent
CREATE TABLE conversations (
    id              SERIAL PRIMARY KEY,
    user_id         VARCHAR(64) NOT NULL,
    agent_id        VARCHAR(64) NOT NULL,
    started_at      TIMESTAMPTZ DEFAULT now(),
    status          VARCHAR(20) DEFAULT 'active'
);

CREATE TABLE messages (
    id              SERIAL PRIMARY KEY,
    conversation_id INTEGER REFERENCES conversations(id),
    role            VARCHAR(20) NOT NULL,  -- 'user', 'assistant', 'system'
    content         TEXT NOT NULL,
    tokens_used     INTEGER,
    created_at      TIMESTAMPTZ DEFAULT now()
);

-- Insert some sample data
INSERT INTO conversations (user_id, agent_id) VALUES
    ('user-001', 'support-bot-v2'),
    ('user-002', 'support-bot-v2'),
    ('user-003', 'onboarding-agent');

INSERT INTO messages (conversation_id, role, content, tokens_used) VALUES
    (1, 'user', 'How do I reset my password?', 12),
    (1, 'assistant', 'You can reset your password by clicking...', 45),
    (2, 'user', 'What are your pricing plans?', 10),
    (2, 'assistant', 'We offer three tiers: Starter, Pro, and Enterprise...', 68),
    (3, 'user', 'Help me get started with the API', 14),
    (3, 'assistant', 'Welcome! Let me walk you through the setup...', 52);

Verify the data:

1
2
3
4


SELECT c.id, c.user_id, c.agent_id, count(m.id) as message_count
FROM conversations c
JOIN messages m ON m.conversation_id = c.id
GROUP BY c.id, c.user_id, c.agent_id;

1
2
3
4
5
6


 id | user_id  |    agent_id      | message_count
----+----------+------------------+---------------
  1 | user-001 | support-bot-v2   |             2
  2 | user-002 | support-bot-v2   |             2
  3 | user-003 | onboarding-agent |             2
(3 rows)

Quick check: Postgres extensions available
#

1
2
3


SELECT name, default_version FROM pg_available_extensions
WHERE name IN ('vector', 'pg_trgm', 'uuid-ossp', 'hstore', 'pgcrypto')
ORDER BY name;

1
2
3
4
5
6
7
8


   name    | default_version
-----------+-----------------
 hstore    | 1.8
 pg_trgm   | 1.6
 pgcrypto  | 1.3
 uuid-ossp | 1.1
 vector    | 0.8.0
(5 rows)

Good news — the essentials are there, including pgvector. We’ll use it in Step 7.

Step 4: Database Branching — Your New Best Friend
#

This is where Lakebase goes beyond standard managed Postgres. Branching creates an instant, isolated copy of your database using copy-on-write — no full data duplication.

Create a development branch
#

In the Databricks UI, navigate to your project and click Create branch from the production branch. Name it dev-new-schema.

Your dev-new-schema branch is an exact copy of production — same data, same schema — but completely isolated. Any changes you make here don’t touch production.

Develop safely on the branch
#

Connect to the branch (it gets its own connection string) and make your schema changes:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


-- Add embedding support to messages (on the dev branch)
ALTER TABLE messages ADD COLUMN embedding vector(1536);

-- Add a feedback table
CREATE TABLE feedback (
    id              SERIAL PRIMARY KEY,
    message_id      INTEGER REFERENCES messages(id),
    rating          SMALLINT CHECK (rating BETWEEN 1 AND 5),
    comment         TEXT,
    created_at      TIMESTAMPTZ DEFAULT now()
);

Verify the schema change on the branch:

1

\d messages

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


                                        Table "public.messages"
     Column      |           Type           | Collation | Nullable |               Default
-----------------+--------------------------+-----------+----------+--------------------------------------
 id              | integer                  |           | not null | nextval('messages_id_seq'::regclass)
 conversation_id | integer                  |           |          |
 role            | character varying(20)    |           | not null |
 content         | text                     |           | not null |
 tokens_used     | integer                  |           |          |
 created_at      | timestamp with time zone |           |          | now()
 embedding       | vector(1536)             |           |          |

The embedding column is live on the branch, but production is untouched. Test thoroughly, run your integration tests against the branch. When you’re confident, use schema diff in the UI to review changes before applying them to production.

Point-in-time branching for recovery
#

Accidentally dropped a table on production at 10:23 AM? Create a branch from 10:22 AM:

Click Create branch on the production branch
Select Point in time
Set the timestamp to just before the incident
Connect to the recovery branch and extract the missing data

No need to restore the entire database. No downtime. This alone is worth the switch for many teams.

Copy-on-write means branches are nearly instant regardless of database size. A 100 GB production database branches in seconds, not hours. You only pay for the incremental storage of changes made on the branch.

Step 5: Register in Unity Catalog
#

Your Lakebase database is invisible to the rest of Databricks until you register it in Unity Catalog. Once registered, it becomes a read-only catalog that the entire platform can query.

Register via the UI
#

In your Lakebase project, click Register in Unity Catalog
Choose a catalog name (e.g., chatbot_app)
Select the target schema mapping

Query from Databricks SQL
#

Once registered, anyone with the right permissions can query your operational data from SQL Warehouses, notebooks, or dashboards:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


-- Query live operational data from a SQL Warehouse
SELECT
    agent_id,
    COUNT(*) as total_conversations,
    AVG(message_count) as avg_messages_per_conversation
FROM (
    SELECT
        c.agent_id,
        c.id,
        COUNT(m.id) as message_count
    FROM chatbot_app.public.conversations c
    JOIN chatbot_app.public.messages m ON m.conversation_id = c.id
    GROUP BY c.agent_id, c.id
) sub
GROUP BY agent_id;

Join operational and analytical data
#

This is the real power — joining OLTP data with your lakehouse in a single query:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


-- Combine live conversation data with historical user analytics
SELECT
    c.user_id,
    c.agent_id,
    m.content as last_message,
    u.subscription_tier,
    u.lifetime_value
FROM chatbot_app.public.conversations c
JOIN chatbot_app.public.messages m
    ON m.conversation_id = c.id
JOIN main.user_analytics.users u
    ON c.user_id = u.user_id
WHERE c.started_at >= current_date - INTERVAL 1 DAY
ORDER BY c.started_at DESC;

1
2
3
4
5


 user_id  |    agent_id      |               last_message                | subscription_tier | lifetime_value
----------+------------------+-------------------------------------------+-------------------+----------------
 user-001 | support-bot-v2   | You can reset your password by clicking...| pro               |        1250.00
 user-002 | support-bot-v2   | We offer three tiers: Starter, Pro, ...   | enterprise        |       45000.00
 user-003 | onboarding-agent | Welcome! Let me walk you through the...   | starter           |           0.00

No ETL. No CDC pipeline. Just SQL across operational and analytical data, governed by Unity Catalog.

Step 6: Lakehouse Sync — CDC Without the Pipeline
#

While the Unity Catalog registration gives you read access to live Lakebase data, Lakehouse Sync goes further: it continuously replicates your Postgres tables into Delta tables using built-in Change Data Capture.

How it works
#

Lakehouse Sync is powered by wal2delta, a Postgres extension that runs inside the Lakebase compute. It uses logical decoding to capture WAL (Write-Ahead Log) changes and writes them directly to Delta tables in Unity Catalog.

1
2
3
4
5


Lakebase Postgres
    └── WAL (Write-Ahead Log)
        └── wal2delta (logical decoding)
            └── Delta tables in Unity Catalog
                └── SCD Type 2 history

Key facts:

No external compute required — it runs inside Lakebase
No pipelines or jobs to manage — it’s a native feature
SCD Type 2 history — every change is appended, giving you full audit trail
Delta tables are named lb_<table_name>_history

Enable Lakehouse Sync
#

Enable it from the Lakebase project settings. Choose which tables to sync and the target catalog/schema. The sync starts immediately and runs continuously.

Once synced, you can query the history:

1
2
3
4
5


-- See the full change history for conversations
SELECT *
FROM main.lakebase_sync.lb_conversations_history
ORDER BY _lb_commit_timestamp DESC
LIMIT 20;

1
2
3
4
5


 id | user_id  | agent_id         | started_at                    | status   | _lb_op | _lb_commit_timestamp
----+----------+------------------+-------------------------------+----------+--------+------------------------------
  1 | user-001 | support-bot-v2   | 2026-04-18 10:15:22.123+00    | closed   | UPDATE | 2026-04-18 10:45:03.456+00
  1 | user-001 | support-bot-v2   | 2026-04-18 10:15:22.123+00    | active   | INSERT | 2026-04-18 10:15:22.789+00
  2 | user-002 | support-bot-v2   | 2026-04-18 10:20:11.456+00    | active   | INSERT | 2026-04-18 10:20:12.012+00

Every row change is captured with the operation type (_lb_op) and commit timestamp. This gives you a complete audit trail of every INSERT, UPDATE, and DELETE — without writing a single line of pipeline code.

Analytics on top of sync history
#

1
2
3
4
5
6
7
8
9


-- Agent performance over time (from Delta, not hitting Postgres)
SELECT
    date_trunc('hour', _lb_commit_timestamp) as hour,
    count(CASE WHEN _lb_op = 'INSERT' THEN 1 END) as new_conversations,
    count(CASE WHEN _lb_op = 'UPDATE' AND status = 'closed' THEN 1 END) as resolved
FROM main.lakebase_sync.lb_conversations_history
WHERE _lb_commit_timestamp >= current_date
GROUP BY 1
ORDER BY 1;

This runs on Delta Lake — zero load on your production Postgres. Your analysts can query freely without impacting application performance.

Step 7: pgvector for AI Applications
#

Lakebase ships with pgvector pre-installed, making it a natural fit for AI applications that need vector similarity search alongside traditional relational data.

Enable the extension
#

1

CREATE EXTENSION IF NOT EXISTS vector;

Create a knowledge base with embeddings
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


CREATE TABLE knowledge_base (
    id          SERIAL PRIMARY KEY,
    title       VARCHAR(255) NOT NULL,
    content     TEXT NOT NULL,
    embedding   vector(1536),  -- OpenAI ada-002 dimensions
    category    VARCHAR(64),
    created_at  TIMESTAMPTZ DEFAULT now()
);

-- Create an index for fast similarity search
CREATE INDEX ON knowledge_base
    USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100);

Semantic search query
#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


-- Find the 5 most relevant knowledge base articles
-- (assuming you've generated an embedding for the query)
SELECT
    title,
    content,
    1 - (embedding <=> $1::vector) as similarity
FROM knowledge_base
WHERE category = 'product-docs'
ORDER BY embedding <=> $1::vector
LIMIT 5;

Complete Python RAG example
#

Here’s a complete, runnable example — an AI agent that stores conversations and retrieves relevant context from a knowledge base using pgvector:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87


import psycopg2
from psycopg2.extras import RealDictCursor
import openai
import json

# --- Config ---
DB_CONN = "postgresql://your_role:[email protected]/databricks_postgres?sslmode=require"
oai = openai.OpenAI()

def get_embedding(text: str) -> list[float]:
    """Generate embedding using OpenAI ada-002."""
    resp = oai.embeddings.create(input=text, model="text-embedding-ada-002")
    return resp.data[0].embedding

def seed_knowledge_base(conn):
    """Insert sample knowledge base articles with embeddings."""
    articles = [
        ("Password Reset", "To reset your password, go to Settings > Security > Reset Password. You'll receive an email with a reset link."),
        ("Pricing Plans", "We offer Starter ($9/mo), Pro ($49/mo), and Enterprise (custom). Pro includes API access and priority support."),
        ("API Rate Limits", "Free tier: 100 req/min. Pro: 1000 req/min. Enterprise: custom. Rate limit headers are X-RateLimit-Remaining."),
        ("Data Export", "Export your data from Settings > Data > Export. Formats: CSV, JSON, Parquet. Exports over 1GB are async."),
    ]
    with conn.cursor() as cur:
        for title, content in articles:
            embedding = get_embedding(content)
            cur.execute(
                """INSERT INTO knowledge_base (title, content, embedding, category)
                   VALUES (%s, %s, %s::vector, 'product-docs')
                   ON CONFLICT DO NOTHING""",
                (title, content, str(embedding))
            )
    conn.commit()

def retrieve_context(conn, query: str, top_k: int = 3) -> list[dict]:
    """Semantic search — find most relevant knowledge base articles."""
    query_embedding = get_embedding(query)
    with conn.cursor(cursor_factory=RealDictCursor) as cur:
        cur.execute("""
            SELECT title, content,
                   1 - (embedding <=> %s::vector) as similarity
            FROM knowledge_base
            WHERE category = 'product-docs'
            ORDER BY embedding <=> %s::vector
            LIMIT %s
        """, (str(query_embedding), str(query_embedding), top_k))
        return cur.fetchall()

def save_message(conn, conversation_id: int, role: str, content: str):
    """Persist message to conversation history."""
    with conn.cursor() as cur:
        cur.execute(
            """INSERT INTO messages (conversation_id, role, content, tokens_used)
               VALUES (%s, %s, %s, %s)""",
            (conversation_id, role, content, len(content.split()))
        )
    conn.commit()

# --- Main flow ---
conn = psycopg2.connect(DB_CONN)
seed_knowledge_base(conn)

# User asks a question
user_question = "How do I export my data as Parquet?"

# 1. Retrieve relevant context from pgvector
context_docs = retrieve_context(conn, user_question)
print("Retrieved context:")
for doc in context_docs:
    print(f"  [{doc['similarity']:.3f}] {doc['title']}")

# 2. Build prompt with retrieved context
context_text = "\n\n".join(f"### {d['title']}\n{d['content']}" for d in context_docs)
response = oai.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": f"Answer using this context:\n\n{context_text}"},
        {"role": "user", "content": user_question}
    ]
)
answer = response.choices[0].message.content

# 3. Persist the conversation (ACID-guaranteed)
save_message(conn, 1, "user", user_question)
save_message(conn, 1, "assistant", answer)

print(f"\nAgent: {answer}")
conn.close()

1
2
3
4
5
6
7
8


Retrieved context:
  [0.932] Data Export
  [0.841] API Rate Limits
  [0.823] Pricing Plans

Agent: You can export your data as Parquet by going to Settings > Data > Export
and selecting "Parquet" as the format. Note that exports over 1GB will be
processed asynchronously — you'll receive a notification when it's ready.

This is a full RAG pipeline — retrieval, generation, and persistence — running against a single Lakebase instance. No separate vector database. No external state store.

LangGraph checkpointer
#

If you’re using LangGraph for agent orchestration, Lakebase works directly as a Postgres checkpointer:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


from langgraph.checkpoint.postgres import PostgresSaver

# Lakebase IS Postgres — LangGraph's native checkpointer just works
checkpointer = PostgresSaver.from_conn_string(DB_CONN)
checkpointer.setup()  # Creates checkpoint tables automatically

# Build your agent graph with persistent state
from langgraph.graph import StateGraph
graph = StateGraph(...)
# ... define nodes and edges ...
app = graph.compile(checkpointer=checkpointer)

# Each thread_id gets its own conversation state, persisted in Lakebase
config = {"configurable": {"thread_id": "user-001-session-42"}}
result = app.invoke({"messages": [("user", "Check order status")]}, config)

# Resume later — state is loaded from Lakebase automatically
result = app.invoke({"messages": [("user", "Now cancel it")]}, config)

Every LangGraph interaction is stored at the thread level in Lakebase. Your agent can resume investigations across sessions, and the entire conversation state is available in Delta Lake via Lakehouse Sync for analytics.

When to Use Lakebase vs. External Postgres
#

Lakebase isn’t a replacement for every Postgres deployment. Here’s a practical decision framework:

Use Lakebase when…	Use external Postgres when…
Building apps on Databricks (Databricks Apps)	You need Postgres extensions beyond what Lakebase supports
Your app data needs to join with lakehouse analytics	Your app is completely independent of the Databricks platform
You want branching/PITR without managing infrastructure	You need multi-region active-active replication
AI agents need relational + vector storage	You’re already deeply invested in RDS/Cloud SQL tooling
You want unified governance via Unity Catalog	Your compliance requires a specific managed Postgres provider

Architecture Pattern: Full-Stack AI App on Databricks
#

Here’s how the pieces fit together for a production AI application:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29


                    ┌─────────────────────────┐
                    │   Databricks App (UI)    │
                    └────────────┬────────────┘
                                 │
                    ┌────────────▼────────────┐
                    │   AI Agent (LangGraph)   │
                    │   - Tool calling         │
                    │   - State management     │
                    └────────────┬────────────┘
                                 │
              ┌──────────────────┼──────────────────┐
              │                  │                   │
   ┌──────────▼─────────┐ ┌─────▼──────┐ ┌─────────▼────────┐
   │  Lakebase           │ │  Model     │ │  Unity Catalog    │
   │  - Conversations    │ │  Serving   │ │  - Delta tables   │
   │  - User state       │ │  (LLM)    │ │  - Feature store  │
   │  - Vector search    │ │            │ │  - ML models      │
   │  (pgvector)         │ │            │ │                   │
   └──────────┬──────────┘ └────────────┘ └───────────────────┘
              │
              │ Lakehouse Sync (wal2delta)
              │
   ┌──────────▼──────────┐
   │  Delta Lake          │
   │  - Conversation      │
   │    history (SCD2)    │
   │  - Analytics         │
   │  - ML training data  │
   └──────────────────────┘

This is the architecture Databricks is betting on. Lakebase is the missing piece that makes “full-stack lakehouse” real.

Minimal FastAPI app backed by Lakebase
#

Here’s a working Databricks App skeleton — a conversation API backed entirely by Lakebase:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87


# app.py — deploy as a Databricks App
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import psycopg2
from psycopg2.extras import RealDictCursor
import os

app = FastAPI(title="AI Chatbot API")

def get_db():
    return psycopg2.connect(
        host=os.environ["LAKEBASE_HOST"],
        port=5432,
        dbname="databricks_postgres",
        user=os.environ["LAKEBASE_USER"],
        password=os.environ["LAKEBASE_PASSWORD"],
        sslmode="require"
    )

class MessageRequest(BaseModel):
    conversation_id: int
    role: str
    content: str

@app.post("/conversations")
def create_conversation(user_id: str, agent_id: str = "default-agent"):
    conn = get_db()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                "INSERT INTO conversations (user_id, agent_id) VALUES (%s, %s) RETURNING *",
                (user_id, agent_id)
            )
            result = cur.fetchone()
        conn.commit()
        return result
    finally:
        conn.close()

@app.get("/conversations/{conv_id}/messages")
def get_messages(conv_id: int):
    conn = get_db()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                "SELECT role, content, tokens_used, created_at FROM messages "
                "WHERE conversation_id = %s ORDER BY created_at",
                (conv_id,)
            )
            return cur.fetchall()
    finally:
        conn.close()

@app.post("/messages")
def post_message(msg: MessageRequest):
    conn = get_db()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                "INSERT INTO messages (conversation_id, role, content, tokens_used) "
                "VALUES (%s, %s, %s, %s) RETURNING *",
                (msg.conversation_id, msg.role, msg.content, len(msg.content.split()))
            )
            result = cur.fetchone()
        conn.commit()
        return result
    finally:
        conn.close()

@app.get("/stats")
def agent_stats():
    """Real-time agent performance — straight from Lakebase."""
    conn = get_db()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute("""
                SELECT c.agent_id,
                       count(DISTINCT c.id) as conversations,
                       count(m.id) as total_messages,
                       round(avg(m.tokens_used), 1) as avg_tokens
                FROM conversations c
                JOIN messages m ON m.conversation_id = c.id
                GROUP BY c.agent_id
            """)
            return cur.fetchall()
    finally:
        conn.close()

Test it locally:

1
2
3
4
5


curl -X POST "http://localhost:8000/conversations?user_id=demo-user"
# {"id": 4, "user_id": "demo-user", "agent_id": "default-agent", "started_at": "2026-04-18T...", "status": "active"}

curl http://localhost:8000/stats
# [{"agent_id": "support-bot-v2", "conversations": 2, "total_messages": 4, "avg_tokens": 33.8}, ...]

Deploy to Databricks Apps and add Lakebase as a resource — credentials rotate automatically.

Current Limitations
#

No product overview is complete without the caveats:

AWS GA, Azure public preview — GCP support coming later in 2026
8 TB maximum per instance (for now)
Read-only Unity Catalog registration — you can’t write to Lakebase via SQL Warehouses
Extension support is curated — not every Postgres extension is available (pgvector is, PostGIS is not yet)
Autoscaling cold starts — scale-to-zero wake-up takes a few seconds, which may matter for latency-sensitive applications
No cross-region replication — single-region only at this stage

These are expected to improve rapidly as Lakebase matures from public preview to full GA across all clouds.

Conclusion
#

Lakebase is a paradigm shift for the Databricks platform. For the first time, you can run operational workloads — your application’s transactional database — alongside analytical and AI workloads on a single platform, governed by a single catalog.

The practical implications are significant:

Fewer moving parts — eliminate external databases and CDC pipelines
Faster development — database branching and scale-to-zero change the development workflow
Unified governance — one set of access controls for operational and analytical data
AI-native — pgvector + relational storage in one place for agent state and RAG

If you’re building AI applications on Databricks, Lakebase should be your default starting point for the transactional layer. The days of “Databricks for analytics, something else for OLTP” are numbered.

Lakebase is GA on AWS with Autoscaling (since March 2026) and in public preview on Azure. Check the Databricks documentation for the latest availability and feature updates.

Author

Mariusz Derela

Cyber Security Specialist | DevSecOps | AI/ML

Databricks Deep Dive - This article is part of a series.

Part : Agent Bricks: Building Enterprise AI Agents on Databricks

Part : This Article

Part : Lakeflow GA: The Unified Data Engineering Platform You've Been Waiting For

Part : MLflow 3: Rebuilt for the GenAI Era

Part : Mosaic AI Gateway: Unified AI Governance at Scale

Part : Serverless GPU & AI Runtime: Running GPU Workloads Without Cluster Management on Databricks

Part : Unity Catalog Deep Dive: Fine-Grained Access Control

Why Lakebase Matters #

Prerequisites #

Autoscaling vs. Provisioned: Which One? #

Step 1: Create a Lakebase Project #

Step 2: Connect with psql #

Connect from Python #

Step 3: Build Your Schema #

Quick check: Postgres extensions available #

Step 4: Database Branching — Your New Best Friend #

Create a development branch #

Develop safely on the branch #

Point-in-time branching for recovery #

Step 5: Register in Unity Catalog #

Register via the UI #

Query from Databricks SQL #

Join operational and analytical data #

Step 6: Lakehouse Sync — CDC Without the Pipeline #

How it works #

Enable Lakehouse Sync #

Analytics on top of sync history #

Step 7: pgvector for AI Applications #

Enable the extension #

Create a knowledge base with embeddings #

Semantic search query #

Complete Python RAG example #

LangGraph checkpointer #

When to Use Lakebase vs. External Postgres #

Architecture Pattern: Full-Stack AI App on Databricks #

Minimal FastAPI app backed by Lakebase #

Current Limitations #

Conclusion #

Related