LangGraph Unity Catalog Checkpoint

Production-ready Unity Catalog persistence for LangChain and LangGraph applications using Databricks as the storage backend.

Following the LangGraph checkpoint-postgres pattern for consistency with the LangGraph ecosystem.

---

🚀 Overview

This package provides enterprise-grade implementations of LangGraph's persistence interfaces backed by Databricks Unity Catalog:

• UnityCatalogStore / AsyncUnityCatalogStore: Implements langgraph.store.base.BaseStore for key-value storage
• UnityCatalogCheckpointSaver / AsyncUnityCatalogCheckpointSaver: Implements BaseCheckpointSaver for graph state persistence

All implementations use Databricks Unity Catalog Delta tables via the WorkspaceClient SQL API, providing:

• ✅ Enterprise-grade reliability with ACID transactions
• ✅ Scalability with Delta Lake optimization
• ✅ Governance with built-in access control and audit trails
• ✅ Time-travel for debugging and recovery
• ✅ Seamless Databricks integration for production ML workflows
• ✅ Performance optimized with batch operations (2-10x faster)

---

📦 Installation

Prerequisites

• Python 3.10+
• Databricks workspace with Unity Catalog enabled
• SQL warehouse with appropriate permissions

Install from PyPI

The easiest way to install is from PyPI:

pip install langgraph-unity-catalog-checkpoint

This will automatically install all required dependencies including:

• databricks-sdk
• langchain
• langgraph
• loguru

Install from Source

For development or to get the latest unreleased features:

# Clone the repository
git clone https://github.com/natefleming/langgraph_unity_catalog_checkpoint.git
cd langgraph_unity_catalog_checkpoint

# Install in editable mode
pip install -e .

# Or with development dependencies
pip install -e ".[dev]"

---

⚡ Quick Start

1. Configure Databricks Authentication

Set up environment variables:

export DATABRICKS_HOST="https://your-workspace.databricks.com"
export DATABRICKS_TOKEN="your-access-token"
export DATABRICKS_WAREHOUSE_ID="your-warehouse-id"
export UC_CATALOG="your_catalog"
export UC_SCHEMA="your_schema"

Or use ~/.databrickscfg:

[DEFAULT]
host = https://your-workspace.databricks.com
token = your-access-token

2. Using the Store for Key-Value Storage

from databricks.sdk import WorkspaceClient
from langgraph_unity_catalog_checkpoint import UnityCatalogStore

# Initialize the store
workspace_client = WorkspaceClient()
store = UnityCatalogStore(
    workspace_client=workspace_client,
    catalog="main",
    schema="langgraph",
    table="my_store",  # Default: "store"
    warehouse_id="your-warehouse-id",  # Optional
)

# Store values with namespaced keys
store.put(("users", "123"), "preferences", {"theme": "dark", "language": "en"})

# Retrieve values
prefs = store.get(("users", "123"), "preferences")
print(prefs)  # {"theme": "dark", "language": "en"}

# Search within a namespace
items = store.search(("users",), limit=10)
for item in items:
    print(f"Key: {item.key}, Namespace: {item.namespace}")

# Delete a key
store.delete(("users", "123"), "preferences")

3. Using the Checkpointer for Graph Persistence

from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, BaseMessage
from typing_extensions import TypedDict
from typing import Annotated
from langgraph_unity_catalog_checkpoint import UnityCatalogCheckpointSaver

# Define your graph state
class State(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]

# Create a simple chatbot node
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")

def chatbot(state: State) -> dict:
    response = llm.invoke(state["messages"])
    return {"messages": [response]}

# Create the checkpointer
workspace_client = WorkspaceClient()
checkpointer = UnityCatalogCheckpointSaver(
    workspace_client=workspace_client,
    catalog="main",
    schema="langgraph",
    # Default tables: "checkpoints", "checkpoint_blobs", "checkpoint_writes"
    warehouse_id="your-warehouse-id",  # Optional
)

# Build the graph
graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)

# Compile with checkpointer for persistence
graph = graph_builder.compile(checkpointer=checkpointer)

# Run conversation with persistence
config = {"configurable": {"thread_id": "conversation_1"}}

# First interaction
result = graph.invoke(
    {"messages": [HumanMessage(content="Hello! What's the weather like?")]},
    config=config
)

# Second interaction - conversation history is maintained!
result = graph.invoke(
    {"messages": [HumanMessage(content="What did I just ask you?")]},
    config=config
)
# The bot remembers the previous question! 🎉

4. Async Usage for High Performance

from langgraph_unity_catalog_checkpoint import AsyncUnityCatalogCheckpointSaver
import asyncio

# Create async checkpointer
async_checkpointer = AsyncUnityCatalogCheckpointSaver(
    workspace_client=workspace_client,
    catalog="main",
    schema="langgraph",
    warehouse_id="your-warehouse-id",
)

# Async chatbot node
async def async_chatbot(state: State) -> dict:
    response = await llm.ainvoke(state["messages"])
    return {"messages": [response]}

