diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index c407f26f1..54d235bbe 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -25,6 +25,7 @@ from langgraph.store.base import ( BaseStore, GetOp, IndexConfig, + InvalidNamespaceError, Item, ListNamespacesOp, Op, @@ -1386,6 +1387,19 @@ 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): + if isinstance(op, (GetOp, PutOp)): + paths = (op.namespace,) + elif isinstance(op, SearchOp): + paths = (op.namespace_prefix,) + else: + paths = tuple(condition.path for condition in op.match_conditions or ()) + for path in paths: + for label in path: + if not isinstance(label, str) or not label or "." in label: + raise InvalidNamespaceError( + f"Invalid namespace label {label!r}: labels must be nonempty " + "strings without periods ('.')." + ) grouped_ops[type(op)].append((idx, op)) tot += 1 return grouped_ops, tot diff --git a/libs/checkpoint-postgres/tests/test_namespace_validation.py b/libs/checkpoint-postgres/tests/test_namespace_validation.py new file mode 100644 index 000000000..1d41ed3bc --- /dev/null +++ b/libs/checkpoint-postgres/tests/test_namespace_validation.py @@ -0,0 +1,82 @@ +import pytest +from langgraph.store.base import ( + GetOp, + InvalidNamespaceError, + ListNamespacesOp, + MatchCondition, + PutOp, + SearchOp, +) + +from langgraph.store.postgres import PostgresStore +from langgraph.store.postgres.aio import AsyncPostgresStore +from tests.conftest import DEFAULT_URI + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +@pytest.mark.parametrize("kind", ["get", "put", "delete", "search", "prefix", "suffix"]) +def test_invalid_namespace_batch(namespace, kind): + operations = { + "get": GetOp(namespace, "key"), + "put": PutOp(namespace, "key", {"changed": True}), + "delete": PutOp(namespace, "key", None), + "search": SearchOp(namespace), + "prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)), + "suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)), + } + with PostgresStore.from_conn_string(DEFAULT_URI) as store: + store.setup() + store.put(("foo", "bar"), "key", {"original": True}) + with pytest.raises(InvalidNamespaceError): + store.batch([PutOp(("valid",), "key", {}), operations[kind]]) + item = store.get(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + assert store.get(("valid",), "key") is None + + +@pytest.mark.parametrize("method", ["get", "delete", "search"]) +def test_invalid_namespace_public(method): + with PostgresStore.from_conn_string(DEFAULT_URI) as store: + store.setup() + store.put(("foo", "bar"), "key", {"original": True}) + args = (("foo.bar",),) if method == "search" else (("foo.bar",), "key") + with pytest.raises(InvalidNamespaceError): + getattr(store, method)(*args) + item = store.get(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + + +async def test_async_namespace_validation(): + async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store: + await store.setup() + await store.aput(("foo", "bar"), "key", {"original": True}) + for op in ( + GetOp(("foo.bar",), "key"), + PutOp(("foo.bar",), "key", {}), + PutOp(("foo.bar",), "key", None), + SearchOp(("foo.bar",)), + ListNamespacesOp((MatchCondition("prefix", ("foo.bar",)),)), + ): + with pytest.raises(InvalidNamespaceError): + await store.abatch([op]) + with pytest.raises(InvalidNamespaceError): + await store.aget(("foo.bar",), "key") + with pytest.raises(InvalidNamespaceError): + await store.adelete(("foo.bar",), "key") + with pytest.raises(InvalidNamespaceError): + await store.asearch(("foo.bar",)) + item = await store.aget(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + + +def test_hierarchical_search_and_listing_preserved(): + with PostgresStore.from_conn_string(DEFAULT_URI) as store: + store.setup() + namespaces = [("user_%",), ("user_%", "child"), ("userX%",), ("USER_%",)] + for namespace in namespaces: + store.put(namespace, "key", {}) + assert {item.namespace for item in store.search(())} == set(namespaces) + assert {item.namespace for item in store.search(("user_%",))} == set( + namespaces[:2] + ) + assert store.list_namespaces(prefix=("user_%", "*")) == [("user_%", "child")] diff --git a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py index fbc23beda..5df32b641 100644 --- a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py +++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py @@ -17,6 +17,7 @@ from langgraph.store.base import ( BaseStore, GetOp, IndexConfig, + InvalidNamespaceError, Item, ListNamespacesOp, Op, @@ -257,6 +258,19 @@ 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): + if isinstance(op, (GetOp, PutOp)): + paths = (op.namespace,) + elif isinstance(op, SearchOp): + paths = (op.namespace_prefix,) + else: + paths = tuple(condition.path for condition in op.match_conditions or ()) + for path in paths: + for label in path: + if not isinstance(label, str) or not label or "." in label: + raise InvalidNamespaceError( + f"Invalid namespace label {label!r}: labels must be nonempty " + "strings without periods ('.')." + ) grouped_ops[type(op)].append((idx, op)) tot += 1 return grouped_ops, tot diff --git a/libs/checkpoint-sqlite/tests/test_namespace_validation.py b/libs/checkpoint-sqlite/tests/test_namespace_validation.py new file mode 100644 index 000000000..05fbeee8a --- /dev/null +++ b/libs/checkpoint-sqlite/tests/test_namespace_validation.py @@ -0,0 +1,77 @@ +import pytest +from langgraph.store.base import ( + GetOp, + InvalidNamespaceError, + ListNamespacesOp, + MatchCondition, + PutOp, + SearchOp, +) + +from langgraph.store.sqlite import SqliteStore +from langgraph.store.sqlite.aio import AsyncSqliteStore + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +@pytest.mark.parametrize("kind", ["get", "put", "delete", "search", "prefix", "suffix"]) +def test_invalid_namespace_batch(namespace, kind): + operations = { + "get": GetOp(namespace, "key"), + "put": PutOp(namespace, "key", {"changed": True}), + "delete": PutOp(namespace, "key", None), + "search": SearchOp(namespace), + "prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)), + "suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)), + } + with SqliteStore.from_conn_string(":memory:") as store: + store.put(("foo", "bar"), "key", {"original": True}) + with pytest.raises(InvalidNamespaceError): + store.batch([PutOp(("valid",), "key", {}), operations[kind]]) + item = store.get(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + assert store.get(("valid",), "key") is None + + +@pytest.mark.parametrize("method", ["get", "delete", "search"]) +def test_invalid_namespace_public(method): + with SqliteStore.from_conn_string(":memory:") as store: + store.put(("foo", "bar"), "key", {"original": True}) + args = (("foo.bar",),) if method == "search" else (("foo.bar",), "key") + with pytest.raises(InvalidNamespaceError): + getattr(store, method)(*args) + item = store.get(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + + +async def test_async_namespace_validation(): + async with AsyncSqliteStore.from_conn_string(":memory:") as store: + await store.aput(("foo", "bar"), "key", {"original": True}) + for op in ( + GetOp(("foo.bar",), "key"), + PutOp(("foo.bar",), "key", {}), + PutOp(("foo.bar",), "key", None), + SearchOp(("foo.bar",)), + ListNamespacesOp((MatchCondition("prefix", ("foo.bar",)),)), + ): + with pytest.raises(InvalidNamespaceError): + await store.abatch([op]) + with pytest.raises(InvalidNamespaceError): + await store.aget(("foo.bar",), "key") + with pytest.raises(InvalidNamespaceError): + await store.adelete(("foo.bar",), "key") + with pytest.raises(InvalidNamespaceError): + await store.asearch(("foo.bar",)) + item = await store.aget(("foo", "bar"), "key") + assert item is not None and item.value == {"original": True} + + +def test_hierarchical_search_and_listing_preserved(): + with SqliteStore.from_conn_string(":memory:") as store: + namespaces = [("user_%",), ("user_%", "child"), ("userX%",), ("USER_%",)] + for namespace in namespaces: + store.put(namespace, "key", {}) + assert {item.namespace for item in store.search(())} == set(namespaces) + assert {item.namespace for item in store.search(("user_%",))} == set( + namespaces[:2] + ) + assert store.list_namespaces(prefix=("user_%", "*")) == [("user_%", "child")]