Add Postgres Store Implementation (#1906)

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -17,9 +17,7 @@ from langgraph.checkpoint.base import (
     CheckpointTuple,
     get_checkpoint_id,
 )
-from langgraph.checkpoint.postgres.base import (
-    BasePostgresSaver,
-)
+from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
 Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
@@ -167,15 +165,17 @@ class PostgresSaver(BasePostgresSaver):
                         value["pending_sends"],
                     ),
                     self._load_metadata(value["metadata"]),
-                    {
-                        "configurable": {
-                            "thread_id": value["thread_id"],
-                            "checkpoint_ns": value["checkpoint_ns"],
-                            "checkpoint_id": value["parent_checkpoint_id"],
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": value["thread_id"],
+                                "checkpoint_ns": value["checkpoint_ns"],
+                                "checkpoint_id": value["parent_checkpoint_id"],
+                            }
                         }
-                    }
-                    if value["parent_checkpoint_id"]
-                    else None,
+                        if value["parent_checkpoint_id"]
+                        else None
+                    ),
                     self._load_writes(value["pending_writes"]),
                 )
 
@@ -246,15 +246,17 @@ class PostgresSaver(BasePostgresSaver):
                         value["pending_sends"],
                     ),
                     self._load_metadata(value["metadata"]),
-                    {
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": value["parent_checkpoint_id"],
+                    (
+                        {
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": value["parent_checkpoint_id"],
+                            }
                         }
-                    }
-                    if value["parent_checkpoint_id"]
-                    else None,
+                        if value["parent_checkpoint_id"]
+                        else None
+                    ),
                     self._load_writes(value["pending_writes"]),
                 )
 
@@ -384,3 +386,6 @@ class PostgresSaver(BasePostgresSaver):
             else:
                 with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
+
+
+__all__ = ["PostgresSaver", "Conn"]

libs/checkpoint-postgres/langgraph/store/postgres/__init__.py

@@ -0,0 +1,4 @@
+from langgraph.store.postgres.aio import AsyncPostgresStore
+from langgraph.store.postgres.base import PostgresStore
+
+__all__ = ["AsyncPostgresStore", "PostgresStore"]

libs/checkpoint-postgres/langgraph/store/postgres/aio.py

