From 05dbc1d4982df5d39f5bd044ce8d21f6627ff3b9 Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Sun, 6 Oct 2024 14:09:35 -0700
Subject: [PATCH] Validate in async batched store (#2017)

---
 libs/checkpoint/langgraph/store/base/batch.py | 11 +++++-
 libs/checkpoint/tests/test_store.py           | 38 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py
index 8283a7a66..079888222 100644
--- a/libs/checkpoint/langgraph/store/base/batch.py
+++ b/libs/checkpoint/langgraph/store/base/batch.py
@@ -2,7 +2,15 @@ import asyncio
 import weakref
 from typing import Any, Optional
 
-from langgraph.store.base import BaseStore, GetOp, Item, Op, PutOp, SearchOp
+from langgraph.store.base import (
+    BaseStore,
+    GetOp,
+    Item,
+    Op,
+    PutOp,
+    SearchOp,
+    _validate_namespace,
+)
 
 
 class AsyncBatchedBaseStore(BaseStore):
@@ -46,6 +54,7 @@ class AsyncBatchedBaseStore(BaseStore):
         key: str,
         value: dict[str, Any],
     ) -> None:
+        _validate_namespace(namespace)
         fut = self._loop.create_future()
         self._aqueue[fut] = PutOp(namespace, key, value)
         return await fut
diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py
index e1d3720a7..d9fbc5084 100644
--- a/libs/checkpoint/tests/test_store.py
+++ b/libs/checkpoint/tests/test_store.py
@@ -312,3 +312,41 @@ async def test_cannot_put_empty_namespace() -> None:
     assert store.search(("langgraph", "foo"))[0].value == doc
     store.delete(("langgraph", "foo"), "bar")
     assert store.get(("langgraph", "foo"), "bar") is None
+
+    class MockAsyncBatchedStore(AsyncBatchedBaseStore):
+        def __init__(self):
+            super().__init__()
+            self._store = InMemoryStore()
+
+        def batch(self, ops: Iterable[Op]) -> list[Result]:
+            return self._store.batch(ops)
+
+        async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+            return self._store.batch(ops)
+
+    async_store = MockAsyncBatchedStore()
+    doc = {"foo": "bar"}
+
+    with pytest.raises(InvalidNamespaceError):
+        await async_store.aput((), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await async_store.aput(("the", "thing.about"), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await async_store.aput(("some", "fun", ""), "foo", doc)
+
+    with pytest.raises(InvalidNamespaceError):
+        await async_store.aput(("langgraph", "foo"), "bar", doc)
+
+    await async_store.aput(("foo", "langgraph", "foo"), "bar", doc)
+    assert (await async_store.aget(("foo", "langgraph", "foo"), "bar")).value == doc
+    assert (await async_store.asearch(("foo", "langgraph", "foo")))[0].value == doc
+    await async_store.adelete(("foo", "langgraph", "foo"), "bar")
+    assert (await async_store.aget(("foo", "langgraph", "foo"), "bar")) is None
+
+    await async_store.abatch([PutOp(("valid", "namespace"), "key", doc)])
+    assert (await async_store.aget(("valid", "namespace"), "key")).value == doc
+    assert (await async_store.asearch(("valid", "namespace")))[0].value == doc
+    await async_store.adelete(("valid", "namespace"), "key")
+    assert (await async_store.aget(("valid", "namespace"), "key")) is None