Use a lock for all operations on sqlite checkpointer

- Otherwise when used in multiple subgraphs in parallel separate queries can interfere w each other

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 (?, ?, ?, ?, ?, ?, ?, ?)",
                 [

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 (?, ?, ?, ?, ?, ?, ?, ?)",
                 [

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