From 1e237bf33a27eee0ced858a31855213932eb3a96 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 7 Aug 2024 08:54:08 -0700 Subject: [PATCH 1/3] postgres: Add migration tracking --- .../langgraph/checkpoint/postgres/__init__.py | 21 +++-- .../langgraph/checkpoint/postgres/aio.py | 27 ++++--- .../langgraph/checkpoint/postgres/base.py | 81 ++++++++++--------- .../tests/compose-postgres.yml | 4 +- libs/checkpoint-postgres/tests/conftest.py | 2 +- libs/langgraph/tests/compose-postgres.yml | 4 +- libs/langgraph/tests/conftest.py | 2 +- 7 files changed, 78 insertions(+), 63 deletions(-) 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" ) From 51b62ca0bdfe06b2733828005c71bf32ce46c269 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 7 Aug 2024 12:18:55 -0400 Subject: [PATCH 2/3] remove setup --- .../langgraph/checkpoint/postgres/__init__.py | 3 --- .../langgraph/checkpoint/postgres/aio.py | 4 ---- libs/checkpoint-postgres/tests/test_async.py | 4 +++- libs/checkpoint-postgres/tests/test_sync.py | 2 ++ libs/langgraph/tests/conftest.py | 8 ++++++++ 5 files changed, 13 insertions(+), 8 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 97f0f0570..5649e4e8c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -97,7 +97,6 @@ class PostgresSaver(BasePostgresSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> Iterator[CheckpointTuple]: - self.setup() where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" if limit: @@ -129,7 +128,6 @@ class PostgresSaver(BasePostgresSaver): ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - self.setup() thread_id = config["configurable"]["thread_id"] checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") @@ -240,7 +238,6 @@ class PostgresSaver(BasePostgresSaver): @contextmanager def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]: - self.setup() if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index e7f080002..7b2895b1c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -99,7 +99,6 @@ class AsyncPostgresSaver(BasePostgresSaver): before: Optional[RunnableConfig] = None, limit: Optional[int] = None, ) -> AsyncIterator[CheckpointTuple]: - await self.setup() where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" if limit: @@ -133,7 +132,6 @@ class AsyncPostgresSaver(BasePostgresSaver): ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: - await self.setup() thread_id = config["configurable"]["thread_id"] checkpoint_id = get_checkpoint_id(config) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") @@ -186,7 +184,6 @@ class AsyncPostgresSaver(BasePostgresSaver): metadata: CheckpointMetadata, new_versions: ChannelVersions, ) -> RunnableConfig: - await self.setup() configurable = config["configurable"].copy() thread_id = configurable.pop("thread_id") checkpoint_ns = configurable.pop("checkpoint_ns") @@ -249,7 +246,6 @@ class AsyncPostgresSaver(BasePostgresSaver): @asynccontextmanager async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]: - await self.setup() if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index e392ddd65..e94cf32ae 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -13,7 +13,7 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver class TestAsyncPostgresSaver: @pytest.fixture(autouse=True) - def setup(self): + async def setup(self): # objects for test setup self.config_1: RunnableConfig = { "configurable": { @@ -55,6 +55,8 @@ class TestAsyncPostgresSaver: "score": None, } self.metadata_3: CheckpointMetadata = {} + async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: + await saver.setup() async def test_asearch(self): async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver: diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index 0d642ee62..dfae82907 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -55,6 +55,8 @@ class TestPostgresSaver: "score": None, } self.metadata_3: CheckpointMetadata = {} + with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: + saver.setup() def test_search(self): with PostgresSaver.from_conn_string(DEFAULT_URI) as saver: diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 44b441982..eb8f20ce9 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -6,6 +6,8 @@ from psycopg.errors import UndefinedTable from psycopg.rows import dict_row from pytest_mock import MockerFixture +from langgraph.checkpoint.postgres import PostgresSaver + DEFAULT_POSTGRES_URI = ( "postgres://postgres:postgres@localhost:5442/postgres?sslmode=disable" ) @@ -27,6 +29,12 @@ async def conn(): yield conn +@pytest.fixture(scope="session", autouse=True) +def setup_before_all_tests(): + with PostgresSaver.from_conn_string(DEFAULT_POSTGRES_URI) as checkpointer: + checkpointer.setup() + + @pytest.fixture(scope="function", autouse=True) async def clear_test_db(conn): """Delete all tables before each test.""" From 54da68cfe71d7def27f15519182dc716d0fd0bc0 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 7 Aug 2024 12:25:34 -0400 Subject: [PATCH 3/3] remove flag --- .../langgraph/checkpoint/postgres/__init__.py | 7 ------- .../langgraph/checkpoint/postgres/aio.py | 7 ------- 2 files changed, 14 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 5649e4e8c..6b7c465fd 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -24,8 +24,6 @@ from langgraph.checkpoint.serde.base import SerializerProtocol class PostgresSaver(BasePostgresSaver): lock: threading.Lock - is_setup: bool - def __init__( self, conn: Connection, @@ -36,7 +34,6 @@ class PostgresSaver(BasePostgresSaver): self.conn = conn self.pipe = pipe self.lock = threading.Lock() - self.is_setup = False @classmethod @contextmanager @@ -68,8 +65,6 @@ class PostgresSaver(BasePostgresSaver): already exist. It is called automatically when needed and should not be called directly by the user. """ - if self.is_setup: - return with self.lock: with self.conn.cursor(binary=True) as cur: try: @@ -87,8 +82,6 @@ class PostgresSaver(BasePostgresSaver): if self.pipe: self.pipe.sync() - self.is_setup = True - def list( self, config: Optional[RunnableConfig], diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 7b2895b1c..0fc86ec3b 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -22,8 +22,6 @@ from langgraph.checkpoint.serde.base import SerializerProtocol class AsyncPostgresSaver(BasePostgresSaver): lock: asyncio.Lock - is_setup: bool - def __init__( self, conn: AsyncConnection, @@ -34,7 +32,6 @@ class AsyncPostgresSaver(BasePostgresSaver): self.conn = conn self.pipe = pipe self.lock = asyncio.Lock() - self.is_setup = False @classmethod @asynccontextmanager @@ -66,8 +63,6 @@ class AsyncPostgresSaver(BasePostgresSaver): already exist. It is called automatically when needed and should not be called directly by the user. """ - if self.is_setup: - return async with self.lock: async with self.conn.cursor(binary=True) as cur: try: @@ -89,8 +84,6 @@ class AsyncPostgresSaver(BasePostgresSaver): if self.pipe: await self.pipe.sync() - self.is_setup = True - async def alist( self, config: Optional[RunnableConfig],