From f09b1308370b9673d1c95bd6b3738639bb5a95c3 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 6 Sep 2024 12:56:55 -0700
Subject: [PATCH] Use pool for postgres checkpointer

---
 libs/scheduler-kafka/tests/conftest.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/libs/scheduler-kafka/tests/conftest.py b/libs/scheduler-kafka/tests/conftest.py
index 4296b9851..1d50ecc94 100644
--- a/libs/scheduler-kafka/tests/conftest.py
+++ b/libs/scheduler-kafka/tests/conftest.py
@@ -4,6 +4,7 @@ from uuid import uuid4
 import kafka.admin
 import pytest
 from psycopg import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 from langgraph.scheduler.kafka.types import Topics
@@ -45,9 +46,10 @@ async def checkpointer() -> AsyncIterator[AsyncPostgresSaver]:
         await conn.execute(f"CREATE DATABASE {database}")
     try:
         # yield checkpointer
-        async with AsyncPostgresSaver.from_conn_string(
-            DEFAULT_POSTGRES_URI + database
-        ) as checkpointer:
+        async with AsyncConnectionPool(
+            DEFAULT_POSTGRES_URI + database, max_size=10, kwargs={"autocommit": True}
+        ) as pool:
+            checkpointer = AsyncPostgresSaver(pool)
             await checkpointer.setup()
             yield checkpointer
     finally: