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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: langgraph-checkpoint-postgres
3
- Version: 1.0.9
3
+ Version: 1.0.11
4
4
  Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
5
5
  Home-page: https://www.github.com/langchain-ai/langgraph
6
6
  License: MIT
@@ -17,9 +17,7 @@ from langgraph.checkpoint.base import (
17
17
  CheckpointTuple,
18
18
  get_checkpoint_id,
19
19
  )
20
- from langgraph.checkpoint.postgres.base import (
21
- BasePostgresSaver,
22
- )
20
+ from langgraph.checkpoint.postgres.base import BasePostgresSaver
23
21
  from langgraph.checkpoint.serde.base import SerializerProtocol
24
22
 
25
23
  Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
@@ -167,15 +165,17 @@ class PostgresSaver(BasePostgresSaver):
167
165
  value["pending_sends"],
168
166
  ),
169
167
  self._load_metadata(value["metadata"]),
170
- {
171
- "configurable": {
172
- "thread_id": value["thread_id"],
173
- "checkpoint_ns": value["checkpoint_ns"],
174
- "checkpoint_id": value["parent_checkpoint_id"],
168
+ (
169
+ {
170
+ "configurable": {
171
+ "thread_id": value["thread_id"],
172
+ "checkpoint_ns": value["checkpoint_ns"],
173
+ "checkpoint_id": value["parent_checkpoint_id"],
174
+ }
175
175
  }
176
- }
177
- if value["parent_checkpoint_id"]
178
- else None,
176
+ if value["parent_checkpoint_id"]
177
+ else None
178
+ ),
179
179
  self._load_writes(value["pending_writes"]),
180
180
  )
181
181
 
@@ -246,15 +246,17 @@ class PostgresSaver(BasePostgresSaver):
246
246
  value["pending_sends"],
247
247
  ),
248
248
  self._load_metadata(value["metadata"]),
249
- {
250
- "configurable": {
251
- "thread_id": thread_id,
252
- "checkpoint_ns": checkpoint_ns,
253
- "checkpoint_id": value["parent_checkpoint_id"],
249
+ (
250
+ {
251
+ "configurable": {
252
+ "thread_id": thread_id,
253
+ "checkpoint_ns": checkpoint_ns,
254
+ "checkpoint_id": value["parent_checkpoint_id"],
255
+ }
254
256
  }
255
- }
256
- if value["parent_checkpoint_id"]
257
- else None,
257
+ if value["parent_checkpoint_id"]
258
+ else None
259
+ ),
258
260
  self._load_writes(value["pending_writes"]),
259
261
  )
260
262
 
@@ -384,3 +386,6 @@ class PostgresSaver(BasePostgresSaver):
384
386
  else:
385
387
  with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
386
388
  yield cur
