Use thread_id for partition key, ignore tasks for stale checkpoints

This commit is contained in:
Nuno Campos
2024-09-10 16:17:22 -07:00
parent 030d6d2def
commit 8f91402642
10 changed files with 107 additions and 33 deletions
@@ -298,18 +298,16 @@ class AsyncPostgresSaver(BasePostgresSaver):
if all(w[0] in WRITES_IDX_MAP for w in writes)
else self.INSERT_CHECKPOINT_WRITES_SQL
)
params = await asyncio.to_thread(
self._dump_writes,
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
writes,
)
async with self._cursor(pipeline=True) as cur:
await cur.executemany(
query,
await asyncio.to_thread(
self._dump_writes,
config["configurable"]["thread_id"],
config["configurable"]["checkpoint_ns"],
config["configurable"]["checkpoint_id"],
task_id,
writes,
),
)
await cur.executemany(query, params)
@asynccontextmanager
async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]: