langgraph-checkpoint-postgres 1.0.3__tar.gz → 1.0.4__tar.gz

{langgraph_checkpoint_postgres-1.0.3 → langgraph_checkpoint_postgres-1.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.3
+Version: 1.0.4
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT
@@ -13,7 +13,8 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: langgraph-checkpoint (>=1.0.1,<2.0.0)
 Requires-Dist: orjson (>=3.10.1)
-Requires-Dist: psycopg[binary] (>=3.1.19)
+Requires-Dist: psycopg (>=3.0.0,<4.0.0)
+Requires-Dist: psycopg-pool (>=3.0.0,<4.0.0)
 Project-URL: Repository, https://www.github.com/langchain-ai/langgraph
 Description-Content-Type: text/markdown
 
@@ -21,6 +22,10 @@ Description-Content-Type: text/markdown
 
 Implementation of LangGraph CheckpointSaver that uses Postgres.
 
+## Dependencies
+
+By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`).
+
 ## Usage
 
 > [!IMPORTANT]
@@ -63,7 +68,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -106,7 +110,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint

{langgraph_checkpoint_postgres-1.0.3 → langgraph_checkpoint_postgres-1.0.4}/README.md

@@ -2,6 +2,10 @@
 
 Implementation of LangGraph CheckpointSaver that uses Postgres.
 
+## Dependencies
+
+By default `langgraph-checkpoint-postgres` installs `psycopg` (Psycopg 3) without any extras. However, you can choose a specific installation that best suits your needs [here](https://www.psycopg.org/psycopg3/docs/basic/install.html) (for example, `psycopg[binary]`).
+
 ## Usage
 
 > [!IMPORTANT]
@@ -44,7 +48,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -87,7 +90,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint

{langgraph_checkpoint_postgres-1.0.3 → langgraph_checkpoint_postgres-1.0.4}/langgraph/checkpoint/postgres/__init__.py

@@ -1,12 +1,13 @@
 import threading
 from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional
+from typing import Any, Iterator, List, Optional, Union
 
 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
+from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.base import (
     ChannelVersions,
@@ -21,16 +22,32 @@ from langgraph.checkpoint.postgres.base import (
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
 
+@contextmanager
+def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]:
+    if isinstance(conn, Connection):
+        yield conn
+    elif isinstance(conn, ConnectionPool):
+        with conn.connection() as conn:
+            yield conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")
+
+
 class PostgresSaver(BasePostgresSaver):
     lock: threading.Lock
 
     def __init__(
         self,
-        conn: Connection,
+        conn: Union[Connection, ConnectionPool],
         pipe: Optional[Pipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
         super().__init__(serde=serde)
+        if isinstance(conn, ConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single Connection, not ConnectionPool."
+            )
+
         self.conn = conn
         self.pipe = pipe
         self.lock = threading.Lock()
@@ -65,22 +82,21 @@ class PostgresSaver(BasePostgresSaver):
         already exist and runs database migrations. It MUST be called directly by the user
         the first time checkpointer is used.
         """
-        with self.lock:
-            with self.conn.cursor(binary=True) as cur:
-                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()
+        with self._cursor() as cur:
+            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()
 
     def list(
         self,
@@ -127,30 +143,33 @@ class PostgresSaver(BasePostgresSaver):
         if limit:
             query += f" LIMIT {limit}"
         # if we change this to use .stream() we need to make sure to close the cursor
-        for value in self.conn.execute(query, args, binary=True):
-            yield CheckpointTuple(
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["checkpoint_id"],
-                    }
-                },
-                {
-                    **self._load_checkpoint(value["checkpoint"]),
-                    "channel_values": self._load_blobs(value["channel_values"]),
-                },
-                self._load_metadata(value["metadata"]),
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["parent_checkpoint_id"],
+        with self._cursor() as cur:
+            cur.execute(query, args, binary=True)
+            for value in cur:
+                yield CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["checkpoint_id"],
+                        }
+                    },
+                    {
+                        **self._load_checkpoint(value["checkpoint"]),
+                        "channel_values": self._load_blobs(value["channel_values"]),
+                    },
+                    self._load_metadata(value["metadata"]),
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["parent_checkpoint_id"],
+                        }
                     }
-                }
-                if value["parent_checkpoint_id"]
-                else None,
-            )
+                    if value["parent_checkpoint_id"]
+                    else None,
+                    self._load_writes(value["pending_writes"]),
+                )
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database.
@@ -198,7 +217,7 @@ class PostgresSaver(BasePostgresSaver):
             where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
 
         with self._cursor() as cur:
-            cur = self.conn.execute(
+            cur.execute(
                 self.SELECT_SQL + where,
                 args,
                 binary=True,
@@ -316,7 +335,7 @@ class PostgresSaver(BasePostgresSaver):
             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:
+        with self._cursor(pipeline=True) as cur:
             cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 self._dump_writes(
@@ -330,21 +349,24 @@ class PostgresSaver(BasePostgresSaver):
 
     @contextmanager
     def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
-        if self.pipe:
-            # a connection in pipeline mode can be used concurrently
-            # in multiple threads/coroutines, but only one cursor can be
-            # used at a time
-            try:
-                with self.conn.cursor(binary=True) as cur:
+        with _get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    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
+                with self.lock, conn.pipeline(), conn.cursor(
+                    binary=True, row_factory=dict_row
+                ) as cur:
+                    yield cur
+            else:
+                with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
                     yield cur
-            finally:
-                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
-            with self.lock, self.conn.pipeline(), self.conn.cursor(binary=True) as cur:
-                yield cur
-        else:
-            with self.lock, self.conn.cursor(binary=True) as cur:
-                yield cur

{langgraph_checkpoint_postgres-1.0.3 → langgraph_checkpoint_postgres-1.0.4}/langgraph/checkpoint/postgres/aio.py

@@ -1,12 +1,13 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncIterator, Optional, Union
 
 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
+from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.base import (
     ChannelVersions,
@@ -19,16 +20,34 @@ from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
 
+@asynccontextmanager
+async def _get_connection(
+    conn: Union[AsyncConnection, AsyncConnectionPool],
+) -> AsyncIterator[AsyncConnection]:
+    if isinstance(conn, AsyncConnection):
+        yield conn
+    elif isinstance(conn, AsyncConnectionPool):
+        async with conn.connection() as conn:
+            yield conn
+    else:
+        raise TypeError(f"Invalid connection type: {type(conn)}")
+
+
 class AsyncPostgresSaver(BasePostgresSaver):
     lock: asyncio.Lock
 
     def __init__(
         self,
-        conn: AsyncConnection,
+        conn: Union[AsyncConnection, AsyncConnectionPool],
         pipe: Optional[AsyncPipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
         super().__init__(serde=serde)
+        if isinstance(conn, AsyncConnectionPool) and pipe is not None:
+            raise ValueError(
+                "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool."
+            )
+
         self.conn = conn
         self.pipe = pipe
         self.lock = asyncio.Lock()
@@ -63,25 +82,22 @@ class AsyncPostgresSaver(BasePostgresSaver):
         already exist and runs database migrations. It MUST be called directly by the user
         the first time checkpointer is used.
         """
-        async with self.lock:
-            async with self.conn.cursor(binary=True) as cur:
-                try:
-                    results = await cur.execute(
-                        "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                    )
-                    version = (await results.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()
+        async with self._cursor() as cur:
+            try:
+                results = await cur.execute(
+                    "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
+                )
+                version = (await results.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()
 
     async def alist(
         self,
@@ -110,32 +126,35 @@ class AsyncPostgresSaver(BasePostgresSaver):
         if limit:
             query += f" LIMIT {limit}"
         # if we change this to use .stream() we need to make sure to close the cursor
-        async for value in await self.conn.execute(query, args, binary=True):
-            yield CheckpointTuple(
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["checkpoint_id"],
-                    }
-                },
-                {
-                    **self._load_checkpoint(value["checkpoint"]),
-                    "channel_values": await asyncio.to_thread(
-                        self._load_blobs, value["channel_values"]
-                    ),
-                },
-                self._load_metadata(value["metadata"]),
-                {
-                    "configurable": {
-                        "thread_id": value["thread_id"],
-                        "checkpoint_ns": value["checkpoint_ns"],
-                        "checkpoint_id": value["parent_checkpoint_id"],
+        async with self._cursor() as cur:
+            await cur.execute(query, args, binary=True)
+            async for value in cur:
+                yield CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["checkpoint_id"],
+                        }
+                    },
+                    {
+                        **self._load_checkpoint(value["checkpoint"]),
+                        "channel_values": await asyncio.to_thread(
+                            self._load_blobs, value["channel_values"]
+                        ),
+                    },
+                    self._load_metadata(value["metadata"]),
+                    {
+                        "configurable": {
+                            "thread_id": value["thread_id"],
+                            "checkpoint_ns": value["checkpoint_ns"],
+                            "checkpoint_id": value["parent_checkpoint_id"],
+                        }
                     }
-                }
-                if value["parent_checkpoint_id"]
-                else None,
-            )
+                    if value["parent_checkpoint_id"]
+                    else None,
+                    await asyncio.to_thread(self._load_writes, value["pending_writes"]),
+                )
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the database asynchronously.
@@ -162,7 +181,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
             where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
 
         async with self._cursor() as cur:
-            cur = await self.conn.execute(
+            await cur.execute(
                 self.SELECT_SQL + where,
                 args,
                 binary=True,
@@ -272,7 +291,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
             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:
+        async with self._cursor(pipeline=True) as cur:
             await cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 await asyncio.to_thread(
@@ -287,23 +306,26 @@ class AsyncPostgresSaver(BasePostgresSaver):
 
     @asynccontextmanager
     async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
-        if self.pipe:
-            # a connection in pipeline mode can be used concurrently
-            # in multiple threads/coroutines, but only one cursor can be
-            # used at a time
-            try:
-                async with self.conn.cursor(binary=True) as cur:
+        async with _get_connection(self.conn) as conn:
+            if self.pipe:
+                # a connection in pipeline mode can be used concurrently
+                # in multiple threads/coroutines, but only one cursor can be
+                # used at a time
+                try:
+                    async with conn.cursor(binary=True, row_factory=dict_row) as cur:
+                        yield cur
+                finally:
+                    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
+                async with self.lock, conn.pipeline(), conn.cursor(
+                    binary=True, row_factory=dict_row
+                ) as cur:
+                    yield cur
+            else:
+                async with self.lock, conn.cursor(
+                    binary=True, row_factory=dict_row
+                ) as cur:
                     yield cur
-            finally:
-                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
-            async with self.lock, self.conn.pipeline(), self.conn.cursor(
-                binary=True
-            ) as cur:
-                yield cur
-        else:
-            async with self.lock, self.conn.cursor(binary=True) as cur:
-                yield cur

{langgraph_checkpoint_postgres-1.0.3 → langgraph_checkpoint_postgres-1.0.4}/langgraph/checkpoint/postgres/base.py

@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
     EmptyChannelError,
@@ -112,6 +113,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
     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:
@@ -199,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 checkpoint_ns,
                 checkpoint_id,
                 task_id,
-                idx,
+                WRITES_IDX_MAP.get(channel, idx),
                 channel,
                 *self.serde.dumps_typed(value),
             )
@@ -253,9 +255,13 @@ class BasePostgresSaver(BaseCheckpointSaver):
         if config:
             wheres.append("thread_id = %s ")
             param_values.append(config["configurable"]["thread_id"])
-            checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
-            wheres.append("checkpoint_ns = %s")
-            param_values.append(checkpoint_ns)
+            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+                wheres.append("checkpoint_ns = %s")
+                param_values.append(checkpoint_ns)
+
+            if checkpoint_id := get_checkpoint_id(config):
+                wheres.append("checkpoint_id = %s ")
+                param_values.append(checkpoint_id)
 
         # construct predicate for metadata filter
         if filter:

{langgraph_checkpoint_postgres-1.0.3 → langgraph_checkpoint_postgres-1.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.3"
+version = "1.0.4"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -12,7 +12,8 @@ packages = [{ include = "langgraph" }]
 python = "^3.9.0,<4.0"
 langgraph-checkpoint = "^1.0.1"
 orjson = ">=3.10.1"
-psycopg = {extras = ["binary"], version = ">=3.1.19"}
+psycopg = "^3.0.0"
+psycopg-pool = "^3.0.0"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.1.4"
@@ -23,7 +24,7 @@ pytest-asyncio = "^0.21.1"
 pytest-mock = "^3.11.1"
 pytest-watch = "^4.2.0"
 mypy = "^1.10.0"
-psycopg-pool = "^3.2.2"
+psycopg = {extras = ["binary"], version = ">=3.0.0"}
 langgraph-checkpoint = {path = "../checkpoint", develop = true}
 
 [tool.pytest.ini_options]