From 7184d7e3faf33ce526fc0c8ee68c31c550533577 Mon Sep 17 00:00:00 2001 From: Baihao Wang Date: Tue, 6 Oct 2026 18:16:20 +0000 Subject: [PATCH] Move namespace validation to shared store methods Co-authored-by: open-swe[bot] --- .../langgraph/store/postgres/base.py | 14 ---- .../tests/test_namespace_validation.py | 82 ------------------- .../langgraph/store/sqlite/base.py | 14 ---- .../tests/test_namespace_validation.py | 77 ----------------- .../langgraph/store/base/__init__.py | 22 ++++- libs/checkpoint/langgraph/store/base/batch.py | 6 ++ .../tests/test_store_namespace_validation.py | 76 +++++++++++++++++ 7 files changed, 100 insertions(+), 191 deletions(-) delete mode 100644 libs/checkpoint-postgres/tests/test_namespace_validation.py delete mode 100644 libs/checkpoint-sqlite/tests/test_namespace_validation.py create mode 100644 libs/checkpoint/tests/test_store_namespace_validation.py diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 54d235bbe..c407f26f1 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -25,7 +25,6 @@ from langgraph.store.base import ( BaseStore, GetOp, IndexConfig, - InvalidNamespaceError, Item, ListNamespacesOp, Op, @@ -1387,19 +1386,6 @@ 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 deleted file mode 100644 index 1d41ed3bc..000000000 --- a/libs/checkpoint-postgres/tests/test_namespace_validation.py +++ /dev/null @@ -1,82 +0,0 @@ -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 5df32b641..fbc23beda 100644 --- a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py +++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py @@ -17,7 +17,6 @@ from langgraph.store.base import ( BaseStore, GetOp, IndexConfig, - InvalidNamespaceError, Item, ListNamespacesOp, Op, @@ -258,19 +257,6 @@ 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 deleted file mode 100644 index 05fbeee8a..000000000 --- a/libs/checkpoint-sqlite/tests/test_namespace_validation.py +++ /dev/null @@ -1,77 +0,0 @@ -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")] diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index 727da8987..fdfdf7299 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -772,6 +772,7 @@ class BaseStore(ABC): Returns: The retrieved item or `None` if not found. """ + _validate_namespace_labels(namespace) return self.batch( [GetOp(namespace, str(key), _ensure_refresh(self.ttl_config, refresh_ttl))] )[0] @@ -840,6 +841,7 @@ class BaseStore(ABC): Natural language search support depends on your store implementation and requires proper embedding configuration. """ + _validate_namespace_labels(namespace_prefix) return self.batch( [ SearchOp( @@ -941,6 +943,7 @@ class BaseStore(ABC): namespace: Hierarchical path for the item. key: Unique identifier within the namespace. """ + _validate_namespace_labels(namespace) self.batch([PutOp(namespace, str(key), None, ttl=None)]) def list_namespaces( @@ -984,6 +987,8 @@ class BaseStore(ABC): # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] ``` """ + _validate_namespace_labels(prefix or ()) + _validate_namespace_labels(suffix or ()) match_conditions = [] if prefix: match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) @@ -1014,6 +1019,7 @@ class BaseStore(ABC): Returns: The retrieved item or `None` if not found. """ + _validate_namespace_labels(namespace) return ( await self.abatch( [ @@ -1091,6 +1097,7 @@ class BaseStore(ABC): Natural language search support depends on your store implementation and requires proper embedding configuration. """ + _validate_namespace_labels(namespace_prefix) return ( await self.abatch( [ @@ -1202,6 +1209,7 @@ class BaseStore(ABC): namespace: Hierarchical path for the item. key: Unique identifier within the namespace. """ + _validate_namespace_labels(namespace) await self.abatch([PutOp(namespace, str(key), None)]) async def alist_namespaces( @@ -1245,6 +1253,8 @@ class BaseStore(ABC): # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] ``` """ + _validate_namespace_labels(prefix or ()) + _validate_namespace_labels(suffix or ()) match_conditions = [] if prefix: match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) @@ -1263,6 +1273,14 @@ class BaseStore(ABC): def _validate_namespace(namespace: tuple[str, ...]) -> None: if not namespace: raise InvalidNamespaceError("Namespace cannot be empty.") + _validate_namespace_labels(namespace) + if namespace[0] == "langgraph": + raise InvalidNamespaceError( + f'Root label for namespace cannot be "langgraph". Got: {namespace}' + ) + + +def _validate_namespace_labels(namespace: tuple[str, ...]) -> None: for label in namespace: if not isinstance(label, str): raise InvalidNamespaceError( @@ -1277,10 +1295,6 @@ def _validate_namespace(namespace: tuple[str, ...]) -> None: raise InvalidNamespaceError( f"Namespace labels cannot be empty strings. Got {label} in {namespace}" ) - if namespace[0] == "langgraph": - raise InvalidNamespaceError( - f'Root label for namespace cannot be "langgraph". Got: {namespace}' - ) def _ensure_refresh( diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 9deb02e2a..8b591df76 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -25,6 +25,7 @@ from langgraph.store.base import ( _ensure_refresh, _ensure_ttl, _validate_namespace, + _validate_namespace_labels, ) F = TypeVar("F", bound=Callable) @@ -86,6 +87,7 @@ class AsyncBatchedBaseStore(BaseStore): *, refresh_ttl: bool | None = None, ) -> Item | None: + _validate_namespace_labels(namespace) self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait( @@ -111,6 +113,7 @@ class AsyncBatchedBaseStore(BaseStore): offset: int = 0, refresh_ttl: bool | None = None, ) -> list[SearchItem]: + _validate_namespace_labels(namespace_prefix) self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait( @@ -155,6 +158,7 @@ class AsyncBatchedBaseStore(BaseStore): namespace: tuple[str, ...], key: str, ) -> None: + _validate_namespace_labels(namespace) self._ensure_task() fut = self._loop.create_future() self._aqueue.put_nowait((fut, PutOp(namespace, key, None))) @@ -169,6 +173,8 @@ class AsyncBatchedBaseStore(BaseStore): limit: int = 100, offset: int = 0, ) -> list[tuple[str, ...]]: + _validate_namespace_labels(prefix or ()) + _validate_namespace_labels(suffix or ()) self._ensure_task() fut = self._loop.create_future() match_conditions = [] diff --git a/libs/checkpoint/tests/test_store_namespace_validation.py b/libs/checkpoint/tests/test_store_namespace_validation.py new file mode 100644 index 000000000..2ab0c86c8 --- /dev/null +++ b/libs/checkpoint/tests/test_store_namespace_validation.py @@ -0,0 +1,76 @@ +import pytest + +from langgraph.store.base import BaseStore, InvalidNamespaceError +from langgraph.store.base.batch import AsyncBatchedBaseStore + + +class RecordingStore(BaseStore): + def __init__(self): + self.operations = [] + + def batch(self, ops): + self.operations.extend(ops) + return [None for _ in self.operations] + + async def abatch(self, ops): + return self.batch(ops) + + +class RecordingBatchedStore(AsyncBatchedBaseStore): + async def abatch(self, ops): + return [None for _ in ops] + + +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), (123,)]) +@pytest.mark.parametrize("method", ["get", "delete", "search", "list_namespaces"]) +def test_base_rejects_invalid_labels(namespace, method): + store = RecordingStore() + args = (namespace, "key") if method in ("get", "delete") else (namespace,) + with pytest.raises(InvalidNamespaceError): + if method == "list_namespaces": + store.list_namespaces(prefix=namespace) + else: + getattr(store, method)(*args) + assert store.operations == [] + + +@pytest.mark.parametrize("store_type", [RecordingStore, RecordingBatchedStore]) +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), (123,)]) +@pytest.mark.parametrize("method", ["aget", "adelete", "asearch", "alist_namespaces"]) +async def test_async_rejects_invalid_labels(store_type, namespace, method): + store = store_type() + args = (namespace, "key") if method in ("aget", "adelete") else (namespace,) + try: + with pytest.raises(InvalidNamespaceError): + if method == "alist_namespaces": + await store.alist_namespaces(suffix=namespace) + else: + await getattr(store, method)(*args) + if isinstance(store, RecordingBatchedStore): + assert store._aqueue.empty() + else: + assert store.operations == [] + finally: + if isinstance(store, RecordingBatchedStore) and store._task: + store._task.cancel() + + +def test_search_and_listing_preserve_prefixes(): + store = RecordingStore() + store.search(()) + store.search(("tenant", "a_%")) + store.list_namespaces(prefix=("tenant", "*"), suffix=("*",)) + assert store.operations[0].namespace_prefix == () + assert store.operations[1].namespace_prefix == ("tenant", "a_%") + assert store.operations[2].match_conditions[0].path == ("tenant", "*") + + +@pytest.mark.parametrize("store_type", [RecordingStore, RecordingBatchedStore]) +async def test_async_empty_search_and_listing_wildcards(store_type): + store = store_type() + try: + await store.asearch(()) + await store.alist_namespaces(prefix=("tenant", "*")) + finally: + if isinstance(store, RecordingBatchedStore) and store._task: + store._task.cancel()