From 4793b3f5e1c690a29b0773e36a4fe37bc53b2b78 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Mon, 2 Sep 2024 19:40:02 -0700
Subject: [PATCH] Use a lock for all operations on sqlite checkpointer

- Otherwise when used in multiple subgraphs in parallel separate queries can interfere w each other
---
 .../langgraph/checkpoint/sqlite/__init__.py   | 21 ++++++++++---------
 .../langgraph/checkpoint/sqlite/aio.py        | 10 +++++----
 libs/langgraph/tests/test_pregel.py           |  1 -
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
index a5f657824..31bffd19c 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py
@@ -165,14 +165,15 @@ class SqliteSaver(BaseCheckpointSaver):
         Yields:
             sqlite3.Cursor: A cursor for the SQLite database.
         """
-        self.setup()
-        cur = self.conn.cursor()
-        try:
-            yield cur
-        finally:
-            if transaction:
-                self.conn.commit()
-            cur.close()
+        with self.lock:
+            self.setup()
+            cur = self.conn.cursor()
+            try:
+                yield cur
+            finally:
+                if transaction:
+                    self.conn.commit()
+                cur.close()
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database.
@@ -398,7 +399,7 @@ class SqliteSaver(BaseCheckpointSaver):
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
         serialized_metadata = self.jsonplus_serde.dumps(metadata)
-        with self.lock, self.cursor() as cur:
+        with self.cursor() as cur:
             cur.execute(
                 "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
                 (
@@ -434,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver):
             writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
         """
-        with self.lock, self.cursor() as cur:
+        with self.cursor() as cur:
             cur.executemany(
                 "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
index 7cdc7c8fb..28ad76109 100644
--- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
+++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py
@@ -276,7 +276,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         """
         await self.setup()
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
-        async with self.conn.cursor() as cur:
+        async with self.lock, self.conn.cursor() as cur:
             # find the latest checkpoint for the thread_id
             if checkpoint_id := get_checkpoint_id(config):
                 await cur.execute(
@@ -371,7 +371,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur:
+        async with self.lock, self.conn.execute(
+            query, params
+        ) as cur, self.conn.cursor() as wcur:
             async for (
                 thread_id,
                 checkpoint_ns,
@@ -438,7 +440,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
         serialized_metadata = self.jsonplus_serde.dumps(metadata)
-        async with self.conn.execute(
+        async with self.lock, self.conn.execute(
             "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
             (
                 str(config["configurable"]["thread_id"]),
@@ -475,7 +477,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         await self.setup()
-        async with self.conn.cursor() as cur:
+        async with self.lock, self.conn.cursor() as cur:
             await cur.executemany(
                 "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 6109f6ce4..96696fa13 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -9784,7 +9784,6 @@ def test_doubly_nested_graph_state(
     ]
 
 
-@pytest.mark.repeat(10)
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_send_to_nested_graphs(
     request: pytest.FixtureRequest, checkpointer_name: str