From 82257015114be2541bdff395a6e094965bf6c3c2 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Thu, 5 Dec 2024 18:40:17 -0800
Subject: [PATCH] Better handle sync calls

---
 .../langgraph/store/postgres/aio.py           |  7 +-
 .../tests/test_async_store.py                 | 92 +++++++++++++++++++
 2 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py
index e62d360cc..42f7360fa 100644
--- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py
+++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py
@@ -156,7 +156,12 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
         return results
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:
-        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
+        futures = []
+        for op in ops:
+            fut = self._loop.create_future()
+            self._aqueue[fut] = op
+            futures.append(fut)
+        return [fut.result() for fut in asyncio.as_completed(futures)]
 
     @classmethod
     @asynccontextmanager
diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py
index eda0e2820..31af9603c 100644
--- a/libs/checkpoint-postgres/tests/test_async_store.py
+++ b/libs/checkpoint-postgres/tests/test_async_store.py
@@ -1,8 +1,10 @@
 # type: ignore
+import asyncio
 import itertools
 import sys
 import uuid
 from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from typing import Any, Optional
 
@@ -63,6 +65,96 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
             await conn.execute(f"DROP DATABASE {database}")
 
 
+async def test_large_batches(store: AsyncPostgresStore) -> None:
+    N = 1000
+    M = 10
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for m in range(M):
+            for i in range(N):
+                _ = [
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        value={"foo": "bar" + str(i)},
+                    ),
+                    executor.submit(
+                        store.get,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                    ),
+                    executor.submit(
+                        store.list_namespaces,
+                        prefix=None,
+                        max_depth=m + 1,
+                    ),
+                    executor.submit(
+                        store.search,
+                        ("test",),
+                    ),
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        value={"foo": "bar" + str(i)},
+                    ),
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        None,
+                    ),
+                ]
+
+
+async def test_large_batches_async(store: AsyncPostgresStore) -> None:
+    N = 1000
+    M = 10
+    coros = []
+    for m in range(M):
+        for i in range(N):
+            coros.append(
+                store.aput(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                    value={"foo": "bar" + str(i)},
+                )
+            )
+            coros.append(
+                store.aget(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                )
+            )
+            coros.append(
+                store.alist_namespaces(
+                    prefix=None,
+                    max_depth=m + 1,
+                )
+            )
+            coros.append(
+                store.asearch(
+                    ("test",),
+                )
+            )
+            coros.append(
+                store.aput(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                    value={"foo": "bar" + str(i)},
+                )
+            )
+            coros.append(
+                store.adelete(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                )
+            )
+
+    await asyncio.gather(*coros)
+
+
 async def test_abatch_order(store: AsyncPostgresStore) -> None:
     # Setup test data
     await store.aput(("test", "foo"), "key1", {"data": "value1"})