389
+
390
+
391
+ __all__ = ["PostgresSaver", "Conn"]
@@ -0,0 +1,4 @@
1
+ from langgraph.store.postgres.aio import AsyncPostgresStore
2
+ from langgraph.store.postgres.base import PostgresStore
3
+
4
+ __all__ = ["AsyncPostgresStore", "PostgresStore"]
@@ -0,0 +1,214 @@
1
+ import asyncio
2
+ import logging
3
+ from contextlib import asynccontextmanager
4
+ from typing import (
5
+ Any,
6
+ AsyncIterator,
7
+ Callable,
8
+ Iterable,
9
+ Optional,
10
+ Sequence,
11
+ Union,
12
+ cast,
13
+ )
14
+
15
+ import orjson
16
+ from psycopg import AsyncConnection, AsyncCursor
17
+ from psycopg.errors import UndefinedTable
18
+ from psycopg.rows import dict_row
19
+
20
+ from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
21
+ from langgraph.store.postgres.base import (
22
+ BasePostgresStore,
23
+ Row,
24
+ _group_ops,
25
+ _row_to_item,
26
+ )
27
+
28
+ logger = logging.getLogger(__name__)
29
+
30
+
31
+ class AsyncPostgresStore(BasePostgresStore[AsyncConnection]):
32
+ def __init__(
33
+ self,
34
+ conn: AsyncConnection[Any],
35
+ *,
36
+ deserializer: Optional[
37
+ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
38
+ ] = None,
39
+ ) -> None:
40
+ super().__init__(deserializer=deserializer)
41
+ self.conn = conn
42
+ self.conn = conn
43
+ self.loop = asyncio.get_running_loop()
44
+
45
+ async def abatch(self, ops: Iterable[Op]) -> list[Result]:
46
+ grouped_ops, num_ops = _group_ops(ops)
47
+ results: list[Result] = [None] * num_ops
48
+
49
+ async with self.conn.pipeline():
50
+ tasks = []
51
+
52
+ if GetOp in grouped_ops:
53
+ tasks.append(
54
+ self._batch_get_ops(
55
+ cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
56
+ )
57
+ )
58
+
59
+ if PutOp in grouped_ops:
60
+ tasks.append(
61
+ self._batch_put_ops(
62
+ cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
63
+ )
64
+ )
65
+
66
+ if SearchOp in grouped_ops:
67
+ tasks.append(
68
+ self._batch_search_ops(
69
+ cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
70
+ results,
71
+ )
72
+ )
73
+
74
+ if ListNamespacesOp in grouped_ops:
75
+ tasks.append(
76
+ self._batch_list_namespaces_ops(
77
+ cast(
78
+ Sequence[tuple[int, ListNamespacesOp]],
79
+ grouped_ops[ListNamespacesOp],
80
+ ),
81
+ results,
82
+ )
83
+ )
84
+
85
+ await asyncio.gather(*tasks)
86
+
87
+ return results
88
+
89
+ def batch(self, ops: Iterable[Op]) -> list[Result]:
90
+ return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
91
+
92
+ async def _batch_get_ops(
93
+ self,
94
+ get_ops: Sequence[tuple[int, GetOp]],
95
+ results: list[Result],
96
+ ) -> None:
97
+ cursors = []
98
+ for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
99
+ cur = self.conn.cursor(binary=True)
100
+ await cur.execute(query, params)
101
+ cursors.append((cur, namespace, items))
102
+
103
+ for cur, namespace, items in cursors:
104
+ rows = cast(list[Row], await cur.fetchall())
105
+ key_to_row = {row["key"]: row for row in rows}
106
+ for idx, key in items:
107
+ row = key_to_row.get(key)
108
+ if row:
109
+ results[idx] = _row_to_item(
110
+ namespace, row, loader=self._deserializer
111
+ )
112
+ else:
113
+ results[idx] = None
114
+
115
+ async def _batch_put_ops(
116
+ self,
117
+ put_ops: Sequence[tuple[int, PutOp]],
118
+ ) -> None:
119
+ queries = self._get_batch_PUT_queries(put_ops)
120
+ for query, params in queries:
121
+ cur = self.conn.cursor(binary=True)
122
+ await cur.execute(query, params)
123
+
124
+ async def _batch_search_ops(
125
+ self,
126
+ search_ops: Sequence[tuple[int, SearchOp]],
127
+ results: list[Result],
128
+ ) -> None:
129
+ queries = self._get_batch_search_queries(search_ops)
130
+ cursors: list[tuple[AsyncCursor[Any], int, SearchOp]] = []
131
+
132
+ for (query, params), (idx, op) in zip(queries, search_ops):
133
+ cur = self.conn.cursor(binary=True)
134
+ await cur.execute(query, params)
135
+ cursors.append((cur, idx, op))
136
+
137
+ for cur, idx, op in cursors:
138
+ rows = cast(list[Row], await cur.fetchall())
139
+ items = [
140
+ _row_to_item(op.namespace_prefix, row, loader=self._deserializer)
141
+ for row in rows
142
+ ]
143
+ results[idx] = items
144
+
145
+ async def _batch_list_namespaces_ops(
146
+ self,
147
+ list_ops: Sequence[tuple[int, ListNamespacesOp]],
148
+ results: list[Result],
149
+ ) -> None:
150
+ queries = self._get_batch_list_namespaces_queries(list_ops)
151
+ cursors: list[tuple[AsyncCursor[Any], int]] = []
152
+ for (query, params), (idx, _) in zip(queries, list_ops):
153
+ cur = self.conn.cursor(binary=True)
154
+ await cur.execute(query, params)
155
+ cursors.append((cur, idx))
156
+
157
+ for cur, idx in cursors:
158
+ rows = cast(list[dict], await cur.fetchall())
159
+ namespaces = [
160
+ tuple(row["truncated_prefix"].decode()[1:].split(".")) for row in rows
161
+ ]
162
+ results[idx] = namespaces
163
+
164
+ @classmethod
165
+ @asynccontextmanager
166
+ async def from_conn_string(
167
+ cls,
168
+ conn_string: str,
169
+ ) -> AsyncIterator["AsyncPostgresStore"]:
170
+ """Create a new AsyncPostgresStore instance from a connection string.
171
+
172
+ Args:
173
+ conn_string (str): The Postgres connection info string.
174
+
175
+ Returns:
176
+ AsyncPostgresStore: A new AsyncPostgresStore instance.
177
+ """
178
+ async with await AsyncConnection.connect(
179
+ conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
180
+ ) as conn:
181
+ yield cls(conn=conn)
182
+
183
+ async def setup(self) -> None:
184
+ """Set up the store database asynchronously.
185
+
186
+ This method creates the necessary tables in the Postgres database if they don't
187
+ already exist and runs database migrations. It MUST be called directly by the user
188
+ the first time the store is used.
189
+ """
190
+ async with self.conn.cursor() as cur:
191
+ try:
192
+ await cur.execute(
193
+ "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
194
+ )
195
+ row = cast(dict, await cur.fetchone())
196
+ if row is None:
197
+ version = -1
198
+ else:
199
+ version = row["v"]
200
+ except UndefinedTable:
201
+ version = -1
202
+ # Create store_migrations table if it doesn't exist
203
+ await cur.execute(
204
+ """
205
+ CREATE TABLE IF NOT EXISTS store_migrations (
206
+ v INTEGER PRIMARY KEY
207
+ )
208
+ """
209
+ )
210
+ for v, migration in enumerate(
211
+ self.MIGRATIONS[version + 1 :], start=version + 1
212
+ ):
213
+ await cur.execute(migration)
214
+ await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
@@ -0,0 +1,435 @@
1
+ import asyncio
2
+ import json
3
+ import logging
4
+ from collections import defaultdict
5
+ from contextlib import contextmanager
6
+ from datetime import datetime
7
+ from typing import (
8
+ Any,
9
+ Callable,
10
+ Generic,
11
+ Iterable,
12
+ Iterator,
13
+ Optional,
14
+ Sequence,
15
+ TypeVar,
16
+ Union,
17
+ cast,
18
+ )
19
+
20
+ import orjson
21
+ from psycopg import BaseConnection, Connection, Cursor
22
+ from psycopg.errors import UndefinedTable
23
+ from psycopg.rows import dict_row
24
+ from psycopg.types.json import Jsonb
25
+ from typing_extensions import TypedDict
26
+
27
+ from langgraph.store.base import (
28
+ BaseStore,
29
+ GetOp,
30
+ Item,
31
+ ListNamespacesOp,
32
+ Op,
33
+ PutOp,
34
+ Result,
35
+ SearchOp,
36
+ )
37
+
38
+ logger = logging.getLogger(__name__)
39
+
40
+
41
+ MIGRATIONS = [
42
+ """
43
+ CREATE EXTENSION IF NOT EXISTS ltree;
44
+ """,
45
+ """
46
+ CREATE TABLE IF NOT EXISTS store (
47
+ -- 'prefix' represents the doc's 'namespace'
48
+ prefix ltree NOT NULL,
49
+ key text NOT NULL,
50
+ value jsonb NOT NULL,
51
+ created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
52
+ updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
53
+ PRIMARY KEY (prefix, key)
54
+ );
55
+ """,
56
+ """
57
+ -- For faster listing of namespaces & lookups by namespace with prefix/suffix matching
58
+ CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING gist (prefix);
59
+ """,
60
+ ]
61
+
62
+ C = TypeVar("C", bound=BaseConnection)
63
+
64
+
65
+ class BasePostgresStore(BaseStore, Generic[C]):
66
+ MIGRATIONS = MIGRATIONS
67
+ conn: C
68
+ __slots__ = ("_deserializer",)
69
+
70
+ def __init__(
71
+ self,
72
+ *,
73
+ deserializer: Optional[
74
+ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
75
+ ] = None,
76
+ ) -> None:
77
+ super().__init__()
78
+ self._deserializer = deserializer
79
+
80
+ def _get_batch_GET_ops_queries(
81
+ self,
82
+ get_ops: Sequence[tuple[int, GetOp]],
83
+ ) -> list[tuple[str, tuple, tuple[str, ...], list]]:
84
+ namespace_groups = defaultdict(list)
85
+ for idx, op in get_ops:
86
+ namespace_groups[op.namespace].append((idx, op.key))
87
+ results = []
88
+ for namespace, items in namespace_groups.items():
89
+ _, keys = zip(*items)
90
+ keys_to_query = ",".join(["%s"] * len(keys))
91
+ query = f"""
92
+ SELECT key, value, created_at, updated_at
93
+ FROM store
94
+ WHERE prefix = %s AND key IN ({keys_to_query})
95
+ """
96
+ params = (_namespace_to_ltree(namespace), *keys)
97
+ results.append((query, params, namespace, items))
98
+ return results
99
+
100
+ def _get_batch_PUT_queries(
101
+ self,
102
+ put_ops: Sequence[tuple[int, PutOp]],
103
+ ) -> list[tuple[str, Sequence]]:
104
+ inserts: list[PutOp] = []
105
+ deletes: list[PutOp] = []
106
+ for _, op in put_ops:
107
+ if op.value is None:
108
+ deletes.append(op)
109
+ else:
110
+ inserts.append(op)
111
+
112
+ queries: list[tuple[str, Sequence]] = []
113
+
114
+ if deletes:
115
+ namespace_groups: dict[tuple[str, ...], list[str]] = defaultdict(list)
116
+ for op in deletes:
117
+ namespace_groups[op.namespace].append(op.key)
118
+ for namespace, keys in namespace_groups.items():
119
+ placeholders = ",".join(["%s"] * len(keys))
120
+ query = (
121
+ f"DELETE FROM store WHERE prefix = %s AND key IN ({placeholders})"
122
+ )
123
+ params = (_namespace_to_ltree(namespace), *keys)
124
+ queries.append((query, params))
125
+ if inserts:
126
+ values = []
127
+ insertion_params = []
128
+ for op in inserts:
129
+ values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
130
+ insertion_params.extend(
131
+ [
132
+ _namespace_to_ltree(op.namespace),
133
+ op.key,
134
+ Jsonb(op.value),
135
+ ]
136
+ )
137
+ values_str = ",".join(values)
138
+ query = f"""
139
+ INSERT INTO store (prefix, key, value, created_at, updated_at)
140
+ VALUES {values_str}
141
+ ON CONFLICT (prefix, key) DO UPDATE
142
+ SET value = EXCLUDED.value, updated_at = CURRENT_TIMESTAMP
143
+ """
144
+ queries.append((query, insertion_params))
145
+
146
+ return queries
147
+
148
+ def _get_batch_search_queries(
149
+ self,
150
+ search_ops: Sequence[tuple[int, SearchOp]],
151
+ ) -> list[tuple[str, Sequence]]:
152
+ queries: list[tuple[str, Sequence]] = []
153
+ for _, op in search_ops:
154
+ query = """
155
+ SELECT key, value, created_at, updated_at
156
+ FROM store
157
+ WHERE prefix <@ %s
158
+ """
159
+ params: list = [_namespace_to_ltree(op.namespace_prefix)]
160
+
161
+ if op.filter:
162
+ filter_conditions = []
163
+ for key, value in op.filter.items():
164
+ if isinstance(value, list):
165
+ filter_conditions.append("value->%s @> %s::jsonb")
166
+ params.extend([key, json.dumps(value)])
167
+ else:
168
+ filter_conditions.append("value->%s = %s::jsonb")
169
+ params.extend([key, json.dumps(value)])
170
+ query += " AND " + " AND ".join(filter_conditions)
171
+
172
+ query += " LIMIT %s OFFSET %s"
173
+ params.extend([op.limit, op.offset])
174
+
175
+ queries.append((query, params))
176
+ return queries
177
+
178
+ def _get_batch_list_namespaces_queries(
179
+ self,
180
+ list_ops: Sequence[tuple[int, ListNamespacesOp]],
181
+ ) -> list[tuple[str, Sequence]]:
182
+ queries: list[tuple[str, Sequence]] = []
183
+ for _, op in list_ops:
184
+ query = "SELECT DISTINCT subltree(prefix, 0, LEAST(nlevel(prefix), %s)) AS truncated_prefix FROM store"
185
+ # https://www.postgresql.org/docs/current/ltree.html
186
+ # The length of a label path cannot exceed 65535 labels.
187
+ params: list[Any] = [op.max_depth if op.max_depth is not None else 65536]
188
+
189
+ conditions = []
190
+ if op.match_conditions:
191
+ for condition in op.match_conditions:
192
+ if condition.match_type == "prefix":
193
+ conditions.append("prefix ~ %s::lquery")
194
+ lquery_pattern = f"{_namespace_to_ltree(condition.path)}.*"
195
+ params.append(lquery_pattern)
196
+ elif condition.match_type == "suffix":
197
+ conditions.append("prefix ~ %s::lquery")
198
+ lquery_pattern = f"*.{_namespace_to_ltree(condition.path)}"
199
+ params.append(lquery_pattern)
200
+ else:
201
+ logger.warning(
202
+ f"Unknown match_type in list_namespaces: {condition.match_type}"
203
+ )
204
+
205
+ if conditions:
206
+ query += " WHERE " + " AND ".join(conditions)
207
+
208
+ query += " ORDER BY truncated_prefix LIMIT %s OFFSET %s"
209
+ params.extend([op.limit, op.offset])
210
+
211
+ queries.append((query, params))
212
+ return queries
213
+
214
+
215
+ class PostgresStore(BasePostgresStore[Connection]):
216
+ def __init__(
217
+ self,
218
+ conn: Connection[Any],
219
+ *,
220
+ deserializer: Optional[
221
+ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
222
+ ] = None,
223
+ ) -> None:
224
+ super().__init__(deserializer=deserializer)
225
+ self.conn = conn
226
+
227
+ def batch(self, ops: Iterable[Op]) -> list[Result]:
228
+ grouped_ops, num_ops = _group_ops(ops)
229
+ results: list[Result] = [None] * num_ops
230
+
231
+ with self.conn.pipeline():
232
+ if GetOp in grouped_ops:
233
+ self._batch_get_ops(
234
+ cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results
235
+ )
236
+
237
+ if PutOp in grouped_ops:
238
+ self._batch_put_ops(
239
+ cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp])
240
+ )
241
+
242
+ if SearchOp in grouped_ops:
243
+ self._batch_search_ops(
244
+ cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
245
+ results,
246
+ )
247
+
248
+ if ListNamespacesOp in grouped_ops:
249
+ self._batch_list_namespaces_ops(
250
+ cast(
251
+ Sequence[tuple[int, ListNamespacesOp]],
252
+ grouped_ops[ListNamespacesOp],
253
+ ),
254
+ results,
255
+ )
256
+
257
+ return results
258
+
259
+ async def abatch(self, ops: Iterable[Op]) -> list[Result]:
260
+ return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops)
261
+
262
+ def _batch_get_ops(
263
+ self,
264
+ get_ops: Sequence[tuple[int, GetOp]],
265
+ results: list[Result],
266
+ ) -> None:
267
+ cursors = []
268
+ for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops):
269
+ cur = self.conn.cursor(binary=True)
270
+ cur.execute(query, params)
271
+ cursors.append((cur, namespace, items))
272
+
273
+ for cur, namespace, items in cursors:
274
+ rows = cast(list[Row], cur.fetchall())
275
+ key_to_row = {row["key"]: row for row in rows}
276
+ for idx, key in items:
277
+ row = key_to_row.get(key)
278
+ if row:
279
+ results[idx] = _row_to_item(
280
+ namespace, row, loader=self._deserializer
281
+ )
282
+ else:
283
+ results[idx] = None
284
+
285
+ def _batch_put_ops(
286
+ self,
287
+ put_ops: Sequence[tuple[int, PutOp]],
288
+ ) -> None:
289
+ queries = self._get_batch_PUT_queries(put_ops)
290
+ for query, params in queries:
291
+ cur = self.conn.cursor(binary=True)
292
+ cur.execute(query, params)
293
+
294
+ def _batch_search_ops(
295
+ self,
296
+ search_ops: Sequence[tuple[int, SearchOp]],
297
+ results: list[Result],
298
+ ) -> None:
299
+ queries = self._get_batch_search_queries(search_ops)
300
+ cursors: list[tuple[Cursor[Any], int, SearchOp]] = []
301
+
302
+ for (query, params), (idx, op) in zip(queries, search_ops):
303
+ cur = self.conn.cursor(binary=True)
304
+ cur.execute(query, params)
305
+ cursors.append((cur, idx, op))
306
+
307
+ for cur, idx, op in cursors:
308
+ rows = cast(list[Row], cur.fetchall())
309
+ items = [
310
+ _row_to_item(op.namespace_prefix, row, loader=self._deserializer)
311
+ for row in rows
312
+ ]
313
+ results[idx] = items
314
+
315
+ def _batch_list_namespaces_ops(
316
+ self,
317
+ list_ops: Sequence[tuple[int, ListNamespacesOp]],
318
+ results: list[Result],
319
+ ) -> None:
320
+ queries = self._get_batch_list_namespaces_queries(list_ops)
321
+ cursors: list[tuple[Cursor[Any], int]] = []
322
+ for (query, params), (idx, _) in zip(queries, list_ops):
323
+ cur = self.conn.cursor(binary=True)
324
+ cur.execute(query, params)
325
+ cursors.append((cur, idx))
326
+
327
+ for cur, idx in cursors:
328
+ rows = cast(list[dict], cur.fetchall())
329
+ namespaces = [
330
+ tuple(row["truncated_prefix"].decode()[1:].split(".")) for row in rows
331
+ ]
332
+ results[idx] = namespaces
333
+
334
+ @classmethod
335
+ @contextmanager
336
+ def from_conn_string(
337
+ cls,
338
+ conn_string: str,
339
+ ) -> Iterator["PostgresStore"]:
340
+ """Create a new BasePostgresStore instance from a connection string.
341
+
342
+ Args:
343
+ conn_string (str): The Postgres connection info string.
344
+
345
+ Returns:
346
+ BasePostgresStore: A new BasePostgresStore instance.
347
+ """
348
+ with Connection.connect(
349
+ conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
350
+ ) as conn:
351
+ yield cls(conn=conn)
352
+
353
+ def setup(self) -> None:
354
+ """Set up the store database.
355
+
356
+ This method creates the necessary tables in the Postgres database if they don't
357
+ already exist and runs database migrations. It MUST be called directly by the user
358
+ the first time the store is used.
359
+ """
360
+ with self.conn.cursor(binary=True) as cur:
361
+ try:
362
+ cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
363
+ row = cast(dict, cur.fetchone())
364
+ if row is None:
365
+ version = -1
366
+ else:
367
+ version = row["v"]
368
+ except UndefinedTable:
369
+ self.conn.rollback()
370
+ version = -1
371
+ # Create store_migrations table if it doesn't exist
372
+ cur.execute(
373
+ """
374
+ CREATE TABLE IF NOT EXISTS store_migrations (
375
+ v INTEGER PRIMARY KEY
376
+ )
377
+ """
378
+ )
379
+ for v, migration in enumerate(
380
+ self.MIGRATIONS[version + 1 :], start=version + 1
381
+ ):
382
+ cur.execute(migration)
383
+ cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
384
+
385
+
386
+ class Row(TypedDict):
387
+ key: str
388
+ value: Any
389
+ prefix: bytes
390
+ created_at: datetime
391
+ updated_at: datetime
392
+
393
+
394
+ def _namespace_to_ltree(namespace: tuple[str, ...]) -> str:
395
+ """Convert namespace tuple to ltree-compatible string."""
396
+ return ".".join(namespace)
397
+
398
+
399
+ def _row_to_item(
400
+ namespace: tuple[str, ...],
401
+ row: Row,
402
+ *,
403
+ loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
404
+ ) -> Item:
405
+ """Convert a row from the database into an Item."""
406
+ loader = loader or _json_loads
407
+ val = row["value"]
408
+ return Item(
409
+ value=val if isinstance(val, dict) else loader(val),
410
+ key=row["key"],
411
+ namespace=namespace,
412
+ created_at=row["created_at"],
413
+ updated_at=row["updated_at"],
414
+ )
415
+
416
+
417
+ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
418
+ grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
419
+ tot = 0
420
+ for idx, op in enumerate(ops):
421
+ grouped_ops[type(op)].append((idx, op))
422
+ tot += 1
423
+ return grouped_ops, tot
424
+
425
+
426
+ def _json_loads(content: Union[bytes, orjson.Fragment]) -> Any:
427
+ if isinstance(content, orjson.Fragment):
428
+ if hasattr(content, "buf"):
429
+ content = content.buf
430
+ else:
431
+ if isinstance(content.contents, bytes):
432
+ content = content.contents
433
+ else:
434
+ content = content.contents.encode()
435
+ return orjson.loads(cast(bytes, content))
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "langgraph-checkpoint-postgres"
3
- version = "1.0.9"
3
+ version = "1.0.11"
4
4
  description = "Library with a Postgres implementation of LangGraph checkpoint saver."
5
5
  authors = []
6
6
  license = "MIT"