DatabricksStore: PostgresStore Wrapper SDK

.github/workflows/main.yml

@@ -92,6 +92,7 @@ jobs:
         run: |
           pytest tests/databricks_ai_bridge/test_lakebase.py
           pytest integrations/langchain/tests/unit_tests/test_checkpoint.py
+          pytest integrations/langchain/tests/unit_tests/test_store.py
 
   langchain_cross_version_test:
     runs-on: ubuntu-latest

integrations/langchain/src/databricks_langchain/__init__.py

@@ -21,6 +21,7 @@
 from databricks_langchain.checkpoint import CheckpointSaver
 from databricks_langchain.embeddings import DatabricksEmbeddings
 from databricks_langchain.genie import GenieAgent
+from databricks_langchain.store import DatabricksStore
 from databricks_langchain.vector_search_retriever_tool import VectorSearchRetrieverTool
 from databricks_langchain.vectorstores import DatabricksVectorSearch
 
@@ -29,6 +30,7 @@
     "ChatDatabricks",
     "CheckpointSaver",
     "DatabricksEmbeddings",
+    "DatabricksStore",
     "DatabricksVectorSearch",
     "GenieAgent",
     "VectorSearchRetrieverTool",

integrations/langchain/src/databricks_langchain/store.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import Any, Iterable, Optional
+
+from databricks.sdk import WorkspaceClient
+
+try:
+    from databricks_ai_bridge.lakebase import LakebasePool
+    from langgraph.store.base import BaseStore, Item, Op, Result
+    from langgraph.store.postgres import PostgresStore
+
+    _store_imports_available = True
+except ImportError:
+    LakebasePool = object
+    PostgresStore = object
+    BaseStore = object
+    Item = object
+    Op = object
+    Result = object
+    _store_imports_available = False
+
+
+class DatabricksStore(BaseStore):
+    """Provides APIs for working with long-term memory on Databricks using Lakebase.
+    Extends LangGraph BaseStore interface using Databricks Lakebase for connection pooling.
+
+    Operations borrow a connection from the pool, create a short-lived PostgresStore,
+    execute the operation, and return the connection to the pool.
+    """
+
+    def __init__(
+        self,
+        *,
+        instance_name: str,
+        workspace_client: Optional[WorkspaceClient] = None,
+        **pool_kwargs: Any,
+    ) -> None:
+        if not _store_imports_available:
+            raise ImportError(
+                "DatabricksStore requires databricks-langchain[memory]. "
+                "Install with: pip install 'databricks-langchain[memory]'"
+            )
+
+        # Store initialization parameters for lazy initialization, otherwise
+        # if we directly initialize pool during deployment it will fail
+        self._instance_name = instance_name
+        self._workspace_client = workspace_client
+        self._pool_kwargs = pool_kwargs
+        self._lakebase: Optional[LakebasePool] = None
+        self._pool = None
+
+    def _ensure_initialized(self) -> None:
+        """Lazy initialization of LakebasePool on first use after deployment is ready."""
+        if self._lakebase is None:
+            self._lakebase = LakebasePool(
+                instance_name=self._instance_name,
+                workspace_client=self._workspace_client,
+                **self._pool_kwargs,
+            )
+            self._pool = self._lakebase.pool
+
+    def _with_store(self, fn, *args, **kwargs):
+        """
+        Borrow a connection, create a short-lived PostgresStore, call fn(store),
+        then return the connection to the pool.
+        """
+        self._ensure_initialized()
+        with self._pool.connection() as conn:
+            store = PostgresStore(conn=conn)
+            return fn(store, *args, **kwargs)
+
+    def setup(self) -> None:
+        """Instantiate the store, setting up necessary persistent storage."""
+        return self._with_store(lambda s: s.setup())
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        """Execute a batch of operations synchronously.
+
+        This is the core method required by BaseStore. All other operations
+        (get, put, search, delete, list_namespaces) are inherited from BaseStore
+        and internally call this batch() method.
+        """
+        return self._with_store(lambda s: s.batch(ops))
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        """Execute a batch of operations asynchronously.
+
+        This is the second abstract method required by BaseStore.
+        Currently delegates to sync batch() - for true async support,
+        would need async-compatible connection pooling.
+        """
+        return self.batch(ops)

integrations/langchain/tests/unit_tests/test_store.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+pytest.importorskip("psycopg")
+pytest.importorskip("psycopg_pool")
+pytest.importorskip("langgraph.checkpoint.postgres")
+
+from databricks_ai_bridge import lakebase
+
+from databricks_langchain import DatabricksStore
+
+
+class TestConnectionPool:
+    def __init__(self, connection_value="conn"):
+        self.connection_value = connection_value
+        self.conninfo = ""
+
+    def __call__(
+        self,
+        *,
+        conninfo,
+        connection_class=None,
+        **kwargs,
+    ):
+        self.conninfo = conninfo
+        return self
+
+    def connection(self):
+        class _Ctx:
+            def __init__(self, outer):
+                self.outer = outer
+
+            def __enter__(self):
+                return self.outer.connection_value
+
+            def __exit__(self, exc_type, exc, tb):
+                pass
+
+        return _Ctx(self)
+
+
+def test_databricks_store_configures_lakebase(monkeypatch):
+    mock_conn = MagicMock()
+    test_pool = TestConnectionPool(connection_value=mock_conn)
+    monkeypatch.setattr(lakebase, "ConnectionPool", test_pool)
+
+    from langgraph.store.postgres import PostgresStore
+
+    monkeypatch.setattr(PostgresStore, "setup", MagicMock())
+
+    workspace = MagicMock()
+    workspace.database.generate_database_credential.return_value = MagicMock(token="stub-token")
+    workspace.database.get_database_instance.return_value.read_write_dns = "db-host"
+    workspace.current_service_principal.me.side_effect = RuntimeError("no sp")
+    workspace.current_user.me.return_value = MagicMock(user_name="[email protected]")
+
+    store = DatabricksStore(
+        instance_name="lakebase-instance",
+        workspace_client=workspace,
+    )
+
+    store.setup()
+
+    assert (
+        test_pool.conninfo
+        == "dbname=databricks_postgres [email protected] host=db-host port=5432 sslmode=require"
+    )
+    assert isinstance(store, DatabricksStore)
+    assert store._lakebase.pool == test_pool
+
+    with store._lakebase.connection() as conn:
+        assert conn == mock_conn