From 97f79fc66b2282f9c49ff9d15a4b9c3e8c811128 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 30 Sep 2024 14:18:58 -0700 Subject: [PATCH] Add Postgres Store Implementation (#1906) --- .../langgraph/checkpoint/postgres/__init__.py | 43 +- .../langgraph/store/postgres/__init__.py | 4 + .../langgraph/store/postgres/aio.py | 190 +++++++ .../langgraph/store/postgres/base.py | 394 +++++++++++++++ libs/checkpoint-postgres/pyproject.toml | 2 +- libs/checkpoint-postgres/tests/conftest.py | 2 + libs/checkpoint-postgres/tests/test_async.py | 2 +- .../tests/test_async_store.py | 465 ++++++++++++++++++ libs/checkpoint-postgres/tests/test_store.py | 458 +++++++++++++++++ libs/checkpoint-postgres/tests/test_sync.py | 4 +- libs/langgraph/tests/conftest.py | 62 +++ libs/langgraph/tests/test_pregel.py | 14 +- libs/langgraph/tests/test_pregel_async.py | 11 +- 13 files changed, 1620 insertions(+), 31 deletions(-) create mode 100644 libs/checkpoint-postgres/langgraph/store/postgres/__init__.py create mode 100644 libs/checkpoint-postgres/langgraph/store/postgres/aio.py create mode 100644 libs/checkpoint-postgres/langgraph/store/postgres/base.py create mode 100644 libs/checkpoint-postgres/tests/test_async_store.py create mode 100644 libs/checkpoint-postgres/tests/test_store.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index a2274cd60..04e3747d2 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/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"] diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/store/postgres/__init__.py new file mode 100644 index 000000000..47b7836d2 --- /dev/null +++ b/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"] diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py new file mode 100644 index 000000000..4e59918e0 --- /dev/null +++ b/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,)) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py new file mode 100644 index 000000000..a22001355 --- /dev/null +++ b/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)) diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index 68cf791f4..b65f545e0 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/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" diff --git a/libs/checkpoint-postgres/tests/conftest.py b/libs/checkpoint-postgres/tests/conftest.py index 36336643f..49061b98c 100644 --- a/libs/checkpoint-postgres/tests/conftest.py +++ b/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 diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index d44376fe5..256bbe8a3 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/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" diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py new file mode 100644 index 000000000..244ad299e --- /dev/null +++ b/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]}") diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py new file mode 100644 index 000000000..47ffce88f --- /dev/null +++ b/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]}") diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index d196d075b..ced755955 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/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" ) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 34b450adb..a58923f14 100644 --- a/libs/langgraph/tests/conftest.py +++ b/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"] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 917057b29..0bd5772d4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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()) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 5ad973ce4..0290bf069 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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())