diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 734ccd7e6..97f0f0570 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -4,6 +4,7 @@ from typing import Any, Iterator, List, Optional from langchain_core.runnables import RunnableConfig from psycopg import Connection, Cursor, Pipeline +from psycopg.errors import UndefinedTable from psycopg.rows import dict_row from psycopg.types.json import Jsonb @@ -70,15 +71,19 @@ class PostgresSaver(BasePostgresSaver): if self.is_setup: return with self.lock: - create_table_queries = [ - self.CREATE_CHECKPOINTS_SQL, - self.CREATE_CHECKPOINT_BLOBS_SQL, - self.CREATE_CHECKPOINT_WRITES_SQL, - ] with self.conn.cursor(binary=True) as cur: - for query in create_table_queries: - cur.execute(query) - + try: + version = cur.execute( + "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" + ).fetchone()["v"] + except UndefinedTable: + version = -1 + for v, migration in zip( + range(version + 1, len(self.MIGRATIONS)), + self.MIGRATIONS[version + 1 :], + ): + cur.execute(migration) + cur.execute(f"INSERT INTO checkpoint_migrations (v) VALUES ({v})") if self.pipe: self.pipe.sync() diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 686f0b305..e7f080002 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -4,6 +4,7 @@ from typing import Any, AsyncIterator, Optional from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline +from psycopg.errors import UndefinedTable from psycopg.rows import dict_row from psycopg.types.json import Jsonb @@ -68,15 +69,23 @@ class AsyncPostgresSaver(BasePostgresSaver): if self.is_setup: return async with self.lock: - create_table_queries = [ - self.CREATE_CHECKPOINTS_SQL, - self.CREATE_CHECKPOINT_BLOBS_SQL, - self.CREATE_CHECKPOINT_WRITES_SQL, - ] - async with self.conn.cursor() as cur: - for query in create_table_queries: - await cur.execute(query) - + async with self.conn.cursor(binary=True) as cur: + try: + version = ( + await cur.execute( + "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1" + ) + ).fetchone()["v"] + except UndefinedTable: + version = -1 + for v, migration in zip( + range(version + 1, len(self.MIGRATIONS)), + self.MIGRATIONS[version + 1 :], + ): + await cur.execute(migration) + await cur.execute( + f"INSERT INTO checkpoint_migrations (v) VALUES ({v})" + ) if self.pipe: await self.pipe.sync() diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 5a88f61ac..0734707c8 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -15,6 +15,46 @@ from langgraph.checkpoint.serde.types import ChannelProtocol MetadataInput = Optional[dict[str, Any]] +""" +To add a new migration, add a new string to the MIGRATIONS list. +The position of the migration in the list is the version number. +""" +MIGRATIONS = [ + """CREATE TABLE IF NOT EXISTS checkpoint_migrations ( + v INTEGER PRIMARY KEY +);""", + """CREATE TABLE IF NOT EXISTS checkpoints ( + thread_id TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL DEFAULT '', + checkpoint_id TEXT NOT NULL, + parent_checkpoint_id TEXT, + type TEXT, + checkpoint JSONB NOT NULL, + metadata JSONB NOT NULL DEFAULT '{}', + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id) +);""", + """CREATE TABLE IF NOT EXISTS checkpoint_blobs ( + thread_id TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL DEFAULT '', + channel TEXT NOT NULL, + version TEXT NOT NULL, + type TEXT NOT NULL, + blob BYTEA NOT NULL, + PRIMARY KEY (thread_id, checkpoint_ns, channel, version) +);""", + """CREATE TABLE IF NOT EXISTS checkpoint_writes ( + thread_id TEXT NOT NULL, + checkpoint_ns TEXT NOT NULL DEFAULT '', + checkpoint_id TEXT NOT NULL, + task_id TEXT NOT NULL, + idx INTEGER NOT NULL, + channel TEXT NOT NULL, + type TEXT, + blob BYTEA NOT NULL, + PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) +);""", +] + SELECT_SQL = """ select thread_id, @@ -42,43 +82,6 @@ select ) as pending_writes from checkpoints """ -CREATE_CHECKPOINTS_SQL = """ - CREATE TABLE IF NOT EXISTS checkpoints ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - checkpoint_id TEXT NOT NULL, - parent_checkpoint_id TEXT, - type TEXT, - checkpoint JSONB NOT NULL, - metadata JSONB NOT NULL DEFAULT '{}', - PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id) -); -""" -CREATE_CHECKPOINT_BLOBS_SQL = """ - CREATE TABLE IF NOT EXISTS checkpoint_blobs ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - channel TEXT NOT NULL, - version TEXT NOT NULL, - type TEXT NOT NULL, - blob BYTEA NOT NULL, - PRIMARY KEY (thread_id, checkpoint_ns, channel, version) -);""" - -CREATE_CHECKPOINT_WRITES_SQL = """ - CREATE TABLE IF NOT EXISTS checkpoint_writes ( - thread_id TEXT NOT NULL, - checkpoint_ns TEXT NOT NULL DEFAULT '', - checkpoint_id TEXT NOT NULL, - task_id TEXT NOT NULL, - idx INTEGER NOT NULL, - channel TEXT NOT NULL, - type TEXT, - blob BYTEA NOT NULL, - PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) -); -""" - UPSERT_CHECKPOINT_BLOBS_SQL = """ INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob) VALUES (%s, %s, %s, %s, %s, %s) @@ -103,9 +106,7 @@ UPSERT_CHECKPOINT_WRITES_SQL = """ class BasePostgresSaver(BaseCheckpointSaver): SELECT_SQL = SELECT_SQL - CREATE_CHECKPOINTS_SQL = CREATE_CHECKPOINTS_SQL - CREATE_CHECKPOINT_BLOBS_SQL = CREATE_CHECKPOINT_BLOBS_SQL - CREATE_CHECKPOINT_WRITES_SQL = CREATE_CHECKPOINT_WRITES_SQL + MIGRATIONS = MIGRATIONS UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL diff --git a/libs/checkpoint-postgres/tests/compose-postgres.yml b/libs/checkpoint-postgres/tests/compose-postgres.yml index 079b412b4..42d8c3781 100644 --- a/libs/checkpoint-postgres/tests/compose-postgres.yml +++ b/libs/checkpoint-postgres/tests/compose-postgres.yml @@ -2,7 +2,7 @@ services: postgres-test: image: postgres:16 ports: - - "5432:5432" + - "5441:5432" environment: POSTGRES_DB: postgres POSTGRES_USER: postgres @@ -13,4 +13,4 @@ services: timeout: 1s retries: 5 interval: 60s - start_interval: 1s \ No newline at end of file + start_interval: 1s diff --git a/libs/checkpoint-postgres/tests/conftest.py b/libs/checkpoint-postgres/tests/conftest.py index d103a4530..6100be1d4 100644 --- a/libs/checkpoint-postgres/tests/conftest.py +++ b/libs/checkpoint-postgres/tests/conftest.py @@ -3,7 +3,7 @@ from psycopg import AsyncConnection from psycopg.errors import UndefinedTable from psycopg.rows import dict_row -DEFAULT_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" +DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable" @pytest.fixture(scope="function") diff --git a/libs/langgraph/tests/compose-postgres.yml b/libs/langgraph/tests/compose-postgres.yml index 079b412b4..80904ce90 100644 --- a/libs/langgraph/tests/compose-postgres.yml +++ b/libs/langgraph/tests/compose-postgres.yml @@ -2,7 +2,7 @@ services: postgres-test: image: postgres:16 ports: - - "5432:5432" + - "5442:5432" environment: POSTGRES_DB: postgres POSTGRES_USER: postgres @@ -13,4 +13,4 @@ services: timeout: 1s retries: 5 interval: 60s - start_interval: 1s \ No newline at end of file + start_interval: 1s diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index e882f92a3..44b441982 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -7,7 +7,7 @@ from psycopg.rows import dict_row from pytest_mock import MockerFixture DEFAULT_POSTGRES_URI = ( - "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable" + "postgres://postgres:postgres@localhost:5442/postgres?sslmode=disable" )