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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.0
+Version: 1.0.2
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT
@@ -23,6 +23,9 @@ Implementation of LangGraph CheckpointSaver that uses Postgres.
 
 ## Usage
 
+> [!IMPORTANT]
+> When using Postgres checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
+
 ```python
 from langgraph.checkpoint.postgres import PostgresSaver
 
@@ -31,6 +34,8 @@ read_config = {"configurable": {"thread_id": "1"}}
 
 DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
 with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
+    # call .setup() the first time you're using the checkpointer
+    checkpointer.setup()
     checkpoint = {
         "v": 1,
         "ts": "2024-07-31T20:14:19.804150+00:00",

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.2}/README.md

@@ -4,6 +4,9 @@ Implementation of LangGraph CheckpointSaver that uses Postgres.
 
 ## Usage
 
+> [!IMPORTANT]
+> When using Postgres checkpointers for the first time, make sure to call `.setup()` method on them to create required tables. See example below.
+
 ```python
 from langgraph.checkpoint.postgres import PostgresSaver
 
@@ -12,6 +15,8 @@ read_config = {"configurable": {"thread_id": "1"}}
 
 DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
 with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
+    # call .setup() the first time you're using the checkpointer
+    checkpointer.setup()
     checkpoint = {
         "v": 1,
         "ts": "2024-07-31T20:14:19.804150+00:00",

{langgraph_checkpoint_postgres-1.0.0 → langgraph_checkpoint_postgres-1.0.2}/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
@@ -63,27 +61,27 @@ class PostgresSaver(BasePostgresSaver):
     def setup(self) -> None:
         """Set up the checkpoint database asynchronously.
 
-        This method creates the necessary tables in the SQLite database if they don't
-        already exist. It is called automatically when needed and should not be called
-        directly by the user.
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
         """
-        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,38 @@ class PostgresSaver(BasePostgresSaver):
         before: Optional[RunnableConfig] = None,
         limit: Optional[int] = None,
     ) -> Iterator[CheckpointTuple]:
-        self.setup()
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (RunnableConfig): The config to use for listing the checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata. Defaults to None.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None.
+
+        Yields:
+            Iterator[CheckpointTuple]: An iterator of checkpoint tuples.
+
+        Examples:
+            >>> from langgraph.checkpoint.postgres import PostgresSaver
+            >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+            >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
+            ... # Run a graph, then list the checkpoints
+            >>>     config = {"configurable": {"thread_id": "1"}}
+            >>>     checkpoints = list(memory.list(config, limit=2))
+            >>> print(checkpoints)
+            [CheckpointTuple(...), CheckpointTuple(...)]
+
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> before = {"configurable": {"checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875"}}
+            >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
+            ... # Run a graph, then list the checkpoints
+            >>>     checkpoints = list(memory.list(config, before=before))
+            >>> print(checkpoints)
+            [CheckpointTuple(...), ...]
+        """
         where, args = self._search_where(config, filter, before)
         query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
         if limit:
@@ -111,7 +140,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 +153,40 @@ class PostgresSaver(BasePostgresSaver):
             )
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
-        self.setup()
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and timestamp is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+
+        Examples:
+
+            Basic:
+            >>> config = {"configurable": {"thread_id": "1"}}
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+
+            With timestamp:
+
+            >>> config = {
+            ...    "configurable": {
+            ...        "thread_id": "1",
+            ...        "checkpoint_ns": "",
+            ...        "checkpoint_id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+            ...    }
+            ... }
+            >>> checkpoint_tuple = memory.get_tuple(config)
+            >>> print(checkpoint_tuple)
+            CheckpointTuple(...)
+        """  # noqa
         thread_id = config["configurable"]["thread_id"]
         checkpoint_id = get_checkpoint_id(config)
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -155,7 +217,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,
@@ -175,6 +237,31 @@ class PostgresSaver(BasePostgresSaver):
         metadata: CheckpointMetadata,
         new_versions: ChannelVersions,
     ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+
