postgres: Add migration tracking

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()
 

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()
 

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

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
+      start_interval: 1s

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")

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
+      start_interval: 1s

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"
 )