langgraph-checkpoint-postgres 1.0.0__tar.gz → 1.0.1__tar.gz

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.0
+Version: 1.0.1
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.1}/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
 
@@ -23,8 +24,6 @@ from langgraph.checkpoint.serde.base import SerializerProtocol
 class PostgresSaver(BasePostgresSaver):
     lock: threading.Lock
 
-    is_setup: bool
-
     def __init__(
         self,
         conn: Connection,
@@ -35,7 +34,6 @@ class PostgresSaver(BasePostgresSaver):
         self.conn = conn
         self.pipe = pipe
         self.lock = threading.Lock()
-        self.is_setup = False
 
     @classmethod
     @contextmanager
@@ -67,23 +65,23 @@ 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:
-            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()
 
-            self.is_setup = True
-
     def list(
         self,
         config: Optional[RunnableConfig],
@@ -92,7 +90,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:
@@ -111,7 +108,7 @@ class PostgresSaver(BasePostgresSaver):
                     **self._load_checkpoint(value["checkpoint"]),
                     "channel_values": self._load_blobs(value["channel_values"]),
                 },
-                value["metadata"],
+                self._load_metadata(value["metadata"]),
                 {
                     "configurable": {
                         "thread_id": value["thread_id"],
@@ -124,7 +121,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", "")
@@ -155,7 +151,7 @@ class PostgresSaver(BasePostgresSaver):
                         **self._load_checkpoint(value["checkpoint"]),
                         "channel_values": self._load_blobs(value["channel_values"]),
                     },
-                    value["metadata"],
+                    self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
                             "thread_id": thread_id,
@@ -210,7 +206,7 @@ class PostgresSaver(BasePostgresSaver):
                     checkpoint["id"],
                     checkpoint_id,
                     Jsonb(self._dump_checkpoint(copy)),
-                    Jsonb(metadata),
+                    self._dump_metadata(metadata),
                 ),
             )
         return next_config
@@ -235,7 +231,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
@@ -244,7 +239,8 @@ class PostgresSaver(BasePostgresSaver):
                 with self.conn.cursor(binary=True) as cur:
                     yield cur
             finally:
-                self.pipe.sync()
+                if pipeline:
+                    self.pipe.sync()
         elif pipeline:
             # a connection not in pipeline mode can only be used by one
             # thread/coroutine at a time, so we acquire a lock

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.1}/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
 
@@ -21,8 +22,6 @@ from langgraph.checkpoint.serde.base import SerializerProtocol
 class AsyncPostgresSaver(BasePostgresSaver):
     lock: asyncio.Lock
 
-    is_setup: bool
-
     def __init__(
         self,
         conn: AsyncConnection,
@@ -33,7 +32,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
         self.conn = conn
         self.pipe = pipe
         self.lock = asyncio.Lock()
-        self.is_setup = False
 
     @classmethod
     @asynccontextmanager
@@ -65,23 +63,27 @@ 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:
-            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()
 
-            self.is_setup = True
-
     async def alist(
         self,
         config: Optional[RunnableConfig],
@@ -90,7 +92,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:
@@ -111,7 +112,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                         self._load_blobs, value["channel_values"]
                     ),
                 },
-                value["metadata"],
+                self._load_metadata(value["metadata"]),
                 {
                     "configurable": {
                         "thread_id": value["thread_id"],
@@ -124,7 +125,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", "")
@@ -157,7 +157,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                             self._load_blobs, value["channel_values"]
                         ),
                     },
-                    value["metadata"],
+                    self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
                             "thread_id": thread_id,
@@ -177,7 +177,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")
@@ -214,7 +213,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     checkpoint["id"],
                     checkpoint_id,
                     Jsonb(self._dump_checkpoint(copy)),
-                    Jsonb(metadata),
+                    self._dump_metadata(metadata),
                 ),
             )
         return next_config
@@ -240,7 +239,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
@@ -249,7 +247,8 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 async with self.conn.cursor(binary=True) as cur:
                     yield cur
             finally:
-                await self.pipe.sync()
+                if pipeline:
+                    await self.pipe.sync()
         elif pipeline:
             # a connection not in pipeline mode can only be used by one
             # thread/coroutine at a time, so we acquire a lock

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.1}/langgraph/checkpoint/postgres/base.py

@@ -11,10 +11,51 @@ from langgraph.checkpoint.base import (
     EmptyChannelError,
     get_checkpoint_id,
 )
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 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 +83,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,12 +107,11 @@ 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
+    jsonplus_serde = JsonPlusSerializer()
 
     def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint:
         if len(checkpoint["pending_sends"]) == 2 and all(
@@ -202,6 +205,19 @@ class BasePostgresSaver(BaseCheckpointSaver):
             for idx, (channel, value) in enumerate(writes)
         ]
 
+    def _load_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
+        return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata))
+
+    def _dump_metadata(self, metadata) -> str:
+        serialized_metadata_type, serialized_metadata = self.jsonplus_serde.dumps_typed(
+            metadata
+        )
+        if serialized_metadata_type != "json":
+            raise TypeError(
+                f"Failed to properly serialize metadata -- expected 'json', got '{serialized_metadata_type}'"
+            )
+        return serialized_metadata.decode()
+
     def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
         if current is None:
             current_v = 0

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.0"
+version = "1.0.1"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"