# Build and compile with async checkpointer
graph_builder = StateGraph(State)
graph_builder.add_node("chatbot", async_chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile(checkpointer=async_checkpointer)

# Run asynchronously
config = {"configurable": {"thread_id": "async_conversation_1"}}
result = await graph.ainvoke(
    {"messages": [HumanMessage(content="Hello async world!")]},
    config=config
)

---

🎯 Use Cases

1. Conversational AI with Memory

Maintain conversation history across multiple interactions:

# Each user gets their own conversation thread
config = {"configurable": {"thread_id": f"user_{user_id}"}}
graph.invoke({"messages": [HumanMessage(content=user_input)]}, config)

2. Human-in-the-Loop Workflows

Pause execution for human review and resume seamlessly:

# Interrupt before critical nodes
graph = builder.compile(
    checkpointer=checkpointer,
    interrupt_before=["approval_node"]
)

# Execute and pause at approval
result = graph.invoke(input_data, config)

# Human reviews and approves...

# Resume from checkpoint
result = graph.invoke(None, config)  # Continues from where it left off

3. Long-Term Memory with LangMem

Integrate with LangMem for user preferences and memories:

from langchain.agents import create_agent
from langmem.tools import get_langmem_tools

# Create store for LangMem
store = UnityCatalogStore(
    workspace_client=workspace_client,
    catalog="main",
    schema="langgraph",
)

# Get LangMem tools
langmem_tools = get_langmem_tools(store=store)

# Create agent with memory
agent = create_agent(llm, tools + langmem_tools)

# Use with user context
config = {
    "configurable": {
        "langgraph_user_id": "user_123"  # Isolates memories per user
    }
}
agent.invoke({"messages": [HumanMessage(content="I prefer dark mode")]}, config)

4. Production ML Pipelines

Reliable state management for complex workflows:

# Automatic recovery from failures
# Time-travel debugging with Delta Lake
# Full audit trail via Unity Catalog
# Multi-agent coordination with isolated states

---

📊 Performance Optimizations

Batch Write Operations (2-10x Faster)

The implementation uses batched SQL operations to minimize round trips to Unity Catalog:

# Instead of N+1 SQL statements:
# - 1 per blob
# - 1 per write
# - 1 checkpoint

# We use just 3 SQL statements:
# - 1 batch for all blobs
# - 1 batch for all writes  
# - 1 for checkpoint

# For a checkpoint with 5 blobs and 3 writes:
# Before: 9 SQL statements
# After: 3 SQL statements
# Speedup: 3x faster! ⚡

See docs/CHECKPOINT_BATCH_WRITE_OPTIMIZATION.md for details.

---

🏗️ Architecture

┌──────────────────────────────────────┐
│  LangChain/LangGraph Application     │
└──────────────────────────────────────┘
                 ↓
┌──────────────────────────────────────┐
│  BaseStore / BaseCheckpointSaver     │
│  (LangGraph Interfaces)              │
└──────────────────────────────────────┘
                 ↓
┌──────────────────────────────────────┐
│  Unity Catalog Implementation        │
│  - UnityCatalogStore                 │
│  - UnityCatalogCheckpointSaver       │
│  - Async variants                    │
└──────────────────────────────────────┘
                 ↓
┌──────────────────────────────────────┐
│  Databricks WorkspaceClient          │
│  (SQL Statement Execution API)       │
└──────────────────────────────────────┘
                 ↓
┌──────────────────────────────────────┐
│  Unity Catalog Delta Tables          │
│  - ACID transactions                 │
│  - Time-travel                       │
│  - Change Data Feed                  │
└──────────────────────────────────────┘

Data Storage

• Serialization: Checkpoints and values are serialized using LangGraph's JsonPlusSerializer
• Binary Storage: BINARY columns for efficient blob storage (base64 encoded)
• JSON Metadata: Structured metadata for filtering and querying
• Delta Lake: ACID transactions, time-travel, and optimization

Default Table Names

Component	Default Tables
Store	store
Checkpointer	checkpoints, checkpoint_blobs, checkpoint_writes

Tables are automatically created on first use with optimized schemas.

---

📚 Examples

Complete Jupyter Notebooks

Explore the notebooks/ directory for interactive examples:

• store_example.ipynb - Store operations and LangMem integration
• checkpointer_example.ipynb - Synchronous graph checkpointing
• async_checkpointer_example.ipynb - Async graph execution

Run in Databricks

1. Upload a notebook to your Databricks workspace
2. Attach to a cluster with Unity Catalog access
3. Set the required configuration (catalog, schema, warehouse_id)
4. Run all cells

---

🔧 Configuration

Environment Variables

Variable	Description	Required
DATABRICKS_HOST	Workspace URL	Yes
DATABRICKS_TOKEN	Access token	Yes
DATABRICKS_WAREHOUSE_ID	SQL warehouse ID	No
UC_CATALOG	Default catalog name	Recommended
UC_SCHEMA	Default schema name	Recommended

Configuration Precedence

Configuration values are resolved in this order:

1. Environment variables (highest priority)
2. Databricks widgets (for notebooks)
3. Constructor parameters (explicit values)

See docs/CONFIGURATION_PRECEDENCE.md for details.

Warehouse ID

The warehouse_id parameter is optional and defaults to None. If not provided:

• Uses the default warehouse for the workspace
• Can be overridden per-operation if needed

---

🔒 Permissions Required

Ensure your Databricks principal has:

• USE CATALOG on the target catalog
• USE SCHEMA on the target schema
• CREATE TABLE on the target schema (for initialization)
• SELECT, INSERT, UPDATE, DELETE, MODIFY on the tables

---

🧪 Testing

Run Unit Tests

# Run all tests
make test

# Run specific test file
uv run pytest tests/test_unity_catalog_store.py -v

# Run with coverage
uv run pytest --cov=src --cov-report=html

Run Integration Tests

Integration tests require a live Databricks connection:

# Set required environment variables
export DATABRICKS_HOST="..."
export DATABRICKS_TOKEN="..."
export DATABRICKS_WAREHOUSE_ID="..."

# Run integration tests
uv run pytest tests/test_integration.py -v

Linting and Formatting

# Format code
make format

# Run linting
make lint

# Type checking
make type-check

---

📖 Documentation

Core Documentation

• Usage Guide - Comprehensive usage examples
• Implementation Summary - Technical architecture
• Environment Setup - Development environment
• Quick Start - Getting started guide
• Install Guide - Installation instructions

Technical Details

• Checkpoint Batch Write Optimization - Performance optimization details
• Configuration Precedence - Configuration resolution
• Default Table Names - Table naming conventions
• MLflow Autolog Setup - Observability with MLflow
• Logging - Logging configuration

Session Summaries

• Batch Optimization (2025-11-07)
• MLflow Tracing Removal (2025-11-07)

---

🚀 Features

UnityCatalogStore

• ✅ Implements langgraph.store.base.BaseStore interface
• ✅ Batch operations (batch, abatch) for performance
• ✅ Namespaced key-value storage
• ✅ Search with filtering and pagination
• ✅ Automatic table initialization
• ✅ Sync and async implementations
• ✅ Compatible with LangMem for long-term memory

UnityCatalogCheckpointSaver

• ✅ Implements BaseCheckpointSaver interface
• ✅ Full LangGraph checkpoint persistence
• ✅ Support for human-in-the-loop workflows
• ✅ Multi-turn conversation memory
• ✅ State recovery and time-travel
• ✅ Pending writes management
• ✅ Checkpoint listing and filtering
• ✅ Sync and async implementations
• ✅ Optimized batch writes (2-10x faster)
• ✅ Automatic table creation and schema management

---

🛠️ Development

Setup Development Environment

# Clone the repository
git clone https://github.com/natefleming/langgraph_unity_catalog_checkpoint.git
cd langgraph_unity_catalog_checkpoint

# Create virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install with development dependencies
pip install -e ".[dev]"

# Install pre-commit hooks
pre-commit install

Project Structure

langgraph_unity_catalog_checkpoint/
├── src/
│   └── langgraph_unity_catalog_checkpoint/
│       ├── store/              # Store implementations
│       │   ├── unity_catalog.py    # Sync store
│       │   ├── aio.py              # Async store
│       │   └── base.py             # Base store class
│       ├── checkpoint/         # Checkpointer implementations
│       │   ├── unity_catalog.py    # Sync checkpointer
│       │   ├── aio.py              # Async checkpointer
│       │   └── base.py             # Base checkpointer class
│       └── __init__.py         # Public API exports
├── tests/                      # Test suite
├── notebooks/                  # Example notebooks
├── docs/                       # Documentation
├── pyproject.toml             # Project configuration
└── README.md                  # This file

---

🤝 Contributing

Contributions are welcome! Please:

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Make your changes with tests
4. Run the test suite (make test)
5. Format and lint (make format lint)
6. Commit your changes (git commit -m 'Add amazing feature')
7. Push to the branch (git push origin feature/amazing-feature)
8. Open a Pull Request

---

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

---

🙏 Acknowledgments

Built with:

• LangChain - Framework for LLM applications
• LangGraph - Graph-based agent framework
• LangMem - Long-term memory for agents
• Databricks SDK - Databricks API client
• Unity Catalog - Data governance platform

---

📞 Support

For issues and questions:

• GitHub Issues: Open an issue
• Documentation: Check the docs/ directory
• Examples: Review the notebooks/ directory

---

🗺️ Roadmap

Planned enhancements:

• Connection pooling for improved performance
• Configurable TTL for automatic checkpoint cleanup
• Metrics and monitoring integration
• Query optimization hints and caching
• Support for alternative serialization formats
• Bulk import/export utilities
• Multi-region replication support

---

⚡ Quick Links

• Quick Start Guide - Get started in 5 minutes
• Usage Examples - Detailed usage patterns
• Notebooks - Interactive examples
• API Reference - Technical details

---

Made with ❤️ for the LangChain community