From 697f878e3619a145e8ad5e4ee9342ac9b1bb3311 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Mon, 17 Mar 2025 12:38:43 -0700
Subject: [PATCH 1/2] Make expires_at idempotent

---
 libs/checkpoint-postgres/langgraph/store/postgres/base.py | 6 +++---
 libs/checkpoint-postgres/pyproject.toml                   | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py
index d6035bb70..338f5e890 100644
--- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py
+++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py
@@ -79,12 +79,12 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (p
     """
 -- Add expires_at column to store table
 ALTER TABLE store
-ADD COLUMN expires_at TIMESTAMP WITH TIME ZONE,
-ADD COLUMN ttl_minutes INT;
+ADD COLUMN IF NOT EXISTS expires_at TIMESTAMP WITH TIME ZONE,
+ADD COLUMN IF NOT EXISTS ttl_minutes INT;
 """,
     """
 -- Add indexes for efficient TTL sweeping
-CREATE INDEX idx_store_expires_at ON store (expires_at)
+CREATE INDEX IF NOT EXISTS idx_store_expires_at ON store (expires_at)
 WHERE expires_at IS NOT NULL;
 """,
 ]
diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml
index 0d9ac1c12..7c727163a 100644
--- a/libs/checkpoint-postgres/pyproject.toml
+++ b/libs/checkpoint-postgres/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.17"
+version = "2.0.18"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"

From 2a71180c1d1b578fa9777043fdcb30549e522ff6 Mon Sep 17 00:00:00 2001
From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>
Date: Mon, 17 Mar 2025 12:43:21 -0700
Subject: [PATCH 2/2] Add tests for idempotency in migraionts

---
 libs/checkpoint-postgres/tests/test_async_store.py | 4 ++++
 libs/checkpoint-postgres/tests/test_store.py       | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py
index 09502403d..94dc1fa03 100644
--- a/libs/checkpoint-postgres/tests/test_async_store.py
+++ b/libs/checkpoint-postgres/tests/test_async_store.py
@@ -69,6 +69,10 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
                 for mig in store.MIGRATIONS
             ]
             await store.setup()
+            async with store._cursor() as cur:
+                # drop the migration index
+                await cur.execute("DROP TABLE IF EXISTS store_migrations")
+            await store.setup()  # Will fail if migrations aren't idempotent
 
         if request.param == "pipe":
             async with AsyncPostgresStore.from_conn_string(
diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py
index 4ee37484e..a1360174a 100644
--- a/libs/checkpoint-postgres/tests/test_store.py
+++ b/libs/checkpoint-postgres/tests/test_store.py
@@ -415,6 +415,10 @@ def _create_vector_store(
             ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
         ) as store:
             store.setup()
+            with store._cursor() as cur:
+                # drop the migration index
+                cur.execute("DROP TABLE IF EXISTS store_migrations")
+            store.setup()  # Will fail if migrations aren't idempotent
             yield store
     finally:
         with Connection.connect(admin_conn_string, autocommit=True) as conn: