Fixup initial provisioning of aio postgres db (#2571) (#2600)

fixes #2570

---------

Co-authored-by: Tai Groot <tai@taigrr.com>
This commit is contained in:
William FH
2024-12-03 01:55:26 +00:00
committed by GitHub
co-authored by Tai Groot
parent 64b99c187a
commit 4332a9515d
5 changed files with 408 additions and 194 deletions
@@ -5,7 +5,6 @@ from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
@@ -76,16 +75,15 @@ class PostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
with self._cursor() as cur:
try:
row = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
).fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
cur.execute(self.MIGRATIONS[0])
results = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
@@ -5,7 +5,6 @@ from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
@@ -81,17 +80,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
async with self._cursor() as cur:
try:
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = await results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
await cur.execute(self.MIGRATIONS[0])
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = await results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],