Create index concurrently (#2659)

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -57,6 +57,17 @@
     PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
 );""",
     "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
+    """
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
+    """,
 ]
 
 SELECT_SQL = f"""

libs/checkpoint-postgres/langgraph/store/postgres/aio.py

@@ -6,7 +6,6 @@
 
 import orjson
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg_pool import AsyncConnectionPool
 
@@ -219,22 +218,19 @@ async def setup(self) -> None:
         """
 
         async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
-            try:
-                await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
-                row = await cur.fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                await cur.execute(
-                    f"""
-                    CREATE TABLE IF NOT EXISTS {table} (
-                        v INTEGER PRIMARY KEY
-                    )
-                """
+            await cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {table} (
+                    v INTEGER PRIMARY KEY
                 )
+            """
+            )
+            await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
+            row = cast(dict, await cur.fetchone())
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
             return version
 
         async with self._cursor() as cur:

libs/checkpoint-postgres/langgraph/store/postgres/base.py

@@ -21,7 +21,6 @@
 
 import orjson
 from psycopg import Capabilities, Connection, Cursor, Pipeline
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
@@ -73,7 +72,7 @@ class Migration(NamedTuple):
 """,
     """
 -- For faster lookups by prefix
-CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
+CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
 """,
 ]
 
@@ -107,7 +106,7 @@ class Migration(NamedTuple):
     ),
     Migration(
         """
-CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
+CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
     USING %(index_type)s (embedding %(ops)s)%(index_params)s;
 """,
         condition=lambda store: bool(
@@ -847,22 +846,19 @@ def setup(self) -> None:
         """
 
         def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
-            try:
-                cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
-                row = cast(dict, cur.fetchone())
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                cur.execute(
-                    f"""
-                    CREATE TABLE IF NOT EXISTS {table} (
-                        v INTEGER PRIMARY KEY
-                    )
-                """
+            cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {table} (
+                    v INTEGER PRIMARY KEY
                 )
+            """
+            )
+            cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
+            row = cast(dict, cur.fetchone())
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
             return version
 
         with self._cursor() as cur:

libs/checkpoint-postgres/pyproject.toml

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

libs/checkpoint-postgres/tests/test_async.py

@@ -63,9 +63,8 @@ async def _pipe_saver():
             prepare_threshold=0,
             row_factory=dict_row,
         ) as conn:
-            async with conn.pipeline() as pipe:
-                checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
-                await checkpointer.setup()
+            checkpointer = AsyncPostgresSaver(conn)
+            await checkpointer.setup()
             async with conn.pipeline() as pipe:
                 checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
                 yield checkpointer

libs/checkpoint-postgres/tests/test_sync.py

@@ -57,9 +57,8 @@ def _pipe_saver():
             prepare_threshold=0,
             row_factory=dict_row,
         ) as conn:
-            with conn.pipeline() as pipe:
-                checkpointer = PostgresSaver(conn, pipe=pipe)
-                checkpointer.setup()
+            checkpointer = PostgresSaver(conn)
+            checkpointer.setup()
             with conn.pipeline() as pipe:
                 checkpointer = PostgresSaver(conn, pipe=pipe)
                 yield checkpointer