From d03c26dd76ae459e41781d2a3f5733c4fe11437e Mon Sep 17 00:00:00 2001 From: byhow <25713361+byhow@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:48:33 -0700 Subject: [PATCH] refactor(checkpoint,checkpoint-postgres,checkpoint-sqlite): share op namespace validation Add `validate_op_namespace` to `langgraph.store.base` and call it from the SQLite and Postgres `_group_ops` instead of keeping a copy in each store, so every path raises the same errors. The store packages now import it, so their next release must require the `langgraph-checkpoint` release that adds it. Also document `InvalidNamespaceError` on `put`/`aput`, rename the SQL "ambiguous" tests to "invalid", cover every invalid label, suffix listing and `refresh_ttl` in the async direct-batch tests, and use one dispatch style in the base tests. --- .../langgraph/store/postgres/base.py | 28 +------- .../tests/test_async_store.py | 19 ++--- libs/checkpoint-postgres/tests/test_store.py | 2 +- .../langgraph/store/sqlite/base.py | 28 +------- .../tests/test_async_store.py | 19 ++--- libs/checkpoint-sqlite/tests/test_store.py | 2 +- .../langgraph/store/base/__init__.py | 31 +++++++++ libs/checkpoint/tests/test_store.py | 69 +++++++++++++++---- 8 files changed, 116 insertions(+), 82 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 9f7906ab8..b36671b92 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, @@ -37,6 +36,7 @@ from langgraph.store.base import ( ensure_embeddings, get_text_at_path, tokenize_path, + validate_op_namespace, ) from psycopg import Capabilities, Connection, Cursor, Pipeline from psycopg.rows import DictRow, dict_row @@ -1387,36 +1387,12 @@ 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): - _validate_op_namespace(op) + validate_op_namespace(op) grouped_ops[type(op)].append((idx, op)) tot += 1 return grouped_ops, tot -def _validate_op_namespace(op: Op) -> None: - """Reject namespace labels that would alias another namespace once dot-joined. - - `("foo.bar",)` and `("foo", "bar")` flatten to the same `prefix` text. - `BaseStore` methods check labels before batching, but ops passed directly to - `batch`/`abatch` reach here unchecked. - """ - if isinstance(op, (GetOp, PutOp)): - paths = (op.namespace,) - elif isinstance(op, SearchOp): - paths = (op.namespace_prefix,) - elif isinstance(op, ListNamespacesOp): - paths = tuple(condition.path for condition in op.match_conditions or ()) - else: - return - 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} found in {path}. Namespace " - "labels must be non-empty strings without periods ('.')." - ) - - def _json_loads(content: bytes | orjson.Fragment) -> Any: if isinstance(content, orjson.Fragment): if hasattr(content, "buf"): diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index 998f85574..0f0fd5ce6 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -875,17 +875,20 @@ async def test_omit_expired_search_pagination(store: AsyncPostgresStore) -> None assert [i.key for i in page2] == ["c"] -async def test_abatch_rejects_ambiguous_namespace_labels( - store: AsyncPostgresStore, +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +async def test_abatch_rejects_invalid_namespace_labels( + store: AsyncPostgresStore, namespace: tuple ) -> None: await store.aput(("foo", "bar"), "key", {"original": True}) for op in ( - GetOp(("foo.bar",), "key"), - PutOp(("foo.bar",), "key", {"changed": True}), - PutOp(("foo.bar",), "key", None), - SearchOp(("foo.bar",)), - ListNamespacesOp((MatchCondition("prefix", ("foo.bar",)),)), + GetOp(namespace, "key"), + GetOp(namespace, "key", refresh_ttl=True), + PutOp(namespace, "key", {"changed": True}), + PutOp(namespace, "key", None), + SearchOp(namespace), + ListNamespacesOp((MatchCondition("prefix", namespace),)), + ListNamespacesOp((MatchCondition("suffix", namespace),)), ): with pytest.raises(InvalidNamespaceError): await store.abatch([op]) @@ -914,7 +917,7 @@ async def test_invalid_namespace_only_fails_its_own_call( assert isinstance(invalid, InvalidNamespaceError) -async def test_sync_methods_reject_ambiguous_namespace_labels( +async def test_sync_methods_reject_invalid_namespace_labels( store: AsyncPostgresStore, ) -> None: """The sync wrappers run off the event loop thread and must validate too.""" diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index 495b4395d..abe6475e6 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -1171,7 +1171,7 @@ def test_namespace_labels_with_trailing_newline(store) -> None: @pytest.mark.parametrize( "kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"] ) -def test_batch_rejects_ambiguous_namespace_labels( +def test_batch_rejects_invalid_namespace_labels( store, namespace: tuple, kind: str ) -> None: """Ops passed straight to `batch` must not reach another namespace. diff --git a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py index 5e901c49b..c4f8f6d18 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, @@ -29,6 +28,7 @@ from langgraph.store.base import ( ensure_embeddings, get_text_at_path, tokenize_path, + validate_op_namespace, ) _AIO_ERROR_MSG = ( @@ -258,36 +258,12 @@ 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): - _validate_op_namespace(op) + validate_op_namespace(op) grouped_ops[type(op)].append((idx, op)) tot += 1 return grouped_ops, tot -def _validate_op_namespace(op: Op) -> None: - """Reject namespace labels that would alias another namespace once dot-joined. - - `("foo.bar",)` and `("foo", "bar")` flatten to the same `prefix` text. - `BaseStore` methods check labels before batching, but ops passed directly to - `batch`/`abatch` reach here unchecked. - """ - if isinstance(op, (GetOp, PutOp)): - paths = (op.namespace,) - elif isinstance(op, SearchOp): - paths = (op.namespace_prefix,) - elif isinstance(op, ListNamespacesOp): - paths = tuple(condition.path for condition in op.match_conditions or ()) - else: - return - 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} found in {path}. Namespace " - "labels must be non-empty strings without periods ('.')." - ) - - class PreparedGetQuery(NamedTuple): query: str # Main query to execute params: tuple # Parameters for the main query diff --git a/libs/checkpoint-sqlite/tests/test_async_store.py b/libs/checkpoint-sqlite/tests/test_async_store.py index 767ceabf0..4af125c9f 100644 --- a/libs/checkpoint-sqlite/tests/test_async_store.py +++ b/libs/checkpoint-sqlite/tests/test_async_store.py @@ -749,17 +749,20 @@ async def test_async_namespace_segment_boundary(store: AsyncSqliteStore) -> None } -async def test_abatch_rejects_ambiguous_namespace_labels( - store: AsyncSqliteStore, +@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)]) +async def test_abatch_rejects_invalid_namespace_labels( + store: AsyncSqliteStore, namespace: tuple ) -> None: await store.aput(("foo", "bar"), "key", {"original": True}) for op in ( - GetOp(("foo.bar",), "key"), - PutOp(("foo.bar",), "key", {"changed": True}), - PutOp(("foo.bar",), "key", None), - SearchOp(("foo.bar",)), - ListNamespacesOp((MatchCondition("prefix", ("foo.bar",)),)), + GetOp(namespace, "key"), + GetOp(namespace, "key", refresh_ttl=True), + PutOp(namespace, "key", {"changed": True}), + PutOp(namespace, "key", None), + SearchOp(namespace), + ListNamespacesOp((MatchCondition("prefix", namespace),)), + ListNamespacesOp((MatchCondition("suffix", namespace),)), ): with pytest.raises(InvalidNamespaceError): await store.abatch([op]) @@ -788,7 +791,7 @@ async def test_invalid_namespace_only_fails_its_own_call( assert isinstance(invalid, InvalidNamespaceError) -async def test_sync_methods_reject_ambiguous_namespace_labels( +async def test_sync_methods_reject_invalid_namespace_labels( store: AsyncSqliteStore, ) -> None: """The sync wrappers run off the event loop thread and must validate too.""" diff --git a/libs/checkpoint-sqlite/tests/test_store.py b/libs/checkpoint-sqlite/tests/test_store.py index 0d9b16b12..91835b512 100644 --- a/libs/checkpoint-sqlite/tests/test_store.py +++ b/libs/checkpoint-sqlite/tests/test_store.py @@ -1442,7 +1442,7 @@ def test_list_namespaces_metacharacter_labels(store: SqliteStore) -> None: @pytest.mark.parametrize( "kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"] ) -def test_batch_rejects_ambiguous_namespace_labels( +def test_batch_rejects_invalid_namespace_labels( store: SqliteStore, namespace: tuple, kind: str ) -> None: """Ops passed straight to `batch` must not reach another namespace. diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index e6996675c..ff44e8465 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -897,6 +897,11 @@ class BaseStore(ABC): By default, the expiration timer refreshes on both read operations (get/search) and write operations (put/update), whenever the item is included in the operation. + Raises: + InvalidNamespaceError: If the namespace is empty, its root label is + `"langgraph"`, or a label is empty, is not a string, or contains a + period (`.`). + Note: Indexing support depends on your store implementation. If you do not initialize the store with indexing capabilities, @@ -1171,6 +1176,11 @@ class BaseStore(ABC): By default, the expiration timer refreshes on both read operations (get/search) and write operations (put/update), whenever the item is included in the operation. + Raises: + InvalidNamespaceError: If the namespace is empty, its root label is + `"langgraph"`, or a label is empty, is not a string, or contains a + period (`.`). + Note: Indexing support depends on your store implementation. If you do not initialize the store with indexing capabilities, @@ -1329,6 +1339,27 @@ def _validate_namespace_labels(namespace: tuple[str, ...]) -> None: ) +def validate_op_namespace(op: Op) -> None: + """Validate the namespace labels an op carries before a store executes it. + + `BaseStore` methods check labels before batching, but ops passed directly to + `batch`/`abatch` skip those methods. Stores that serialize namespaces as + delimited text should call this for every op they execute, so a label such + as `"foo.bar"` cannot address the namespace `("foo", "bar")`. + + Raises: + InvalidNamespaceError: If a label is empty, is not a string, or contains + a period (`.`). + """ + if isinstance(op, (GetOp, PutOp)): + _validate_namespace_labels(op.namespace) + elif isinstance(op, SearchOp): + _validate_namespace_labels(op.namespace_prefix) + elif isinstance(op, ListNamespacesOp): + for condition in op.match_conditions or (): + _validate_namespace_labels(condition.path) + + def _ensure_refresh( ttl_config: TTLConfig | None, refresh_ttl: bool | None = None ) -> bool: diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index 36b33e2bb..bee5c68d2 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -13,10 +13,14 @@ from langgraph.store.base import ( GetOp, InvalidNamespaceError, Item, + ListNamespacesOp, + MatchCondition, Op, PutOp, Result, + SearchOp, get_text_at_path, + validate_op_namespace, ) from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.memory import InMemoryStore @@ -539,14 +543,16 @@ def test_rejects_invalid_namespace_labels( ) -> None: store = InMemoryStore() batch = mocker.spy(InMemoryStore, "batch") + call = { + "get": lambda: store.get(namespace, "key"), + "delete": lambda: store.delete(namespace, "key"), + "search": lambda: store.search(namespace), + "prefix": lambda: store.list_namespaces(prefix=namespace), + "suffix": lambda: store.list_namespaces(suffix=namespace), + }[method] with pytest.raises(InvalidNamespaceError): - if method in ("prefix", "suffix"): - store.list_namespaces(**{method: namespace}) - elif method == "search": - store.search(namespace) - else: - getattr(store, method)(namespace, "key") + call() batch.assert_not_called() @@ -563,14 +569,16 @@ async def test_async_rejects_invalid_namespace_labels( # `MockAsyncBatchedStore` dispatches through `InMemoryStore.batch`. batch = mocker.spy(InMemoryStore, "batch") abatch = mocker.spy(InMemoryStore, "abatch") + call = { + "get": lambda: store.aget(namespace, "key"), + "delete": lambda: store.adelete(namespace, "key"), + "search": lambda: store.asearch(namespace), + "prefix": lambda: store.alist_namespaces(prefix=namespace), + "suffix": lambda: store.alist_namespaces(suffix=namespace), + }[method] with pytest.raises(InvalidNamespaceError): - if method in ("prefix", "suffix"): - await store.alist_namespaces(**{method: namespace}) - elif method == "search": - await store.asearch(namespace) - else: - await getattr(store, f"a{method}")(namespace, "key") + await call() batch.assert_not_called() abatch.assert_not_called() @@ -608,6 +616,43 @@ async def test_async_search_and_listing_keep_empty_prefixes_and_wildcards( ) == [("tenant", "a_%"), ("tenant", "b", "child")] +@pytest.mark.parametrize("namespace", INVALID_NAMESPACES) +@pytest.mark.parametrize( + "kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"] +) +def test_validate_op_namespace_rejects_invalid_labels( + namespace: tuple, kind: str +) -> None: + op = { + "get": GetOp(namespace, "key"), + "put": PutOp(namespace, "key", {"v": 1}), + "delete": PutOp(namespace, "key", None), + "search": SearchOp(namespace), + "list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)), + "list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)), + }[kind] + + with pytest.raises(InvalidNamespaceError): + validate_op_namespace(op) + + +def test_validate_op_namespace_allows_empty_prefix_and_wildcards() -> None: + for op in ( + SearchOp(()), + ListNamespacesOp(), + ListNamespacesOp( + ( + MatchCondition("prefix", ("tenant", "*")), + MatchCondition("suffix", ("*",)), + ) + ), + GetOp(("tenant", "a_%"), "key"), + # Write-only rules belong to `put`, not to op validation. + PutOp(("langgraph", "x"), "key", {"v": 1}), + ): + validate_op_namespace(op) + + async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None: abatch = mocker.spy(InMemoryStore, "batch") store = MockAsyncBatchedStore()