refactor: MongoDB service

steps/infrastructure/fetch_from_mongodb.py

@@ -1,22 +1,11 @@
-from loguru import logger
 from zenml.steps import step
 
 from second_brain.infrastructure.mongo import MongoDBService
 
 
 @step
 def fetch_from_mongodb(limit: int, collection_name: str) -> list[dict]:
-    try:
-        service = MongoDBService(collection_name=collection_name)
+    with MongoDBService(collection_name=collection_name) as service:
+        documents = service.fetch_documents(limit, query={})
 
-        # Fetch documents
-        query = {}
-        documents = service.fetch_documents(limit, query)
-
-        # Log genre-specific counts
-        # service.verify_genre(genre)
-
-        return documents
-    except Exception as e:
-        logger.error(f"Failed to fetch documents: {e}")
-        raise
+    return documents

steps/infrastructure/ingest_to_mongodb.py

@@ -14,13 +14,12 @@ def ingest_to_mongodb(documents: list[dict], collection_name: str) -> None:
     Raises:
         Exception: If the ingestion process fails.
     """
-    try:
-        service = MongoDBService(collection_name=collection_name)
 
+    with MongoDBService(collection_name=collection_name) as service:
         service.clear_collection()
         service.ingest_documents(documents)
 
-        service.get_collection_count()
-    except Exception as e:
-        logger.error(f"Failed to ingest documents: {e}")
-        raise
+        count = service.get_collection_count()
+        logger.info(
+            f"Successfully ingested {count} documents into MongoDB collection '{collection_name}'"
+        )