feat: add delete_all_documents() to MongoDB DocumentStore (#2401)

* feat: add delete_all_documents() to MongoDB DocumentStore

* add functionality to recreate a collection

---------

Co-authored-by: Amna Mubashar <amnahkhan.ak@gmail.com>
Co-authored-by: David S. Batista <dsbatista@gmail.com>

Author: 3 people

integrations/mongodb_atlas/src/haystack_integrations/document_stores/mongodb_atlas/document_store.py

@@ -423,6 +423,108 @@ async def delete_documents_async(self, document_ids: List[str]) -> None:
             return
         await self._collection_async.delete_many(filter={"id": {"$in": document_ids}})
 
+    def delete_all_documents(self, *, recreate_collection: bool = False) -> None:
+        """
+        Deletes all documents in the document store.
+
+        :param recreate_collection: If True, the collection will be dropped and recreated with the original
+            configuration and indexes. If False, all documents will be deleted while preserving the collection.
+            Recreating the collection is faster for very large collections.
+        """
+        self._ensure_connection_setup()
+        assert self._collection is not None
+        assert self._connection is not None
+
+        try:
+            if recreate_collection:
+                database = self._connection[self.database_name]
+
+                # Save collection configuration
+                collection_info = database.list_collections(filter={"name": self.collection_name})
+                config = next(collection_info, {}).get("options", {})
+
+                # Save index definitions (excluding default _id index)
+                indexes = list(self._collection.list_indexes())
+                custom_indexes = [idx for idx in indexes if idx["name"] != "_id_"]
+
+                # Drop and recreate collection
+                self._collection.drop()
+                database.create_collection(self.collection_name, **config)
+
+                # Recreate indexes
+                for idx in custom_indexes:
+                    keys = list(idx["key"].items())
+                    index_options = {k: v for k, v in idx.items() if k not in ["key", "v", "ns"]}
+                    self._collection.create_index(keys, **index_options)
+
+                logger.info(
+                    "Collection '{collection}' recreated with original configuration.",
+                    collection=self.collection_name,
+                )
+            else:
+                # Delete all documents without recreating collection
+                result = self._collection.delete_many({})
+                logger.info(
+                    "Deleted {n_docs} documents from collection '{collection}'.",
+                    n_docs=result.deleted_count,
+                    collection=self.collection_name,
+                )
+        except Exception as e:
+            msg = f"Failed to delete all documents from MongoDB Atlas: {e!s}"
+            raise DocumentStoreError(msg) from e
+
+    async def delete_all_documents_async(self, *, recreate_collection: bool = False) -> None:
+        """
+        Asynchronously deletes all documents in the document store.
+
+        :param recreate_collection: If True, the collection will be dropped and recreated with the original
+            configuration and indexes. If False, all documents will be deleted while preserving the collection.
+            Recreating the collection is faster for very large collections.
+        """
+        await self._ensure_connection_setup_async()
+        assert self._collection_async is not None
+        assert self._connection_async is not None
+
+        try:
+            if recreate_collection:
+                database = self._connection_async[self.database_name]
+
+                # Save collection configuration
+                collection_info_cursor = await database.list_collections(filter={"name": self.collection_name})
+                config_list = await collection_info_cursor.to_list(length=1)
+                config = config_list[0].get("options", {}) if config_list else {}
+
+                # Save index definitions (excluding default _id index)
+                indexes_cursor = await self._collection_async.list_indexes()
+                indexes = await indexes_cursor.to_list(length=None)
+                custom_indexes = [idx for idx in indexes if idx["name"] != "_id_"]
+
+                # Drop and recreate collection
+                await self._collection_async.drop()
+                await database.create_collection(self.collection_name, **config)
+
+                # Recreate indexes
+                for idx in custom_indexes:
+                    keys = list(idx["key"].items())
+                    index_options = {k: v for k, v in idx.items() if k not in ["key", "v", "ns"]}
+                    await self._collection_async.create_index(keys, **index_options)
+
+                logger.info(
+                    "Collection '{collection}' recreated with original configuration.",
+                    collection=self.collection_name,
+                )
+            else:
+                # Delete all documents without recreating collection
+                result = await self._collection_async.delete_many({})
+                logger.info(
+                    "Deleted {n_docs} documents from collection '{collection}'.",
+                    n_docs=result.deleted_count,
+                    collection=self.collection_name,
+                )
+        except Exception as e:
+            msg = f"Failed to delete all documents from MongoDB Atlas: {e!s}"
+            raise DocumentStoreError(msg) from e
+
     def _embedding_retrieval(
         self,
         query_embedding: List[float],

integrations/mongodb_atlas/tests/test_document_store.py

@@ -333,3 +333,29 @@ def test_custom_content_field(self):
 
         finally:
             database[collection_name].drop()
+
+    def test_delete_all_documents(self, document_store: MongoDBAtlasDocumentStore):
+        docs = [Document(id="1", content="first doc"), Document(id="2", content="second doc")]
+        document_store.write_documents(docs)
+        assert document_store.count_documents() == 2
+        document_store.delete_all_documents()
+        assert document_store.count_documents() == 0
+
+    def test_delete_all_documents_empty_collection(self, document_store: MongoDBAtlasDocumentStore):
+        assert document_store.count_documents() == 0
+        document_store.delete_all_documents()
+        assert document_store.count_documents() == 0
+
+    def test_delete_all_documents_with_recreate_collection(self, document_store: MongoDBAtlasDocumentStore):
+        docs = [Document(id="1", content="first doc"), Document(id="2", content="second doc")]
+        document_store.write_documents(docs)
+        assert document_store.count_documents() == 2
+
+        # Delete all documents with collection recreation
+        document_store.delete_all_documents(recreate_collection=True)
+        assert document_store.count_documents() == 0
+
+        # Verify collection still exists and we can write to it
+        new_docs = [Document(id="3", content="third doc")]
+        document_store.write_documents(new_docs)
+        assert document_store.count_documents() == 1

integrations/mongodb_atlas/tests/test_document_store_async.py

@@ -127,3 +127,29 @@ async def test_delete_documents_async(self, document_store: MongoDBAtlasDocument
         assert await document_store.count_documents_async() == 1
         await document_store.delete_documents_async(document_ids=["1"])
         assert await document_store.count_documents_async() == 0
+
+    async def test_delete_all_documents_async(self, document_store: MongoDBAtlasDocumentStore):
+        docs = [Document(id="1", content="first doc"), Document(id="2", content="second doc")]
+        await document_store.write_documents_async(docs)
+        assert await document_store.count_documents_async() == 2
+        await document_store.delete_all_documents_async()
+        assert await document_store.count_documents_async() == 0
+
+    async def test_delete_all_documents_async_empty_collection(self, document_store: MongoDBAtlasDocumentStore):
+        assert await document_store.count_documents_async() == 0
+        await document_store.delete_all_documents_async()
+        assert await document_store.count_documents_async() == 0
+
+    async def test_delete_all_documents_async_with_recreate_collection(self, document_store: MongoDBAtlasDocumentStore):
+        docs = [Document(id="1", content="first doc"), Document(id="2", content="second doc")]
+        await document_store.write_documents_async(docs)
+        assert await document_store.count_documents_async() == 2
+
+        # Delete all documents with collection recreation
+        await document_store.delete_all_documents_async(recreate_collection=True)
+        assert await document_store.count_documents_async() == 0
+
+        # Verify collection still exists and we can write to it
+        new_docs = [Document(id="3", content="third doc")]
+        await document_store.write_documents_async(new_docs)
+        assert await document_store.count_documents_async() == 1