Merge branch 'main' into vb/update-get-state

This commit is contained in:
vbarda
2024-08-21 15:07:52 -04:00
64 changed files with 4372 additions and 454 deletions
@@ -316,7 +316,17 @@ class PostgresSaver(BasePostgresSaver):
writes (List[Tuple[str, Any]]): List of writes to store.
task_id (str): Identifier for the task creating the writes.
"""
with self._cursor() as cur:
with self._cursor(pipeline=True) as cur:
cur.execute(
self.DELETE_WRITES_SQL,
(
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
len(writes),
),
)
cur.executemany(
self.UPSERT_CHECKPOINT_WRITES_SQL,
self._dump_writes(
@@ -272,7 +272,17 @@ class AsyncPostgresSaver(BasePostgresSaver):
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.
"""
async with self._cursor() as cur:
async with self._cursor(pipeline=True) as cur:
await cur.execute(
self.DELETE_WRITES_SQL,
(
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
len(writes),
),
)
await cur.executemany(
self.UPSERT_CHECKPOINT_WRITES_SQL,
await asyncio.to_thread(
@@ -105,6 +105,15 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
"""
DELETE_WRITES_SQL = """
DELETE FROM checkpoint_writes
WHERE thread_id = %s
AND checkpoint_ns = %s
AND checkpoint_id = %s
AND task_id = %s
AND idx >= %s
"""
class BasePostgresSaver(BaseCheckpointSaver):
SELECT_SQL = SELECT_SQL
@@ -112,6 +121,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
DELETE_WRITES_SQL = DELETE_WRITES_SQL
jsonplus_serde = JsonPlusSerializer()
def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint: