Skip to content

Commit

Permalink
fixup initial provisioning of aio postgres db (#2571)
Browse files Browse the repository at this point in the history
fixes #2570
  • Loading branch information
taigrr authored Dec 3, 2024
1 parent c6fe265 commit 57f18e4
Showing 1 changed file with 8 additions and 11 deletions.
19 changes: 8 additions & 11 deletions libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

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
Expand Down Expand Up @@ -81,17 +80,15 @@ async def setup(self) -> None:
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 :],
Expand Down

0 comments on commit 57f18e4

Please sign in to comment.