langgraph-checkpoint-postgres 1.0.10__tar.gz → 2.0.0__tar.gz

{langgraph_checkpoint_postgres-1.0.10 → langgraph_checkpoint_postgres-2.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.10
+Version: 2.0.0
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT
@@ -11,7 +11,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: langgraph-checkpoint (>=1.0.11,<2.0.0)
+Requires-Dist: langgraph-checkpoint (>=2.0.0,<3.0.0)
 Requires-Dist: orjson (>=3.10.1)
 Requires-Dist: psycopg (>=3.0.0,<4.0.0)
 Requires-Dist: psycopg-pool (>=3.0.0,<4.0.0)

{langgraph_checkpoint_postgres-1.0.10 → langgraph_checkpoint_postgres-2.0.0}/langgraph/store/postgres/aio.py

@@ -1,8 +1,18 @@
 import asyncio
 import logging
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Iterable, Sequence, cast
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Iterable,
+    Optional,
+    Sequence,
+    Union,
+    cast,
+)
 
+import orjson
 from psycopg import AsyncConnection, AsyncCursor
 from psycopg.errors import UndefinedTable
 from psycopg.rows import dict_row
@@ -11,6 +21,7 @@ from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, Sea
 from langgraph.store.postgres.base import (
     BasePostgresStore,
     Row,
+    _decode_ns_bytes,
     _group_ops,
     _row_to_item,
 )
@@ -19,7 +30,16 @@ logger = logging.getLogger(__name__)
 
 
 class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
-    def __init__(self, conn: AsyncConnection[Any]) -> None:
+    def __init__(
+        self,
+        conn: AsyncConnection[Any],
+        *,
+        deserializer: Optional[
+            Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
+        ] = None,
+    ) -> None:
+        super().__init__(deserializer=deserializer)
+        self.conn = conn
         self.conn = conn
         self.loop = asyncio.get_running_loop()
 
@@ -87,7 +107,9 @@ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
             for idx, key in items:
                 row = key_to_row.get(key)
                 if row:
-                    results[idx] = _row_to_item(namespace, row)
+                    results[idx] = _row_to_item(
+                        namespace, row, loader=self._deserializer
+                    )
                 else:
                     results[idx] = None
 
@@ -106,16 +128,21 @@ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
         results: list[Result],
     ) -> None:
         queries = self._get_batch_search_queries(search_ops)
-        cursors: list[tuple[AsyncCursor[Any], int, SearchOp]] = []
+        cursors: list[tuple[AsyncCursor[Any], int]] = []
 
-        for (query, params), (idx, op) in zip(queries, search_ops):
+        for (query, params), (idx, _) in zip(queries, search_ops):
             cur = self.conn.cursor(binary=True)
             await cur.execute(query, params)
-            cursors.append((cur, idx, op))
+            cursors.append((cur, idx))
 
-        for cur, idx, op in cursors:
+        for cur, idx in cursors:
             rows = cast(list[Row], await cur.fetchall())
