diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 97f0f0570..5649e4e8c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -97,7 +97,6 @@ class PostgresSaver(BasePostgresSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: - self.setup() where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" if limit: @@ -129,7 +128,6 @@ class PostgresSaver(BasePostgresSaver): ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - self.setup() thread_id = config["configurable"]["thread_id"] checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") @@ -240,7 +238,6 @@ class PostgresSaver(BasePostgresSaver): @contextmanager def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]: - self.setup() if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index e7f080002..7b2895b1c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -99,7 +99,6 @@ class AsyncPostgresSaver(BasePostgresSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> AsyncIterator[CheckpointTuple]: - await self.setup() where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" if limit: @@ -133,7 +132,6 @@ class AsyncPostgresSaver(BasePostgresSaver): ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - await self.setup() thread_id = config["configurable"]["thread_id"] checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") @@ -186,7 +184,6 @@ class AsyncPostgresSaver(BasePostgresSaver): metadata: CheckpointMetadata, new_versions: ChannelVersions, ) -> RunnableConfig: - await self.setup() configurable = config["configurable"].copy() thread_id = configurable.pop("thread_id") checkpoint_ns = configurable.pop("checkpoint_ns") @@ -249,7 +246,6 @@ class AsyncPostgresSaver(BasePostgresSaver): @asynccontextmanager async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]: - await self.setup() if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index e392ddd65..e94cf32ae 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -13,7 +13,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver class TestAsyncPostgresSaver: @pytest.fixture(autouse=True) - def setup(self): + async def setup(self): # objects for test setup self.config_1: RunnableConfig = { "configurable": { @@ -55,6 +55,8 @@ class TestAsyncPostgresSaver: "score": None, } self.metadata_3: CheckpointMetadata = {} + async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: + await saver.setup() async def test_asearch(self): async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index 0d642ee62..dfae82907 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -55,6 +55,8 @@ class TestPostgresSaver: "score": None, } self.metadata_3: CheckpointMetadata = {} + with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: + saver.setup() def test_search(self): with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 44b441982..eb8f20ce9 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -6,6 +6,8 @@ from psycopg.errors import UndefinedTable from psycopg.rows import dict_row from pytest_mock import MockerFixture +from langgraph.checkpoint.postgres import PostgresSaver + DEFAULT_POSTGRES_URI = ( "postgres://postgres:postgres@localhost:5442/postgres?sslmode=disable" ) @@ -27,6 +29,12 @@ async def conn(): yield conn +@pytest.fixture(scope="session", autouse=True) +def setup_before_all_tests(): + with PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI) as checkpointer: + checkpointer.setup() + + @pytest.fixture(scope="function", autouse=True) async def clear_test_db(conn): """Delete all tables before each test."""