@@ -0,0 +1,190 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Iterable, Sequence, cast
+
+from psycopg import AsyncConnection, AsyncCursor
+from psycopg.errors import UndefinedTable
+from psycopg.rows import dict_row
+
+from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
+from langgraph.store.postgres.base import (
+    BasePostgresStore,
+    Row,
+    _group_ops,
+    _row_to_item,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
+    def __init__(self, conn: AsyncConnection[Any]) -> None:
+        self.conn = conn
+        self.loop = asyncio.get_running_loop()
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        grouped_ops, num_ops = _group_ops(ops)
+        results: list[Result] = [None] * num_ops
+
+        async with self.conn.pipeline():
+            tasks = []
+
+            if GetOp in grouped_ops:
+                tasks.append(
+                    self._batch_get_ops(
+                        cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
+                    )
+                )
+
+            if PutOp in grouped_ops:
+                tasks.append(
+                    self._batch_put_ops(
+                        cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
+                    )
+                )
+
+            if SearchOp in grouped_ops:
+                tasks.append(
+                    self._batch_search_ops(
+                        cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
+                        results,
+                    )
+                )
+
+            if ListNamespacesOp in grouped_ops:
+                tasks.append(
+                    self._batch_list_namespaces_ops(
+                        cast(
+                            Sequence[tuple[int, ListNamespacesOp]],
+                            grouped_ops[ListNamespacesOp],
+                        ),
+                        results,
+                    )
+                )
+
+            await asyncio.gather(*tasks)
+
+        return results
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
+
+    async def _batch_get_ops(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+        results: list[Result],
+    ) -> None:
+        cursors = []
+        for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
+            cur = self.conn.cursor(binary=True)
+            await cur.execute(query, params)
+            cursors.append((cur, namespace, items))
+
+        for cur, namespace, items in cursors:
+            rows = cast(list[Row], await cur.fetchall())
+            key_to_row = {row["key"]: row for row in rows}
+            for idx, key in items:
+                row = key_to_row.get(key)
+                if row:
+                    results[idx] = _row_to_item(namespace, row)
+                else:
+                    results[idx] = None
+
+    async def _batch_put_ops(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+    ) -> None:
+        queries = self._get_batch_PUT_queries(put_ops)
+        for query, params in queries:
+            cur = self.conn.cursor(binary=True)
+            await cur.execute(query, params)
+
+    async def _batch_search_ops(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_search_queries(search_ops)
+        cursors: list[tuple[AsyncCursor[Any], int, SearchOp]] = []
+
+        for (query, params), (idx, op) in zip(queries, search_ops):
+            cur = self.conn.cursor(binary=True)
+            await cur.execute(query, params)
+            cursors.append((cur, idx, op))
+
+        for cur, idx, op in cursors:
+            rows = cast(list[Row], await cur.fetchall())
+            items = [_row_to_item(op.namespace_prefix, row) for row in rows]
+            results[idx] = items
+
+    async def _batch_list_namespaces_ops(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_list_namespaces_queries(list_ops)
+        cursors: list[tuple[AsyncCursor[Any], int]] = []
+        for (query, params), (idx, _) in zip(queries, list_ops):
+            cur = self.conn.cursor(binary=True)
+            await cur.execute(query, params)
+            cursors.append((cur, idx))
+
+        for cur, idx in cursors:
+            rows = cast(list[dict], await cur.fetchall())
+            namespaces = [
+                tuple(row["truncated_prefix"].decode()[1:].split(".")) for row in rows
+            ]
+            results[idx] = namespaces
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls,
+        conn_string: str,
+    ) -> AsyncIterator["AsyncPostgresStore"]:
+        """Create a new AsyncPostgresStore instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+
+        Returns:
+            AsyncPostgresStore: A new AsyncPostgresStore instance.
+        """
+        async with await AsyncConnection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            yield cls(conn=conn)
+
+    async def setup(self) -> None:
+        """Set up the store database asynchronously.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time the store is used.
+        """
+        async with self.conn.cursor() as cur:
+            try:
+                await cur.execute(
+                    "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
+                )
+                row = cast(dict, await cur.fetchone())
+                if row is None:
+                    version = -1
+                else:
+                    version = row["v"]
+            except UndefinedTable:
+                version = -1
+                # Create store_migrations table if it doesn't exist
+                await cur.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS store_migrations (
+                        v INTEGER PRIMARY KEY
+                    )
+                    """
+                )
+            for v, migration in enumerate(
+                self.MIGRATIONS[version + 1 :], start=version + 1
+            ):
+                await cur.execute(migration)
+                await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))

libs/checkpoint-postgres/langgraph/store/postgres/base.py

@@ -0,0 +1,394 @@
+import asyncio
+import json
+import logging
+from collections import defaultdict
+from contextlib import contextmanager
+from datetime import datetime
+from typing import Any, Generic, Iterable, Iterator, Sequence, TypeVar, Union, cast
+
+import orjson
+from psycopg import BaseConnection, Connection, Cursor
+from psycopg.errors import UndefinedTable
+from psycopg.rows import dict_row
+from psycopg.types.json import Jsonb
+from typing_extensions import TypedDict
+
+from langgraph.store.base import (
+    BaseStore,
+    GetOp,
+    Item,
+    ListNamespacesOp,
+    Op,
+    PutOp,
+    Result,
+    SearchOp,
+)
+
+logger = logging.getLogger(__name__)
+
+
+MIGRATIONS = [
+    """
+CREATE EXTENSION IF NOT EXISTS ltree;
+""",
+    """
+CREATE TABLE IF NOT EXISTS store (
+    -- 'prefix' represents the doc's 'namespace'
+    prefix ltree NOT NULL,
+    key text NOT NULL,
+    value jsonb NOT NULL,
+    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
+    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (prefix, key)
+);
+""",
+    """
+-- For faster listing of namespaces & lookups by namespace with prefix/suffix matching
+CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING gist (prefix);
+""",
+]
+
+C = TypeVar("C", bound=BaseConnection)
+
+
+class BasePostgresStore(BaseStore, Generic[C]):
+    MIGRATIONS = MIGRATIONS
+    conn: C
+
+    def _get_batch_GET_ops_queries(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+    ) -> list[tuple[str, tuple, tuple[str, ...], list]]:
+        namespace_groups = defaultdict(list)
+        for idx, op in get_ops:
+            namespace_groups[op.namespace].append((idx, op.key))
+        results = []
+        for namespace, items in namespace_groups.items():
+            _, keys = zip(*items)
+            keys_to_query = ",".join(["%s"] * len(keys))
+            query = f"""
+                SELECT key, value, created_at, updated_at
+                FROM store
+                WHERE prefix = %s AND key IN ({keys_to_query})
+            """
+            params = (_namespace_to_ltree(namespace), *keys)
+            results.append((query, params, namespace, items))
+        return results
+
+    def _get_batch_PUT_queries(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+    ) -> list[tuple[str, Sequence]]:
+        inserts: list[PutOp] = []
+        deletes: list[PutOp] = []
+        for _, op in put_ops:
+            if op.value is None:
+                deletes.append(op)
+            else:
+                inserts.append(op)
+
+        queries: list[tuple[str, Sequence]] = []
+
+        if deletes:
+            namespace_groups: dict[tuple[str, ...], list[str]] = defaultdict(list)
+            for op in deletes:
+                namespace_groups[op.namespace].append(op.key)
+            for namespace, keys in namespace_groups.items():
+                placeholders = ",".join(["%s"] * len(keys))
+                query = (
+                    f"DELETE FROM store WHERE prefix = %s AND key IN ({placeholders})"
+                )
+                params = (_namespace_to_ltree(namespace), *keys)
+                queries.append((query, params))
+        if inserts:
+            values = []
+            insertion_params = []
+            for op in inserts:
+                values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
+                insertion_params.extend(
+                    [
+                        _namespace_to_ltree(op.namespace),
+                        op.key,
+                        Jsonb(op.value),
+                    ]
+                )
+            values_str = ",".join(values)
+            query = f"""
+                INSERT INTO store (prefix, key, value, created_at, updated_at)
+                VALUES {values_str}
+                ON CONFLICT (prefix, key) DO UPDATE
+                SET value = EXCLUDED.value, updated_at = CURRENT_TIMESTAMP
+            """
+            queries.append((query, insertion_params))
+
+        return queries
+
+    def _get_batch_search_queries(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+    ) -> list[tuple[str, Sequence]]:
+        queries: list[tuple[str, Sequence]] = []
+        for _, op in search_ops:
+            query = """
+                SELECT key, value, created_at, updated_at
+                FROM store
+                WHERE prefix <@ %s
+            """
+            params: list = [_namespace_to_ltree(op.namespace_prefix)]
+
+            if op.filter:
+                filter_conditions = []
+                for key, value in op.filter.items():
+                    if isinstance(value, list):
+                        filter_conditions.append("value->%s @> %s::jsonb")
+                        params.extend([key, json.dumps(value)])
+                    else:
+                        filter_conditions.append("value->%s = %s::jsonb")
+                        params.extend([key, json.dumps(value)])
+                query += " AND " + " AND ".join(filter_conditions)
+
+            query += " LIMIT %s OFFSET %s"
+            params.extend([op.limit, op.offset])
+
+            queries.append((query, params))
+        return queries
+
+    def _get_batch_list_namespaces_queries(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+    ) -> list[tuple[str, Sequence]]:
+        queries: list[tuple[str, Sequence]] = []
+        for _, op in list_ops:
+            query = "SELECT DISTINCT subltree(prefix, 0, LEAST(nlevel(prefix), %s)) AS truncated_prefix FROM store"
+            # https://www.postgresql.org/docs/current/ltree.html
+            # The length of a label path cannot exceed 65535 labels.
+            params: list[Any] = [op.max_depth if op.max_depth is not None else 65536]
+
+            conditions = []
+            if op.match_conditions:
+                for condition in op.match_conditions:
+                    if condition.match_type == "prefix":
+                        conditions.append("prefix ~ %s::lquery")
+                        lquery_pattern = f"{_namespace_to_ltree(condition.path)}.*"
+                        params.append(lquery_pattern)
+                    elif condition.match_type == "suffix":
+                        conditions.append("prefix ~ %s::lquery")
+                        lquery_pattern = f"*.{_namespace_to_ltree(condition.path)}"
+                        params.append(lquery_pattern)
+                    else:
+                        logger.warning(
+                            f"Unknown match_type in list_namespaces: {condition.match_type}"
+                        )
+
+            if conditions:
+                query += " WHERE " + " AND ".join(conditions)
+
+            query += " ORDER BY truncated_prefix LIMIT %s OFFSET %s"
+            params.extend([op.limit, op.offset])
+
+            queries.append((query, params))
+        return queries
+
+
+class PostgresStore(BasePostgresStore[Connection]):
+    def __init__(self, conn: Connection[Any]) -> None:
+        self.conn = conn
+
+    def batch(self, ops: Iterable[Op]) -> list[Result]:
+        grouped_ops, num_ops = _group_ops(ops)
+        results: list[Result] = [None] * num_ops
+
+        with self.conn.pipeline():
+            if GetOp in grouped_ops:
+                self._batch_get_ops(
+                    cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
+                )
+
+            if PutOp in grouped_ops:
+                self._batch_put_ops(
+                    cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
+                )
+
+            if SearchOp in grouped_ops:
+                self._batch_search_ops(
+                    cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
+                    results,
+                )
+
+            if ListNamespacesOp in grouped_ops:
+                self._batch_list_namespaces_ops(
+                    cast(
+                        Sequence[tuple[int, ListNamespacesOp]],
+                        grouped_ops[ListNamespacesOp],
+                    ),
+                    results,
+                )
+
+        return results
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
+
+    def _batch_get_ops(
+        self,
+        get_ops: Sequence[tuple[int, GetOp]],
+        results: list[Result],
+    ) -> None:
+        cursors = []
+        for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
+            cur = self.conn.cursor(binary=True)
+            cur.execute(query, params)
+            cursors.append((cur, namespace, items))
+
+        for cur, namespace, items in cursors:
+            rows = cast(list[Row], cur.fetchall())
+            key_to_row = {row["key"]: row for row in rows}
+            for idx, key in items:
+                row = key_to_row.get(key)
+                if row:
+                    results[idx] = _row_to_item(namespace, row)
+                else:
+                    results[idx] = None
+
+    def _batch_put_ops(
+        self,
+        put_ops: Sequence[tuple[int, PutOp]],
+    ) -> None:
+        queries = self._get_batch_PUT_queries(put_ops)
+        for query, params in queries:
+            cur = self.conn.cursor(binary=True)
+            cur.execute(query, params)
+
+    def _batch_search_ops(
+        self,
+        search_ops: Sequence[tuple[int, SearchOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_search_queries(search_ops)
+        cursors: list[tuple[Cursor[Any], int, SearchOp]] = []
+
+        for (query, params), (idx, op) in zip(queries, search_ops):
+            cur = self.conn.cursor(binary=True)
+            cur.execute(query, params)
+            cursors.append((cur, idx, op))
+
+        for cur, idx, op in cursors:
+            rows = cast(list[Row], cur.fetchall())
+            items = [_row_to_item(op.namespace_prefix, row) for row in rows]
+            results[idx] = items
+
+    def _batch_list_namespaces_ops(
+        self,
+        list_ops: Sequence[tuple[int, ListNamespacesOp]],
+        results: list[Result],
+    ) -> None:
+        queries = self._get_batch_list_namespaces_queries(list_ops)
+        cursors: list[tuple[Cursor[Any], int]] = []
+        for (query, params), (idx, _) in zip(queries, list_ops):
+            cur = self.conn.cursor(binary=True)
+            cur.execute(query, params)
+            cursors.append((cur, idx))
+
+        for cur, idx in cursors:
+            rows = cast(list[dict], cur.fetchall())
+            namespaces = [
+                tuple(row["truncated_prefix"].decode()[1:].split(".")) for row in rows
+            ]
+            results[idx] = namespaces
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(
+        cls,
+        conn_string: str,
+    ) -> Iterator["PostgresStore"]:
+        """Create a new BasePostgresStore instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+
+        Returns:
+            BasePostgresStore: A new BasePostgresStore instance.
+        """
+        with Connection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            yield cls(conn=conn)
+
+    def setup(self) -> None:
+        """Set up the store database.
+
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time the store is used.
+        """
+        with self.conn.cursor(binary=True) as cur:
+            try:
+                cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
+                row = cast(dict, cur.fetchone())
+                if row is None:
+                    version = -1
+                else:
+                    version = row["v"]
+            except UndefinedTable:
+                self.conn.rollback()
+                version = -1
+                # Create store_migrations table if it doesn't exist
+                cur.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS store_migrations (
+                        v INTEGER PRIMARY KEY
+                    )
+                """
+                )
+            for v, migration in enumerate(
+                self.MIGRATIONS[version + 1 :], start=version + 1
+            ):
+                cur.execute(migration)
+                cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
+
+
+class Row(TypedDict):
+    key: str
+    value: Any
+    prefix: bytes
+    created_at: datetime
+    updated_at: datetime
+
+
+def _namespace_to_ltree(namespace: tuple[str, ...]) -> str:
+    """Convert namespace tuple to ltree-compatible string."""
+    return ".".join(namespace)
+
+
+def _row_to_item(namespace: tuple[str, ...], row: Row) -> Item:
+    """Convert a row from the database into an Item."""
+    val = row["value"]
+    return Item(
+        value=val if isinstance(val, dict) else _json_loads(val),
+        key=row["key"],
+        namespace=namespace,
+        created_at=row["created_at"],
+        updated_at=row["updated_at"],
+    )
+
+
+def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
+    grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
+    tot = 0
+    for idx, op in enumerate(ops):
+        grouped_ops[type(op)].append((idx, op))
+        tot += 1
+    return grouped_ops, tot
+
+
+def _json_loads(content: Union[bytes, orjson.Fragment]) -> Any:
+    if isinstance(content, orjson.Fragment):
+        if hasattr(content, "buf"):
+            content = content.buf
+        else:
+            if isinstance(content.contents, bytes):
+                content = content.contents
+            else:
+                content = content.contents.encode()
+    return orjson.loads(cast(bytes, content))

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.9"
+version = "1.0.10"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"

libs/checkpoint-postgres/tests/conftest.py

@@ -23,5 +23,7 @@ async def clear_test_db(conn: AsyncConnection[DictRow]) -> None:
         await conn.execute("DELETE FROM checkpoints")
         await conn.execute("DELETE FROM checkpoint_blobs")
         await conn.execute("DELETE FROM checkpoint_writes")
+        await conn.execute("DELETE FROM checkpoint_migrations")
+        await conn.execute("DELETE FROM store_migrations")
     except UndefinedTable:
         pass

libs/checkpoint-postgres/tests/test_async.py

@@ -107,7 +107,7 @@ class TestAsyncPostgresSaver:
             config = await saver.aput(
                 self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {}
             )
-            assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc"
+            assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc"  # type: ignore
             assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
                 0
             ].metadata["my_key"] == "abc"