+        Examples:
+
+            >>> from langgraph.checkpoint.postgres import PostgresSaver
+            >>> DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+            >>> with PostgresSaver.from_conn_string(DB_URI) as memory:
+            >>>     config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+            >>>     checkpoint = {"ts": "2024-05-04T06:32:42.235444+00:00", "id": "1ef4f797-8335-6428-8001-8a1503f9b875", "data": {"key": "value"}}
+            >>>     saved_config = memory.put(config, checkpoint, {"source": "input", "step": 1, "writes": {"key": "value"}}, {})
+            >>> print(saved_config)
+            {'configurable': {'thread_id': '1', 'checkpoint_ns': '', 'checkpoint_id': '1ef4f797-8335-6428-8001-8a1503f9b875'}}
+        """
         configurable = config["configurable"].copy()
         thread_id = configurable.pop("thread_id")
         checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -198,7 +285,6 @@ class PostgresSaver(BasePostgresSaver):
                     thread_id,
                     checkpoint_ns,
                     copy.pop("channel_values"),
-                    copy["channel_versions"],
                     new_versions,
                 ),
             )
@@ -210,7 +296,7 @@ class PostgresSaver(BasePostgresSaver):
                     checkpoint["id"],
                     checkpoint_id,
                     Jsonb(self._dump_checkpoint(copy)),
-                    Jsonb(metadata),
+                    self._dump_metadata(metadata),
                 ),
             )
         return next_config
@@ -221,6 +307,15 @@ class PostgresSaver(BasePostgresSaver):
         writes: List[tuple[str, Any]],
         task_id: str,
     ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the Postgres database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (List[Tuple[str, Any]]): List of writes to store.
+            task_id (str): Identifier for the task creating the writes.
+        """
         with self._cursor() as cur:
             cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
@@ -235,7 +330,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 +338,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.2}/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
@@ -61,27 +59,31 @@ class AsyncPostgresSaver(BasePostgresSaver):
     async def setup(self) -> None:
         """Set up the checkpoint database asynchronously.
 
-        This method creates the necessary tables in the SQLite database if they don't
-        already exist. It is called automatically when needed and should not be called
-        directly by the user.
+        This method creates the necessary tables in the Postgres database if they don't
+        already exist and runs database migrations. It MUST be called directly by the user
+        the first time checkpointer is used.
         """
-        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,20 @@ class AsyncPostgresSaver(BasePostgresSaver):
         before: Optional[RunnableConfig] = None,
         limit: Optional[int] = None,
     ) -> AsyncIterator[CheckpointTuple]:
-        await self.setup()
+        """List checkpoints from the database asynchronously.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Yields:
+            AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
+        """
         where, args = self._search_where(config, filter, before)
         query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
         if limit:
@@ -111,7 +126,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 +139,19 @@ class AsyncPostgresSaver(BasePostgresSaver):
             )
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
-        await self.setup()
+        """Get a checkpoint tuple from the database asynchronously.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
         thread_id = config["configurable"]["thread_id"]
         checkpoint_id = get_checkpoint_id(config)
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -157,7 +184,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                             self._load_blobs, value["channel_values"]
                         ),
                     },
-                    value["metadata"],
+                    self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
                             "thread_id": thread_id,
@@ -177,7 +204,20 @@ class AsyncPostgresSaver(BasePostgresSaver):
         metadata: CheckpointMetadata,
         new_versions: ChannelVersions,
     ) -> RunnableConfig:
-        await self.setup()
+        """Save a checkpoint to the database asynchronously.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
         configurable = config["configurable"].copy()
         thread_id = configurable.pop("thread_id")
         checkpoint_ns = configurable.pop("checkpoint_ns")
@@ -202,7 +242,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     thread_id,
                     checkpoint_ns,
                     copy.pop("channel_values"),
-                    copy["channel_versions"],
                     new_versions,
                 ),
             )
@@ -214,7 +253,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     checkpoint["id"],
                     checkpoint_id,
                     Jsonb(self._dump_checkpoint(copy)),
-                    Jsonb(metadata),
+                    self._dump_metadata(metadata),
                 ),
             )
         return next_config
@@ -225,6 +264,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
         writes: list[tuple[str, Any]],
         task_id: str,
     ) -> None:
+        """Store intermediate writes linked to a checkpoint asynchronously.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
         async with self._cursor() as cur:
             await cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
@@ -240,7 +288,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 +296,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.2}/langgraph/checkpoint/postgres/base.py

@@ -11,10 +11,52 @@ 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,
+    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)
+);""",
+    "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
+]
+
 SELECT_SQL = """
 select
     thread_id,
@@ -42,43 +84,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 +108,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(
@@ -137,6 +141,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
         return {
             k.decode(): self.serde.loads_typed((t.decode(), v))
             for k, t, v in blob_values
+            if t.decode() != "empty"
         }
 
     def _dump_blobs(
@@ -145,24 +150,23 @@ class BasePostgresSaver(BaseCheckpointSaver):
         checkpoint_ns: str,
         values: dict[str, Any],
         versions: dict[str, str],
-        new_versions: Optional[dict[str, str]],
     ) -> list[tuple[str, str, str, str, str, bytes]]:
         if not versions:
             return []
 
-        if new_versions:
-            versions = new_versions
-
         return [
             (
                 thread_id,
                 checkpoint_ns,
                 k,
                 ver,
-                *self.serde.dumps_typed(values[k]),
+                *(
+                    self.serde.dumps_typed(values[k])
+                    if k in values
+                    else ("empty", None)
+                ),
             )
             for k, ver in versions.items()
-            if k in values
         ]
 
     def _load_writes(
@@ -202,6 +206,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.2}/pyproject.toml

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