langgraph-checkpoint-postgres 1.0.11__tar.gz → 2.0.1__tar.gz

{langgraph_checkpoint_postgres-1.0.11 → langgraph_checkpoint_postgres-2.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.11
+Version: 2.0.1
 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.11 → langgraph_checkpoint_postgres-2.0.1}/langgraph/store/postgres/aio.py

@@ -18,9 +18,11 @@ from psycopg.errors import UndefinedTable
 from psycopg.rows import dict_row
 
 from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
+from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.postgres.base import (
     BasePostgresStore,
     Row,
+    _decode_ns_bytes,
     _group_ops,
     _row_to_item,
 )
@@ -28,7 +30,9 @@ from langgraph.store.postgres.base import (
 logger = logging.getLogger(__name__)
 
 
-class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
+class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnection]):
+    __slots__ = ("_deserializer",)
+
     def __init__(
         self,
         conn: AsyncConnection[Any],
@@ -37,7 +41,8 @@ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
     ) -> None:
-        super().__init__(deserializer=deserializer)
+        super().__init__()
+        self._deserializer = deserializer
         self.conn = conn
         self.conn = conn
         self.loop = asyncio.get_running_loop()
@@ -127,17 +132,19 @@ 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, loader=self._deserializer)
+                _row_to_item(
+                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                )
                 for row in rows
             ]
             results[idx] = items
@@ -156,9 +163,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.11 → langgraph_checkpoint_postgres-2.0.1}/langgraph/store/postgres/base.py

@@ -40,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,
@@ -54,28 +51,18 @@ 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);
 """,
 ]
 
 C = TypeVar("C", bound=BaseConnection)
 
 
-class BasePostgresStore(BaseStore, Generic[C]):
+class BasePostgresStore(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
+    _deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]]
 
     def _get_batch_GET_ops_queries(
         self,
@@ -93,7 +80,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
 
@@ -120,7 +107,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 = []
@@ -129,7 +116,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),
                     ]
@@ -152,11 +139,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 = []
@@ -169,7 +156,9 @@ class BasePostgresStore(BaseStore, Generic[C]):
                         params.extend([key, json.dumps(value)])
                 query += " AND " + " AND ".join(filter_conditions)
 
-            query += " LIMIT %s OFFSET %s"
+            # Note: we will need to not do this if sim/keyword search
+            # is used
+            query += " ORDER BY updated_at DESC LIMIT %s OFFSET %s"
             params.extend([op.limit, op.offset])
 
             queries.append((query, params))
@@ -181,22 +170,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}"
@@ -204,15 +210,18 @@ 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]):
+class PostgresStore(BaseStore, BasePostgresStore[Connection]):
+    __slots__ = ("_deserializer",)
+
     def __init__(
         self,
         conn: Connection[Any],
@@ -221,7 +230,8 @@ class PostgresStore(BasePostgresStore[Connection]):
             Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
         ] = None,
     ) -> None:
-        super().__init__(deserializer=deserializer)
+        super().__init__()
+        self._deserializer = deserializer
         self.conn = conn
 
     def batch(self, ops: Iterable[Op]) -> list[Result]:
@@ -297,17 +307,19 @@ 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, loader=self._deserializer)
+                _row_to_item(
+                    _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
+                )
                 for row in rows
             ]
             results[idx] = items
@@ -326,9 +338,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
@@ -386,13 +396,17 @@ 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)
 
 
@@ -433,3 +447,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.11 → langgraph_checkpoint_postgres-2.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.11"
+version = "2.0.1"
 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"