WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

checkpoint-postgres: pass row_factory in cursor (#1433)

Browse Source
This commit is contained in:
Vadym Barda
2024-08-22 17:39:21 -04:00
committed by GitHub
parent 7074604204
commit 8090ca67c5
3 changed files with 71 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/persistence_postgres.ipynb
+3 -4
View File
@@ -122,7 +122,7 @@
"metadata": {},
"outputs": [],
"source": [
"DB_URI = \"postgresql://postgres:postgres@localhost:5441/postgres?sslmode=disable\""
"DB_URI = \"postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable\""
]
},
{
@@ -137,7 +137,6 @@
"connection_kwargs ={\n",
" \"autocommit\": True,\n",
" \"prepare_threshold\": 0,\n",
" \"row_factory\": dict_row,\n",
"}"
]
},
@@ -551,9 +550,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "langgraph-postgres",
"display_name": "langgraph",
"language": "python",
"name": "langgraph-postgres"
"name": "langgraph"
},
"language_info": {
"codemirror_mode": {
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+33 -29
View File
@@ -66,7 +66,7 @@ class PostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
with self.lock:
with self.conn.cursor(binary=True) as cur:
with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
try:
version = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
@@ -127,31 +127,33 @@ class PostgresSaver(BasePostgresSaver):
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
for value in self.conn.execute(query, args, binary=True):
yield CheckpointTuple(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["checkpoint_id"],
with self._cursor() as cur:
cur.execute(query, args, binary=True)
for value in cur:
yield CheckpointTuple(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["checkpoint_id"],
}
},
{
**self._load_checkpoint(value["checkpoint"]),
"channel_values": self._load_blobs(value["channel_values"]),
},
self._load_metadata(value["metadata"]),
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["parent_checkpoint_id"],
}
}
},
{
**self._load_checkpoint(value["checkpoint"]),
"channel_values": self._load_blobs(value["channel_values"]),
},
self._load_metadata(value["metadata"]),
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None,
self._load_writes(value["pending_writes"]),
)
if value["parent_checkpoint_id"]
else None,
self._load_writes(value["pending_writes"]),
)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
@@ -199,7 +201,7 @@ class PostgresSaver(BasePostgresSaver):
where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
with self._cursor() as cur:
cur = self.conn.execute(
cur.execute(
self.SELECT_SQL + where,
args,
binary=True,
@@ -336,7 +338,7 @@ class PostgresSaver(BasePostgresSaver):
# in multiple threads/coroutines, but only one cursor can be
# used at a time
try:
with self.conn.cursor(binary=True) as cur:
with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
finally:
if pipeline:
@@ -344,8 +346,10 @@ class PostgresSaver(BasePostgresSaver):
elif pipeline:
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
with self.lock, self.conn.pipeline(), self.conn.cursor(binary=True) as cur:
with self.lock, self.conn.pipeline(), self.conn.cursor(
binary=True, row_factory=dict_row
) as cur:
yield cur
else:
with self.lock, self.conn.cursor(binary=True) as cur:
with self.lock, self.conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+35 -31
View File
@@ -64,7 +64,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
async with self.lock:
async with self.conn.cursor(binary=True) as cur:
async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
try:
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
@@ -110,33 +110,35 @@ class AsyncPostgresSaver(BasePostgresSaver):
if limit:
query += f" LIMIT {limit}"
# if we change this to use .stream() we need to make sure to close the cursor
async for value in await self.conn.execute(query, args, binary=True):
yield CheckpointTuple(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["checkpoint_id"],
async with self._cursor() as cur:
await cur.execute(query, args, binary=True)
async for value in cur:
yield CheckpointTuple(
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["checkpoint_id"],
}
},
{
**self._load_checkpoint(value["checkpoint"]),
"channel_values": await asyncio.to_thread(
self._load_blobs, value["channel_values"]
),
},
self._load_metadata(value["metadata"]),
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["parent_checkpoint_id"],
}
}
},
{
**self._load_checkpoint(value["checkpoint"]),
"channel_values": await asyncio.to_thread(
self._load_blobs, value["channel_values"]
),
},
self._load_metadata(value["metadata"]),
{
"configurable": {
"thread_id": value["thread_id"],
"checkpoint_ns": value["checkpoint_ns"],
"checkpoint_id": value["parent_checkpoint_id"],
}
}
if value["parent_checkpoint_id"]
else None,
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
if value["parent_checkpoint_id"]
else None,
await asyncio.to_thread(self._load_writes, value["pending_writes"]),
)
async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database asynchronously.
@@ -163,7 +165,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
async with self._cursor() as cur:
cur = await self.conn.execute(
await cur.execute(
self.SELECT_SQL + where,
args,
binary=True,
@@ -293,7 +295,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
# in multiple threads/coroutines, but only one cursor can be
# used at a time
try:
async with self.conn.cursor(binary=True) as cur:
async with self.conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
finally:
if pipeline:
@@ -302,9 +304,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
async with self.lock, self.conn.pipeline(), self.conn.cursor(
binary=True
binary=True, row_factory=dict_row
) as cur:
yield cur
else:
async with self.lock, self.conn.cursor(binary=True) as cur:
async with self.lock, self.conn.cursor(
binary=True, row_factory=dict_row
) as cur:
yield cur