fix(store): reject ambiguous namespace labels before SQL batches

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

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

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