From c31c940bbb0092ba16c74bd7a28ba383adb6c0ee Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Tue, 3 Jun 2025 14:59:56 -0700
Subject: [PATCH] checkpoint-postgres: Use lock also for pipeline mode

---
 .../langgraph/checkpoint/postgres/__init__.py            | 6 ++----
 .../langgraph/checkpoint/postgres/aio.py                 | 9 ++-------
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
index 41f06e869..be6f3c137 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
@@ -433,7 +433,7 @@ class PostgresSaver(BasePostgresSaver):
                 Will be applied regardless of whether the PostgresSaver instance was initialized with a pipeline.
                 If pipeline mode is not supported, will fall back to using transaction context manager.
         """
-        with _internal.get_connection(self.conn) as conn:
+        with self.lock, _internal.get_connection(self.conn) as conn:
             if self.pipe:
                 # a connection in pipeline mode can be used concurrently
                 # in multiple threads/coroutines, but only one cursor can be
@@ -449,7 +449,6 @@ class PostgresSaver(BasePostgresSaver):
                 # thread/coroutine at a time, so we acquire a lock
                 if self.supports_pipeline:
                     with (
-                        self.lock,
                         conn.pipeline(),
                         conn.cursor(binary=True, row_factory=dict_row) as cur,
                     ):
@@ -457,13 +456,12 @@ class PostgresSaver(BasePostgresSaver):
                 else:
                     # Use connection's transaction context manager when pipeline mode not supported
                     with (
-                        self.lock,
                         conn.transaction(),
                         conn.cursor(binary=True, row_factory=dict_row) as cur,
                     ):
                         yield cur
             else:
-                with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
+                with conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
 
 
diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
index 5e5660cb6..8253d4d8f 100644
--- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
@@ -391,7 +391,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline.
                 If pipeline mode is not supported, will fall back to using transaction context manager.
         """
-        async with _ainternal.get_connection(self.conn) as conn:
+        async with self.lock, _ainternal.get_connection(self.conn) as conn:
             if self.pipe:
                 # a connection in pipeline mode can be used concurrently
                 # in multiple threads/coroutines, but only one cursor can be
@@ -407,7 +407,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 # thread/coroutine at a time, so we acquire a lock
                 if self.supports_pipeline:
                     async with (
-                        self.lock,
                         conn.pipeline(),
                         conn.cursor(binary=True, row_factory=dict_row) as cur,
                     ):
@@ -415,16 +414,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 else:
                     # Use connection's transaction context manager when pipeline mode not supported
                     async with (
-                        self.lock,
                         conn.transaction(),
                         conn.cursor(binary=True, row_factory=dict_row) as cur,
                     ):
                         yield cur
             else:
-                async with (
-                    self.lock,
-                    conn.cursor(binary=True, row_factory=dict_row) as cur,
-                ):
+                async with conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
 
     def list(