mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-18 05:57:37 +02:00
fix(checkpoint,checkpoint-postgres,checkpoint-sqlite): reject ambiguous namespace labels in direct batches
`BaseStore` methods validate labels before queueing, but ops passed straight
to `batch`/`abatch` on the SQL stores were flattened unchecked, so
`("foo.bar",)` could read, overwrite, or delete `("foo", "bar")`. Validate
each op in `_group_ops` before any SQL runs. Queued ops are already checked
before they are enqueued, so this only fails a direct caller's own batch.
Also document `InvalidNamespaceError` on the read, delete, and listing
methods, and tighten the base validation tests.
This commit is contained in:
@@ -25,6 +25,7 @@ from langgraph.store.base import (
|
||||
BaseStore,
|
||||
GetOp,
|
||||
IndexConfig,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
Op,
|
||||
@@ -1386,11 +1387,36 @@ 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)
|
||||
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"):
|
||||
|
||||
@@ -13,8 +13,10 @@ import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
PutOp,
|
||||
SearchOp,
|
||||
)
|
||||
@@ -871,3 +873,62 @@ async def test_omit_expired_search_pagination(store: AsyncPostgresStore) -> None
|
||||
page2 = await store.asearch(ns, limit=2, offset=2)
|
||||
assert [i.key for i in page1] == ["a", "b"]
|
||||
assert [i.key for i in page2] == ["c"]
|
||||
|
||||
|
||||
async def test_abatch_rejects_ambiguous_namespace_labels(
|
||||
store: AsyncPostgresStore,
|
||||
) -> 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",)),)),
|
||||
):
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await store.abatch([op])
|
||||
|
||||
item = await store.aget(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
|
||||
|
||||
async def test_invalid_namespace_only_fails_its_own_call(
|
||||
store: AsyncPostgresStore,
|
||||
) -> None:
|
||||
"""Concurrent calls share one `abatch`, which fails every op if it raises.
|
||||
|
||||
Labels are checked before an op is queued, so one caller's bad label cannot
|
||||
fail another caller's request.
|
||||
"""
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
valid, invalid = await asyncio.gather(
|
||||
store.aget(("foo", "bar"), "key"),
|
||||
store.aget(("foo.bar",), "key"),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert isinstance(valid, Item) and valid.value == {"original": True}
|
||||
assert isinstance(invalid, InvalidNamespaceError)
|
||||
|
||||
|
||||
async def test_sync_methods_reject_ambiguous_namespace_labels(
|
||||
store: AsyncPostgresStore,
|
||||
) -> None:
|
||||
"""The sync wrappers run off the event loop thread and must validate too."""
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
for call in (
|
||||
lambda: store.get(("foo.bar",), "key"),
|
||||
lambda: store.search(("foo.bar",)),
|
||||
lambda: store.delete(("foo.bar",), "key"),
|
||||
lambda: store.list_namespaces(prefix=("foo.bar",)),
|
||||
lambda: store.batch([GetOp(("foo.bar",), "key")]),
|
||||
):
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await asyncio.to_thread(call)
|
||||
|
||||
item = await store.aget(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
|
||||
@@ -11,6 +11,7 @@ import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
@@ -1164,3 +1165,51 @@ def test_namespace_labels_with_trailing_newline(store) -> None:
|
||||
assert set(store.list_namespaces(prefix=["users", "alice"], limit=100)) == {
|
||||
("users", "alice"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
|
||||
@pytest.mark.parametrize(
|
||||
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
|
||||
)
|
||||
def test_batch_rejects_ambiguous_namespace_labels(
|
||||
store, namespace: tuple, kind: str
|
||||
) -> None:
|
||||
"""Ops passed straight to `batch` must not reach another namespace.
|
||||
|
||||
Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same
|
||||
text as `("foo", "bar")`. `BaseStore` methods validate labels themselves,
|
||||
but `batch` takes ops as given.
|
||||
"""
|
||||
op = {
|
||||
"get": GetOp(namespace, "key"),
|
||||
"put": PutOp(namespace, "key", {"changed": True}),
|
||||
"delete": PutOp(namespace, "key", None),
|
||||
"search": SearchOp(namespace),
|
||||
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
||||
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
||||
}[kind]
|
||||
store.put(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
store.batch([PutOp(("valid",), "key", {}), op])
|
||||
|
||||
item = store.get(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
# The whole batch is rejected before any SQL runs.
|
||||
assert store.get(("valid",), "key") is None
|
||||
|
||||
|
||||
def test_batch_allows_empty_search_prefix_and_listing_wildcards(
|
||||
store,
|
||||
) -> None:
|
||||
store.put(("foo", "bar"), "key", {"v": 1})
|
||||
|
||||
found, listed = store.batch(
|
||||
[
|
||||
SearchOp(()),
|
||||
ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)),
|
||||
]
|
||||
)
|
||||
|
||||
assert [item.namespace for item in found] == [("foo", "bar")]
|
||||
assert listed == [("foo", "bar")]
|
||||
|
||||
Reference in New Issue
Block a user