INTPYTHON-886 - Fix LangGraph checkpointer index creation logic (langchain-ai#313)

INTPYTHON-886

Author: NoahStapp

libs/langgraph-checkpoint-mongodb/langgraph/checkpoint/mongodb/saver.py

@@ -20,11 +20,49 @@
 from langgraph.checkpoint.serde.base import SerializerProtocol
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from pymongo import ASCENDING, MongoClient, UpdateOne
+from pymongo.collection import Collection
 from pymongo.database import Database as MongoDatabase
 
 from .utils import DRIVER_METADATA, dumps_metadata, loads_metadata
 
 
+def _create_saver_indexes(
+    collection: Collection,
+    compound_index: list[tuple[str, int]],
+    ttl: Optional[int] = None,
+) -> None:
+    """Create indexes for the saver collections.
+
+    This helper function creates the given compound index and TTL index (if required)
+    for the given collection.
+
+    Args:
+        collection (Collection): The MongoDB collection to create indexes on.
+        compound_index (list[tuple[str, int]]): The compound index to create.
+        ttl (int, optional): Time to live in seconds for the TTL index. Defaults to None.
+    """
+
+    def index_key_list(index: Any) -> list[tuple[str, int]]:
+        return list((k, v) for k, v in index["key"].items())
+
+    indexes = list(collection.list_indexes())
+    index_keys = [index_key_list(idx) for idx in indexes]
+    if compound_index not in index_keys:
+        collection.create_index(compound_index, unique=True)
+    if ttl is not None:
+        ttl_index = [("created_at", ASCENDING)]
+        found = False
+        for idx in indexes:
+            if (
+                index_key_list(idx) == tuple(ttl_index)
+                and idx.get("expireAfterSeconds") == ttl
+            ):
+                found = True
+                break
+        if not found:
+            collection.create_index(ttl_index, expireAfterSeconds=ttl)
+
+
 class MongoDBSaver(BaseCheckpointSaver):
     """A checkpointer that stores StateGraph checkpoints in a MongoDB database.
 
@@ -97,34 +135,22 @@ def __init__(
         else:
             self.serde = JsonPlusSerializer()
 
-        # Create indexes if not present
-        if len(self.checkpoint_collection.list_indexes().to_list()) < 2:
-            self.checkpoint_collection.create_index(
-                keys=[("thread_id", 1), ("checkpoint_ns", 1), ("checkpoint_id", -1)],
-                unique=True,
-            )
-            if self.ttl:
-                self.checkpoint_collection.create_index(
-                    keys=[("created_at", ASCENDING)],
-                    expireAfterSeconds=self.ttl,
-                )
-
-        if len(self.writes_collection.list_indexes().to_list()) < 2:
-            self.writes_collection.create_index(
-                keys=[
-                    ("thread_id", 1),
-                    ("checkpoint_ns", 1),
-                    ("checkpoint_id", -1),
-                    ("task_id", 1),
-                    ("idx", 1),
-                ],
-                unique=True,
-            )
-            if self.ttl:
-                self.writes_collection.create_index(
-                    keys=[("created_at", ASCENDING)],
-                    expireAfterSeconds=self.ttl,
-                )
+        _create_saver_indexes(
+            self.checkpoint_collection,
+            [("thread_id", 1), ("checkpoint_ns", 1), ("checkpoint_id", -1)],
+            self.ttl,
+        )
+        _create_saver_indexes(
+            self.writes_collection,
+            [
+                ("thread_id", 1),
+                ("checkpoint_ns", 1),
+                ("checkpoint_id", -1),
+                ("task_id", 1),
+                ("idx", 1),
+            ],
+            self.ttl,
+        )
 
     @classmethod
     @contextmanager

libs/langgraph-checkpoint-mongodb/tests/unit_tests/test_sync.py

@@ -209,3 +209,47 @@ def test_ttl(input_data: dict[str, Any]) -> None:
             saver.checkpoint_collection.drop_indexes()
             saver.writes_collection.delete_many({})
             saver.writes_collection.drop_indexes()
+
+
+def test_init_creates_indexes() -> None:
+    client: MongoClient = MongoClient(MONGODB_URI)
+    db = client[DB_NAME]
+    checkpoint_coll = "checkpoints_test"
+    writes_coll = "writes_test"
+
+    db.drop_collection(checkpoint_coll)
+    db.drop_collection(writes_coll)
+
+    ttl = 100
+    with MongoDBSaver.from_conn_string(
+        MONGODB_URI, DB_NAME, checkpoint_coll, writes_coll, ttl=ttl
+    ) as saver:
+        cp_indexes = saver.checkpoint_collection.index_information()
+        wr_indexes = saver.writes_collection.index_information()
+
+        def _has_index(index_info: Any, keys: list[tuple[str, int]]) -> bool:
+            for _, info in index_info.items():
+                if info.get("key") == keys:
+                    return True
+            return False
+
+        expected_cp_keys = [
+            ("thread_id", 1),
+            ("checkpoint_ns", 1),
+            ("checkpoint_id", -1),
+        ]
+        assert _has_index(cp_indexes, expected_cp_keys)
+        assert _has_index(cp_indexes, [("created_at", 1)])
+
+        expected_wr_keys = [
+            ("thread_id", 1),
+            ("checkpoint_ns", 1),
+            ("checkpoint_id", -1),
+            ("task_id", 1),
+            ("idx", 1),
+        ]
+        assert _has_index(wr_indexes, expected_wr_keys)
+        assert _has_index(wr_indexes, [("created_at", 1)])
+
+    db.drop_collection(checkpoint_coll)
+    db.drop_collection(writes_coll)