Move namespace validation to shared store methods

Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>

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

libs/checkpoint-postgres/tests/test_namespace_validation.py

@@ -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")]

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

libs/checkpoint-sqlite/tests/test_namespace_validation.py

@@ -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")]

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(

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 = []

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()