libs/checkpoint-postgres/tests/test_async_store.py

@@ -0,0 +1,465 @@
+# type: ignore
+import uuid
+from datetime import datetime
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from conftest import DEFAULT_URI  # type: ignore
+
+from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
+from langgraph.store.postgres import AsyncPostgresStore
+
+
+class MockAsyncCursor:
+    def __init__(self, fetch_result: Any) -> None:
+        self.fetch_result = fetch_result
+        self.execute = AsyncMock()
+        self.fetchall = AsyncMock(return_value=self.fetch_result)
+
+
+class MockAsyncConnection:
+    def __init__(self) -> None:
+        self.cursor = MagicMock()
+        self.pipeline = MagicMock(
+            return_value=AsyncMock(__aenter__=AsyncMock(), __aexit__=AsyncMock())
+        )
+
+
+@pytest.fixture
+def mock_connection() -> MockAsyncConnection:
+    return MockAsyncConnection()
+
+
+@pytest.fixture
+async def store(mock_connection: MockAsyncConnection) -> AsyncPostgresStore:
+    return AsyncPostgresStore(mock_connection)
+
+
+async def test_abatch_order(store: AsyncPostgresStore) -> None:
+    mock_connection = store.conn
+    mock_get_cursor = MockAsyncCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+            {
+                "key": "key2",
+                "value": '{"data": "value2"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_search_cursor = MockAsyncCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_list_namespaces_cursor = MockAsyncCursor(
+        [
+            {"truncated_prefix": b"\x01test"},
+        ]
+    )
+
+    failures = []
+
+    def cursor_side_effect(binary: bool = False) -> Any:
+        cursor = MagicMock()
+
+        async def execute_side_effect(query: str, *params: Any) -> None:
+            # My super sophisticated database.
+            if "WHERE prefix <@" in query:
+                cursor.fetchall = mock_search_cursor.fetchall
+            elif "SELECT DISTINCT subltree" in query:
+                cursor.fetchall = mock_list_namespaces_cursor.fetchall
+            elif "WHERE prefix = %s AND key" in query:
+                cursor.fetchall = mock_get_cursor.fetchall
+            elif "INSERT INTO " in query:
+                pass
+            else:
+                e = ValueError(f"Unmatched query: {query}")
+                failures.append(e)
+                raise e
+
+        cursor.execute = AsyncMock(side_effect=execute_side_effect)
+        return cursor
+
+    mock_connection.cursor.side_effect = cursor_side_effect  # type: ignore
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+    results = await store.abatch(ops)
+    assert not failures
+    assert len(results) == 5
+    assert isinstance(results[0], Item)
+    assert isinstance(results[0].value, dict)
+    assert results[0].value == {"data": "value1"}
+    assert results[0].key == "key1"
+    assert results[1] is None
+    assert isinstance(results[2], list)
+    assert len(results[2]) == 1
+    assert isinstance(results[3], list)
+    assert results[3] == [("test",)]
+    assert results[4] is None
+
+    ops_reordered = [
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+        GetOp(namespace=("test",), key="key2"),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
+        PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
+        GetOp(namespace=("test",), key="key1"),
+    ]
+
+    results_reordered = await store.abatch(ops_reordered)
+    assert not failures
+    assert len(results_reordered) == 5
+    assert isinstance(results_reordered[0], list)
+    assert len(results_reordered[0]) == 1
+    assert isinstance(results_reordered[1], Item)
+    assert results_reordered[1].value == {"data": "value2"}
+    assert results_reordered[1].key == "key2"
+    assert isinstance(results_reordered[2], list)
+    assert results_reordered[2] == [("test",)]
+    assert results_reordered[3] is None
+    assert isinstance(results_reordered[4], Item)
+    assert results_reordered[4].value == {"data": "value1"}
+    assert results_reordered[4].key == "key1"
+
+
+async def test_batch_get_ops(store: AsyncPostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockAsyncCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+            {
+                "key": "key2",
+                "value": '{"data": "value2"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        GetOp(namespace=("test",), key="key2"),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 3
+    assert results[0] is not None
+    assert results[1] is not None
+    assert results[2] is None
+    assert results[0].key == "key1"
+    assert results[1].key == "key2"
+
+
+async def test_batch_put_ops(store: AsyncPostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockAsyncCursor([])
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        PutOp(namespace=("test",), key="key3", value=None),
+    ]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 3
+    assert all(result is None for result in results)
+    assert mock_cursor.execute.call_count == 2
+
+
+async def test_batch_search_ops(store: AsyncPostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockAsyncCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+            {
+                "key": "key2",
+                "value": '{"data": "value2"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+    ]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 2
+    assert len(results[0]) == 2
+    assert len(results[1]) == 2
+
+
+async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockAsyncCursor(
+        [
+            {"truncated_prefix": b"\x01test.namespace1"},
+            {"truncated_prefix": b"\x01test.namespace2"},
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
+
+    results = await store.abatch(ops)
+
+    assert len(results) == 1
+    assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
+
+
+# The following use the actual DB connection
+
+
+class TestAsyncPostgresStore:
+    @pytest.fixture(autouse=True)
+    async def setup(self) -> None:
+        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
+            await store.setup()
+
+    async def test_basic_store_ops(self) -> None:
+        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
+            namespace = ("test", "documents")
+            item_id = "doc1"
+            item_value = {"title": "Test Document", "content": "Hello, World!"}
+
+            await store.aput(namespace, item_id, item_value)
+            item = await store.aget(namespace, item_id)
+
+            assert item
+            assert item.namespace == namespace
+            assert item.key == item_id
+            assert item.value == item_value
+
+            updated_value = {
+                "title": "Updated Test Document",
+                "content": "Hello, LangGraph!",
+            }
+            await store.aput(namespace, item_id, updated_value)
+            updated_item = await store.aget(namespace, item_id)
+
+            assert updated_item.value == updated_value
+            assert updated_item.updated_at > item.updated_at
+            different_namespace = ("test", "other_documents")
+            item_in_different_namespace = await store.aget(different_namespace, item_id)
+            assert item_in_different_namespace is None
+
+            new_item_id = "doc2"
+            new_item_value = {"title": "Another Document", "content": "Greetings!"}
+            await store.aput(namespace, new_item_id, new_item_value)
+
+            search_results = await store.asearch(["test"], limit=10)
+            items = search_results
+            assert len(items) == 2
+            assert any(item.key == item_id for item in items)
+            assert any(item.key == new_item_id for item in items)
+
+            namespaces = await store.alist_namespaces(prefix=["test"])
+            assert ("test", "documents") in namespaces
+
+            await store.adelete(namespace, item_id)
+            await store.adelete(namespace, new_item_id)
+            deleted_item = await store.aget(namespace, item_id)
+            assert deleted_item is None
+
+            deleted_item = await store.aget(namespace, new_item_id)
+            assert deleted_item is None
+
+            empty_search_results = await store.asearch(["test"], limit=10)
+            assert len(empty_search_results) == 0
+
+    async def test_list_namespaces(self) -> None:
+        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
+            test_pref = str(uuid.uuid4())
+            test_namespaces = [
+                (test_pref, "test", "documents", "public", test_pref),
+                (test_pref, "test", "documents", "private", test_pref),
+                (test_pref, "test", "images", "public", test_pref),
+                (test_pref, "test", "images", "private", test_pref),
+                (test_pref, "prod", "documents", "public", test_pref),
+                (
+                    test_pref,
+                    "prod",
+                    "documents",
+                    "some",
+                    "nesting",
+                    "public",
+                    test_pref,
+                ),
+                (test_pref, "prod", "documents", "private", test_pref),
+            ]
+
+            for namespace in test_namespaces:
+                await store.aput(namespace, "dummy", {"content": "dummy"})
+
+            prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
+            assert len(prefix_result) == 4
+            assert all([ns[1] == "test" for ns in prefix_result])
+
+            specific_prefix_result = await store.alist_namespaces(
+                prefix=[test_pref, "test", "documents"]
+            )
+            assert len(specific_prefix_result) == 2
+            assert all(
+                [ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
+            )
+
+            suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
+            assert len(suffix_result) == 4
+            assert all(ns[-2] == "public" for ns in suffix_result)
+
+            prefix_suffix_result = await store.alist_namespaces(
+                prefix=[test_pref, "test"], suffix=["public", test_pref]
+            )
+            assert len(prefix_suffix_result) == 2
+            assert all(
+                ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
+            )
+
+            wildcard_prefix_result = await store.alist_namespaces(
+                prefix=[test_pref, "*", "documents"]
+            )
+            assert len(wildcard_prefix_result) == 5
+            assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
+
+            wildcard_suffix_result = await store.alist_namespaces(
+                suffix=["*", "public", test_pref]
+            )
+            assert len(wildcard_suffix_result) == 4
+            assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
+            wildcard_single = await store.alist_namespaces(
+                suffix=["some", "*", "public", test_pref]
+            )
+            assert len(wildcard_single) == 1
+            assert wildcard_single[0] == (
+                test_pref,
+                "prod",
+                "documents",
+                "some",
+                "nesting",
+                "public",
+                test_pref,
+            )
+
+            max_depth_result = await store.alist_namespaces(max_depth=3)
+            assert all([len(ns) <= 3 for ns in max_depth_result])
+
+            max_depth_result = await store.alist_namespaces(
+                max_depth=4, prefix=[test_pref, "*", "documents"]
+            )
+            assert (
+                len(set(tuple(res) for res in max_depth_result))
+                == len(max_depth_result)
+                == 5
+            )
+
+            limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
+            assert len(limit_result) == 3
+
+            offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
+            assert len(offset_result) == len(test_namespaces) - 3
+
+            empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
+            assert len(empty_prefix_result) == len(test_namespaces)
+            assert set(tuple(ns) for ns in empty_prefix_result) == set(
+                tuple(ns) for ns in test_namespaces
+            )
+
+            for namespace in test_namespaces:
+                await store.adelete(namespace, "dummy")
+
+    async def test_search(self):
+        async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
+            test_namespaces = [
+                ("test_search", "documents", "user1"),
+                ("test_search", "documents", "user2"),
+                ("test_search", "reports", "department1"),
+                ("test_search", "reports", "department2"),
+            ]
+            test_items = [
+                {"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
+                {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
+                {"title": "Report A", "author": "John Doe", "tags": ["final"]},
+                {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
+            ]
+
+            for namespace, item in zip(test_namespaces, test_items):
+                await store.aput(namespace, f"item_{namespace[-1]}", item)
+
+            docs_result = await store.asearch(["test_search", "documents"])
+            assert len(docs_result) == 2
+            assert all(item.namespace[1] == "documents" for item in docs_result)
+
+            reports_result = await store.asearch(["test_search", "reports"])
+            assert len(reports_result) == 2
+            assert all(item.namespace[1] == "reports" for item in reports_result)
+
+            limited_result = await store.asearch(["test_search"], limit=2)
+            assert len(limited_result) == 2
+            offset_result = await store.asearch(["test_search"])
+            assert len(offset_result) == 4
+
+            offset_result = await store.asearch(["test_search"], offset=2)
+            assert len(offset_result) == 2
+            assert all(item not in limited_result for item in offset_result)
+
+            john_doe_result = await store.asearch(
+                ["test_search"], filter={"author": "John Doe"}
+            )
+            assert len(john_doe_result) == 2
+            assert all(item.value["author"] == "John Doe" for item in john_doe_result)
+
+            draft_result = await store.asearch(
+                ["test_search"], filter={"tags": ["draft"]}
+            )
+            assert len(draft_result) == 2
+            assert all("draft" in item.value["tags"] for item in draft_result)
+
+            page1 = await store.asearch(["test_search"], limit=2, offset=0)
+            page2 = await store.asearch(["test_search"], limit=2, offset=2)
+            all_items = page1 + page2
+            assert len(all_items) == 4
+            assert len(set(item.key for item in all_items)) == 4
+
+            for namespace in test_namespaces:
+                await store.adelete(namespace, f"item_{namespace[-1]}")

libs/checkpoint-postgres/tests/test_store.py

@@ -0,0 +1,458 @@
+# type: ignore
+import uuid
+from datetime import datetime
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from conftest import DEFAULT_URI  # type: ignore
+
+from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
+from langgraph.store.postgres import PostgresStore
+
+
+class MockCursor:
+    def __init__(self, fetch_result: Any) -> None:
+        self.fetch_result = fetch_result
+        self.execute = MagicMock()
+        self.fetchall = MagicMock(return_value=self.fetch_result)
+
+
+class MockConnection:
+    def __init__(self) -> None:
+        self.cursor = MagicMock()
+        self.pipeline = MagicMock()
+
+
+@pytest.fixture
+def mock_connection() -> MockConnection:
+    return MockConnection()
+
+
+@pytest.fixture
+def store(mock_connection: MockConnection) -> PostgresStore:
+    return PostgresStore(mock_connection)
+
+
+def test_batch_order(store: PostgresStore) -> None:
+    mock_connection = store.conn
+    mock_get_cursor = MockCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+            {
+                "key": "key2",
+                "value": '{"data": "value2"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_search_cursor = MockCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_list_namespaces_cursor = MockCursor(
+        [
+            {"truncated_prefix": b"\x01test"},
+        ]
+    )
+
+    failures = []
+
+    def cursor_side_effect(binary: bool = False) -> Any:
+        cursor = MagicMock()
+
+        def execute_side_effect(query: str, *params: Any) -> None:
+            # My super sophisticated database.
+            if "WHERE prefix <@" in query:
+                cursor.fetchall = mock_search_cursor.fetchall
+            elif "SELECT DISTINCT subltree" in query:
+                cursor.fetchall = mock_list_namespaces_cursor.fetchall
+            elif "WHERE prefix = %s AND key" in query:
+                cursor.fetchall = mock_get_cursor.fetchall
+            elif "INSERT INTO " in query:
+                pass
+            else:
+                e = ValueError(f"Unmatched query: {query}")
+                failures.append(e)
+                raise e
+
+        cursor.execute = MagicMock(side_effect=execute_side_effect)
+        return cursor
+
+    mock_connection.cursor.side_effect = cursor_side_effect
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+    results = store.batch(ops)
+    assert not failures
+    assert len(results) == 5
+    assert isinstance(results[0], Item)
+    assert isinstance(results[0].value, dict)
+    assert results[0].value == {"data": "value1"}
+    assert results[0].key == "key1"
+    assert results[1] is None
+    assert isinstance(results[2], list)
+    assert len(results[2]) == 1
+    assert isinstance(results[3], list)
+    assert results[3] == [("test",)]
+    assert results[4] is None
+
+    ops_reordered = [
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+        GetOp(namespace=("test",), key="key2"),
+        ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0),
+        PutOp(namespace=("test",), key="key3", value={"data": "value3"}),
+        GetOp(namespace=("test",), key="key1"),
+    ]
+
+    results_reordered = store.batch(ops_reordered)
+    assert not failures
+    assert len(results_reordered) == 5
+    assert isinstance(results_reordered[0], list)
+    assert len(results_reordered[0]) == 1
+    assert isinstance(results_reordered[1], Item)
+    assert results_reordered[1].value == {"data": "value2"}
+    assert results_reordered[1].key == "key2"
+    assert isinstance(results_reordered[2], list)
+    assert results_reordered[2] == [("test",)]
+    assert results_reordered[3] is None
+    assert isinstance(results_reordered[4], Item)
+    assert results_reordered[4].value == {"data": "value1"}
+    assert results_reordered[4].key == "key1"
+
+
+def test_batch_get_ops(store: PostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+            {
+                "key": "key2",
+                "value": '{"data": "value2"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        GetOp(namespace=("test",), key="key1"),
+        GetOp(namespace=("test",), key="key2"),
+        GetOp(namespace=("test",), key="key3"),
+    ]
+
+    results = store.batch(ops)
+
+    assert len(results) == 3
+    assert results[0] is not None
+    assert results[1] is not None
+    assert results[2] is None
+    assert results[0].key == "key1"
+    assert results[1].key == "key2"
+
+
+def test_batch_put_ops(store: PostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor([])
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        PutOp(namespace=("test",), key="key1", value={"data": "value1"}),
+        PutOp(namespace=("test",), key="key2", value={"data": "value2"}),
+        PutOp(namespace=("test",), key="key3", value=None),
+    ]
+
+    results = store.batch(ops)
+
+    assert len(results) == 3
+    assert all(result is None for result in results)
+    assert mock_cursor.execute.call_count == 2
+
+
+def test_batch_search_ops(store: PostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            {
+                "key": "key1",
+                "value": '{"data": "value1"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+            {
+                "key": "key2",
+                "value": '{"data": "value2"}',
+                "created_at": datetime.now(),
+                "updated_at": datetime.now(),
+            },
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [
+        SearchOp(
+            namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0
+        ),
+        SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0),
+    ]
+
+    results = store.batch(ops)
+
+    assert len(results) == 2
+    assert len(results[0]) == 2
+    assert len(results[1]) == 2
+
+
+def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
+    mock_connection = store.conn
+    mock_cursor = MockCursor(
+        [
+            {"truncated_prefix": b"\x01test.namespace1"},
+            {"truncated_prefix": b"\x01test.namespace2"},
+        ]
+    )
+    mock_connection.cursor.return_value = mock_cursor
+
+    ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)]
+
+    results = store.batch(ops)
+
+    assert len(results) == 1
+    assert results[0] == [("test", "namespace1"), ("test", "namespace2")]
+
+
+class TestPostgresStore:
+    @pytest.fixture(autouse=True)
+    def setup(self) -> None:
+        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
+            store.setup()
+
+    def test_basic_store_ops(self) -> None:
+        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
+            namespace = ("test", "documents")
+            item_id = "doc1"
+            item_value = {"title": "Test Document", "content": "Hello, World!"}
+
+            store.put(namespace, item_id, item_value)
+            item = store.get(namespace, item_id)
+
+            assert item
+            assert item.namespace == namespace
+            assert item.key == item_id
+            assert item.value == item_value
+
+            updated_value = {
+                "title": "Updated Test Document",
+                "content": "Hello, LangGraph!",
+            }
+            store.put(namespace, item_id, updated_value)
+            updated_item = store.get(namespace, item_id)
+
+            assert updated_item.value == updated_value
+            assert updated_item.updated_at > item.updated_at
+            different_namespace = ("test", "other_documents")
+            item_in_different_namespace = store.get(different_namespace, item_id)
+            assert item_in_different_namespace is None
+
+            new_item_id = "doc2"
+            new_item_value = {"title": "Another Document", "content": "Greetings!"}
+            store.put(namespace, new_item_id, new_item_value)
+
+            search_results = store.search(["test"], limit=10)
+            items = search_results
+            assert len(items) == 2
+            assert any(item.key == item_id for item in items)
+            assert any(item.key == new_item_id for item in items)
+
+            namespaces = store.list_namespaces(prefix=["test"])
+            assert ("test", "documents") in namespaces
+
+            store.delete(namespace, item_id)
+            store.delete(namespace, new_item_id)
+            deleted_item = store.get(namespace, item_id)
+            assert deleted_item is None
+
+            deleted_item = store.get(namespace, new_item_id)
+            assert deleted_item is None
+
+            empty_search_results = store.search(["test"], limit=10)
+            assert len(empty_search_results) == 0
+
+    def test_list_namespaces(self) -> None:
+        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
+            test_pref = str(uuid.uuid4())
+            test_namespaces = [
+                (test_pref, "test", "documents", "public", test_pref),
+                (test_pref, "test", "documents", "private", test_pref),
+                (test_pref, "test", "images", "public", test_pref),
+                (test_pref, "test", "images", "private", test_pref),
+                (test_pref, "prod", "documents", "public", test_pref),
+                (
+                    test_pref,
+                    "prod",
+                    "documents",
+                    "some",
+                    "nesting",
+                    "public",
+                    test_pref,
+                ),
+                (test_pref, "prod", "documents", "private", test_pref),
+            ]
+
+            for namespace in test_namespaces:
+                store.put(namespace, "dummy", {"content": "dummy"})
+
+            prefix_result = store.list_namespaces(prefix=[test_pref, "test"])
+            assert len(prefix_result) == 4
+            assert all([ns[1] == "test" for ns in prefix_result])
+
+            specific_prefix_result = store.list_namespaces(
+                prefix=[test_pref, "test", "documents"]
+            )
+            assert len(specific_prefix_result) == 2
+            assert all(
+                [ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
+            )
+
+            suffix_result = store.list_namespaces(suffix=["public", test_pref])
+            assert len(suffix_result) == 4
+            assert all(ns[-2] == "public" for ns in suffix_result)
+
+            prefix_suffix_result = store.list_namespaces(
+                prefix=[test_pref, "test"], suffix=["public", test_pref]
+            )
+            assert len(prefix_suffix_result) == 2
+            assert all(
+                ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
+            )
+
+            wildcard_prefix_result = store.list_namespaces(
+                prefix=[test_pref, "*", "documents"]
+            )
+            assert len(wildcard_prefix_result) == 5
+            assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
+
+            wildcard_suffix_result = store.list_namespaces(
+                suffix=["*", "public", test_pref]
+            )
+            assert len(wildcard_suffix_result) == 4
+            assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
+            wildcard_single = store.list_namespaces(
+                suffix=["some", "*", "public", test_pref]
+            )
+            assert len(wildcard_single) == 1
+            assert wildcard_single[0] == (
+                test_pref,
+                "prod",
+                "documents",
+                "some",
+                "nesting",
+                "public",
+                test_pref,
+            )
+
+            max_depth_result = store.list_namespaces(max_depth=3)
+            assert all([len(ns) <= 3 for ns in max_depth_result])
+
+            max_depth_result = store.list_namespaces(
+                max_depth=4, prefix=[test_pref, "*", "documents"]
+            )
+            assert (
+                len(set(tuple(res) for res in max_depth_result))
+                == len(max_depth_result)
+                == 5
+            )
+
+            limit_result = store.list_namespaces(prefix=[test_pref], limit=3)
+            assert len(limit_result) == 3
+
+            offset_result = store.list_namespaces(prefix=[test_pref], offset=3)
+            assert len(offset_result) == len(test_namespaces) - 3
+
+            empty_prefix_result = store.list_namespaces(prefix=[test_pref])
+            assert len(empty_prefix_result) == len(test_namespaces)
+            assert set(tuple(ns) for ns in empty_prefix_result) == set(
+                tuple(ns) for ns in test_namespaces
+            )
+
+            for namespace in test_namespaces:
+                store.delete(namespace, "dummy")
+
+    def test_search(self):
+        with PostgresStore.from_conn_string(DEFAULT_URI) as store:
+            test_namespaces = [
+                ("test_search", "documents", "user1"),
+                ("test_search", "documents", "user2"),
+                ("test_search", "reports", "department1"),
+                ("test_search", "reports", "department2"),
+            ]
+            test_items = [
+                {"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
+                {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
+                {"title": "Report A", "author": "John Doe", "tags": ["final"]},
+                {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
+            ]
+
+            for namespace, item in zip(test_namespaces, test_items):
+                store.put(namespace, f"item_{namespace[-1]}", item)
+
+            docs_result = store.search(["test_search", "documents"])
+            assert len(docs_result) == 2
+            assert all(item.namespace[1] == "documents" for item in docs_result)
+
+            reports_result = store.search(["test_search", "reports"])
+            assert len(reports_result) == 2
+            assert all(item.namespace[1] == "reports" for item in reports_result)
+
+            limited_result = store.search(["test_search"], limit=2)
+            assert len(limited_result) == 2
+            offset_result = store.search(["test_search"])
+            assert len(offset_result) == 4
+
+            offset_result = store.search(["test_search"], offset=2)
+            assert len(offset_result) == 2
+            assert all(item not in limited_result for item in offset_result)
+
+            john_doe_result = store.search(
+                ["test_search"], filter={"author": "John Doe"}
+            )
+            assert len(john_doe_result) == 2
+            assert all(item.value["author"] == "John Doe" for item in john_doe_result)
+
+            draft_result = store.search(["test_search"], filter={"tags": ["draft"]})
+            assert len(draft_result) == 2
+            assert all("draft" in item.value["tags"] for item in draft_result)
+
+            page1 = store.search(["test_search"], limit=2, offset=0)
+            page2 = store.search(["test_search"], limit=2, offset=2)
+            all_items = page1 + page2
+            assert len(all_items) == 4
+            assert len(set(item.key for item in all_items)) == 4
+
+            for namespace in test_namespaces:
+                store.delete(namespace, f"item_{namespace[-1]}")

libs/checkpoint-postgres/tests/test_sync.py

@@ -105,8 +105,8 @@ class TestPostgresSaver:
     def test_null_chars(self) -> None:
         with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
             config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
-            assert saver.get_tuple(config).metadata["my_key"] == "abc"
+            assert saver.get_tuple(config).metadata["my_key"] == "abc"  # type: ignore
             assert (
-                list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"]
+                list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"]  # type: ignore
                 == "abc"
             )

libs/langgraph/tests/conftest.py

@@ -15,6 +15,9 @@ from langgraph.checkpoint.postgres import PostgresSaver
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.checkpoint.sqlite import SqliteSaver
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from langgraph.store.base import BaseStore
+from langgraph.store.memory import InMemoryStore
+from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
 from tests.memory_assert import MemorySaverAssertImmutable
 
 DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/"
@@ -222,6 +225,63 @@ async def awith_checkpointer(
         raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}")
 
 
+@asynccontextmanager
+async def _store_postgres_aio():
+    if sys.version_info < (3, 10):
+        pytest.skip("Async Postgres tests require Python 3.10+")
+    database = f"test_{uuid4().hex[:16]}"
+    async with await AsyncConnection.connect(
+        DEFAULT_POSTGRES_URI, autocommit=True
+    ) as conn:
+        await conn.execute(f"CREATE DATABASE {database}")
+    try:
+        async with AsyncPostgresStore.from_conn_string(
+            DEFAULT_POSTGRES_URI + database
+        ) as store:
+            await store.setup()
+            yield store
+    finally:
+        async with await AsyncConnection.connect(
+            DEFAULT_POSTGRES_URI, autocommit=True
+        ) as conn:
+            await conn.execute(f"DROP DATABASE {database}")
+
+
+@pytest.fixture(scope="function")
+def store_postgres():
+    database = f"test_{uuid4().hex[:16]}"
+    # create unique db
+    with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
+        conn.execute(f"CREATE DATABASE {database}")
+    try:
+        # yield store
+        with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store:
+            store.setup()
+            yield store
+    finally:
+        # drop unique db
+        with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
+            conn.execute(f"DROP DATABASE {database}")
+
+
+@pytest.fixture(scope="function")
+def store_in_memory():
+    yield InMemoryStore()
+
+
+@asynccontextmanager
+async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
+    if store_name is None:
+        yield None
+    elif store_name == "in_memory":
+        yield InMemoryStore()
+    elif store_name == "postgres_aio":
+        async with _store_postgres_aio() as store:
+            yield store
+    else:
+        raise NotImplementedError(f"Unknown store {store_name}")
+
+
 ALL_CHECKPOINTERS_SYNC = [
     "memory",
     "sqlite",
@@ -240,3 +300,5 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [
     *ALL_CHECKPOINTERS_ASYNC,
     None,
 ]
+ALL_STORES_SYNC = ["in_memory", "postgres"]
+ALL_STORES_ASYNC = ["in_memory", "postgres_aio"]

libs/langgraph/tests/test_pregel.py

@@ -75,7 +75,11 @@ from langgraph.store.base import BaseStore
 from langgraph.store.memory import InMemoryStore
 from langgraph.types import Interrupt, PregelTask, Send, StreamWriter
 from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence
-from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS
+from tests.conftest import (
+    ALL_CHECKPOINTERS_SYNC,
+    ALL_STORES_SYNC,
+    SHOULD_CHECK_SNAPSHOTS,
+)
 from tests.fake_chat import FakeChatModel
 from tests.fake_tracer import FakeTracer
 from tests.memory_assert import MemorySaverAssertCheckpointMetadata
@@ -11457,8 +11461,12 @@ def test_subgraph_retries():
 
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
-def test_store_injected(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
+@pytest.mark.parametrize("store_name", ALL_STORES_SYNC)
+def test_store_injected(
+    request: pytest.FixtureRequest, checkpointer_name: str, store_name: str
+) -> None:
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+    the_store = request.getfixturevalue(f"store_{store_name}")
 
     class State(TypedDict):
         count: Annotated[int, operator.add]
@@ -11468,7 +11476,6 @@ def test_store_injected(request: pytest.FixtureRequest, checkpointer_name: str)
 
     def node(input: State, config: RunnableConfig, store: BaseStore):
         assert isinstance(store, BaseStore)
-        assert isinstance(store, InMemoryStore)
         store.put(
             ("foo", "bar"),
             doc_id,
@@ -11483,7 +11490,6 @@ def test_store_injected(request: pytest.FixtureRequest, checkpointer_name: str)
     builder = StateGraph(State)
     builder.add_node("node", node)
     builder.add_edge("__start__", "node")
-    the_store = InMemoryStore()
     graph = builder.compile(store=the_store, checkpointer=checkpointer)
 
     thread_1 = str(uuid.uuid4())

libs/langgraph/tests/test_pregel_async.py

@@ -67,8 +67,10 @@ from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSeq
 from tests.conftest import (
     ALL_CHECKPOINTERS_ASYNC,
     ALL_CHECKPOINTERS_ASYNC_PLUS_NONE,
+    ALL_STORES_ASYNC,
     SHOULD_CHECK_SNAPSHOTS,
     awith_checkpointer,
+    awith_store,
 )
 from tests.fake_chat import FakeChatModel
 from tests.fake_tracer import FakeTracer
@@ -9709,7 +9711,8 @@ async def test_checkpointer_null_pending_writes() -> None:
 
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
-async def test_store_injected_async(checkpointer_name: str) -> None:
+@pytest.mark.parametrize("store_name", ALL_STORES_ASYNC)
+async def test_store_injected_async(checkpointer_name: str, store_name: str) -> None:
     class State(TypedDict):
         count: Annotated[int, operator.add]
 
@@ -9718,7 +9721,6 @@ async def test_store_injected_async(checkpointer_name: str) -> None:
 
     async def node(input: State, config: RunnableConfig, store: BaseStore):
         assert isinstance(store, BaseStore)
-        assert isinstance(store, InMemoryStore)
         await store.aput(
             ("foo", "bar"),
             doc_id,
@@ -9733,8 +9735,9 @@ async def test_store_injected_async(checkpointer_name: str) -> None:
     builder = StateGraph(State)
     builder.add_node("node", node)
     builder.add_edge("__start__", "node")
-    the_store = InMemoryStore()
-    async with awith_checkpointer(checkpointer_name) as checkpointer:
+    async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store(
+        store_name
+    ) as the_store:
         graph = builder.compile(store=the_store, checkpointer=checkpointer)
 
         thread_1 = str(uuid.uuid4())