-            items = [_row_to_item(op.namespace_prefix, row) for row in rows]
+            items = [
+                _row_to_item(
+                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                )
+                for row in rows
+            ]
             results[idx] = items
 
     async def _batch_list_namespaces_ops(
@@ -132,9 +159,7 @@ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
 
         for cur, idx in cursors:
             rows = cast(list[dict], await cur.fetchall())
-            namespaces = [
-                tuple(row["truncated_prefix"].decode()[1:].split(".")) for row in rows
-            ]
+            namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
             results[idx] = namespaces
 
     @classmethod

{langgraph_checkpoint_postgres-1.0.10 → langgraph_checkpoint_postgres-2.0.0}/langgraph/store/postgres/base.py

@@ -4,7 +4,18 @@ import logging
 from collections import defaultdict
 from contextlib import contextmanager
 from datetime import datetime
-from typing import Any, Generic, Iterable, Iterator, Sequence, TypeVar, Union, cast
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    Iterator,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import orjson
 from psycopg import BaseConnection, Connection, Cursor
@@ -29,12 +40,9 @@ logger = logging.getLogger(__name__)
 
 MIGRATIONS = [
     """
-CREATE EXTENSION IF NOT EXISTS ltree;
-""",
-    """
 CREATE TABLE IF NOT EXISTS store (
     -- 'prefix' represents the doc's 'namespace'
-    prefix ltree NOT NULL,
+    prefix text NOT NULL,
     key text NOT NULL,
     value jsonb NOT NULL,
     created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
@@ -43,8 +51,8 @@ CREATE TABLE IF NOT EXISTS store (
 );
 """,
     """
--- For faster listing of namespaces & lookups by namespace with prefix/suffix matching
-CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING gist (prefix);
+-- For faster lookups by prefix
+CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
 """,
 ]
 
@@ -54,6 +62,17 @@ C = TypeVar("C", bound=BaseConnection)
 class BasePostgresStore(BaseStore, Generic[C]):
     MIGRATIONS = MIGRATIONS
     conn: C
+    __slots__ = ("_deserializer",)
+
+    def __init__(
+        self,
+        *,
+        deserializer: Optional[
+            Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
+        ] = None,
+    ) -> None:
+        super().__init__()
+        self._deserializer = deserializer
 
     def _get_batch_GET_ops_queries(
         self,
@@ -71,7 +90,7 @@ class BasePostgresStore(BaseStore, Generic[C]):
                 FROM store
                 WHERE prefix = %s AND key IN ({keys_to_query})
             """
-            params = (_namespace_to_ltree(namespace), *keys)
+            params = (_namespace_to_text(namespace), *keys)
             results.append((query, params, namespace, items))
         return results
 
@@ -98,7 +117,7 @@ class BasePostgresStore(BaseStore, Generic[C]):
                 query = (
                     f"DELETE FROM store WHERE prefix = %s AND key IN ({placeholders})"
                 )
-                params = (_namespace_to_ltree(namespace), *keys)
+                params = (_namespace_to_text(namespace), *keys)
                 queries.append((query, params))
         if inserts:
             values = []
@@ -107,7 +126,7 @@ class BasePostgresStore(BaseStore, Generic[C]):
                 values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
                 insertion_params.extend(
                     [
-                        _namespace_to_ltree(op.namespace),
+                        _namespace_to_text(op.namespace),
                         op.key,
                         Jsonb(op.value),
                     ]
@@ -130,11 +149,11 @@ class BasePostgresStore(BaseStore, Generic[C]):
         queries: list[tuple[str, Sequence]] = []
         for _, op in search_ops:
             query = """
-                SELECT key, value, created_at, updated_at
+                SELECT prefix, key, value, created_at, updated_at
                 FROM store
-                WHERE prefix <@ %s
+                WHERE prefix LIKE %s
             """
-            params: list = [_namespace_to_ltree(op.namespace_prefix)]
+            params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
 
             if op.filter:
                 filter_conditions = []
@@ -159,22 +178,39 @@ class BasePostgresStore(BaseStore, Generic[C]):
     ) -> list[tuple[str, Sequence]]:
         queries: list[tuple[str, Sequence]] = []
         for _, op in list_ops:
-            query = "SELECT DISTINCT subltree(prefix, 0, LEAST(nlevel(prefix), %s)) AS truncated_prefix FROM store"
-            # https://www.postgresql.org/docs/current/ltree.html
-            # The length of a label path cannot exceed 65535 labels.
-            params: list[Any] = [op.max_depth if op.max_depth is not None else 65536]
+            query = """
+                SELECT DISTINCT ON (truncated_prefix) truncated_prefix, prefix
+                FROM (
+                    SELECT
+                        prefix,
+                        CASE
+                            WHEN %s::integer IS NOT NULL THEN
+                                (SELECT STRING_AGG(part, '.' ORDER BY idx)
+                                 FROM (
+                                     SELECT part, ROW_NUMBER() OVER () AS idx
+                                     FROM UNNEST(REGEXP_SPLIT_TO_ARRAY(prefix, '\.')) AS part
+                                     LIMIT %s::integer
+                                 ) subquery
+                                )
+                            ELSE prefix
+                        END AS truncated_prefix
+                    FROM store
+            """
+            params: list[Any] = [op.max_depth, op.max_depth]
 
             conditions = []
             if op.match_conditions:
                 for condition in op.match_conditions:
                     if condition.match_type == "prefix":
-                        conditions.append("prefix ~ %s::lquery")
-                        lquery_pattern = f"{_namespace_to_ltree(condition.path)}.*"
-                        params.append(lquery_pattern)
+                        conditions.append("prefix LIKE %s")
+                        params.append(
+                            f"{_namespace_to_text(condition.path, handle_wildcards=True)}%"
+                        )
                     elif condition.match_type == "suffix":
-                        conditions.append("prefix ~ %s::lquery")
-                        lquery_pattern = f"*.{_namespace_to_ltree(condition.path)}"
-                        params.append(lquery_pattern)
+                        conditions.append("prefix LIKE %s")
+                        params.append(
+                            f"%{_namespace_to_text(condition.path, handle_wildcards=True)}"
+                        )
                     else:
                         logger.warning(
                             f"Unknown match_type in list_namespaces: {condition.match_type}"
@@ -182,16 +218,25 @@ class BasePostgresStore(BaseStore, Generic[C]):
 
             if conditions:
                 query += " WHERE " + " AND ".join(conditions)
+            query += ") AS subquery "
 
             query += " ORDER BY truncated_prefix LIMIT %s OFFSET %s"
             params.extend([op.limit, op.offset])
-
             queries.append((query, params))
+
         return queries
 
 
 class PostgresStore(BasePostgresStore[Connection]):
-    def __init__(self, conn: Connection[Any]) -> None:
+    def __init__(
+        self,
+        conn: Connection[Any],
+        *,
+        deserializer: Optional[
+            Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
+        ] = None,
+    ) -> None:
+        super().__init__(deserializer=deserializer)
         self.conn = conn
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:
@@ -246,7 +291,9 @@ class PostgresStore(BasePostgresStore[Connection]):
             for idx, key in items:
                 row = key_to_row.get(key)
                 if row:
-                    results[idx] = _row_to_item(namespace, row)
+                    results[idx] = _row_to_item(
+                        namespace, row, loader=self._deserializer
+                    )
                 else:
                     results[idx] = None
 
@@ -265,16 +312,21 @@ class PostgresStore(BasePostgresStore[Connection]):
         results: list[Result],
     ) -> None:
         queries = self._get_batch_search_queries(search_ops)
-        cursors: list[tuple[Cursor[Any], int, SearchOp]] = []
+        cursors: list[tuple[Cursor[Any], int]] = []
 
-        for (query, params), (idx, op) in zip(queries, search_ops):
+        for (query, params), (idx, _) in zip(queries, search_ops):
             cur = self.conn.cursor(binary=True)
             cur.execute(query, params)
-            cursors.append((cur, idx, op))
+            cursors.append((cur, idx))
 
-        for cur, idx, op in cursors:
+        for cur, idx in cursors:
             rows = cast(list[Row], cur.fetchall())
-            items = [_row_to_item(op.namespace_prefix, row) for row in rows]
+            items = [
+                _row_to_item(
+                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                )
+                for row in rows
+            ]
             results[idx] = items
 
     def _batch_list_namespaces_ops(
@@ -291,9 +343,7 @@ class PostgresStore(BasePostgresStore[Connection]):
 
         for cur, idx in cursors:
             rows = cast(list[dict], cur.fetchall())
-            namespaces = [
-                tuple(row["truncated_prefix"].decode()[1:].split(".")) for row in rows
-            ]
+            namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows]
             results[idx] = namespaces
 
     @classmethod
@@ -351,21 +401,31 @@ class PostgresStore(BasePostgresStore[Connection]):
 class Row(TypedDict):
     key: str
     value: Any
-    prefix: bytes
+    prefix: str
     created_at: datetime
     updated_at: datetime
 
 
-def _namespace_to_ltree(namespace: tuple[str, ...]) -> str:
-    """Convert namespace tuple to ltree-compatible string."""
+def _namespace_to_text(
+    namespace: tuple[str, ...], handle_wildcards: bool = False
+) -> str:
+    """Convert namespace tuple to text string."""
+    if handle_wildcards:
+        namespace = tuple("%" if val == "*" else val for val in namespace)
     return ".".join(namespace)
 
 
-def _row_to_item(namespace: tuple[str, ...], row: Row) -> Item:
+def _row_to_item(
+    namespace: tuple[str, ...],
+    row: Row,
+    *,
+    loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
+) -> Item:
     """Convert a row from the database into an Item."""
+    loader = loader or _json_loads
     val = row["value"]
     return Item(
-        value=val if isinstance(val, dict) else _json_loads(val),
+        value=val if isinstance(val, dict) else loader(val),
         key=row["key"],
         namespace=namespace,
         created_at=row["created_at"],
@@ -392,3 +452,11 @@ def _json_loads(content: Union[bytes, orjson.Fragment]) -> Any:
             else:
                 content = content.contents.encode()
     return orjson.loads(cast(bytes, content))
+
+
+def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
+    if isinstance(namespace, list):
+        return tuple(namespace)
+    if isinstance(namespace, bytes):
+        namespace = namespace.decode()[1:]
+    return tuple(namespace.split("."))

{langgraph_checkpoint_postgres-1.0.10 → langgraph_checkpoint_postgres-2.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.10"
+version = "2.0.0"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
 
 [tool.poetry.dependencies]
 python = "^3.9.0,<4.0"
-langgraph-checkpoint = "^1.0.11"
+langgraph-checkpoint = "^2.0.0"
 orjson = ">=3.10.1"
 psycopg = "^3.0.0"
 psycopg-pool = "^3.0.0"