Firestore to Mongodb

modules/training_content/service.py

@@ -1,9 +1,12 @@
 import json
+import uuid
 from datetime import datetime
 from logging import getLogger
 
 from typing import Dict, List
 
+from pymongo.database import Database
+
 from modules.training_content.dtos import TrainingContentDTO, WeakAreaDTO, QueryDTO, DetailsDTO, TipsDTO
 
 
@@ -19,9 +22,9 @@ class TrainingContentService:
     ]
     # strategy word_link ct_focus reading_skill word_partners writing_skill language_for_writing
 
-    def __init__(self, kb, openai, firestore):
+    def __init__(self, kb, openai, mongo: Database):
         self._training_content_module = kb
-        self._db = firestore
+        self._db: Database = mongo
         self._logger = getLogger(__name__)
         self._llm = openai
 
@@ -37,16 +40,18 @@ class TrainingContentService:
         for area in training_content.weak_areas:
             weak_areas["weak_areas"].append(area.dict())
 
+        new_id = uuid.uuid4()
         training_doc = {
+            'id': new_id,
             'created_at': int(datetime.now().timestamp() * 1000),
             **exam_map,
             **usefull_tips.dict(),
             **weak_areas,
             "user": user
         }
-        doc_ref = self._db.collection('training').add(training_doc)
+        self._db.training.insert_one(training_doc)
         return {
-            "id": doc_ref[1].id
+            "id": new_id
         }
 
     @staticmethod
@@ -400,10 +405,5 @@ class TrainingContentService:
         return result
 
     def _get_doc_by_id(self, collection: str, doc_id: str):
-        collection_ref = self._db.collection(collection)
-        doc_ref = collection_ref.document(doc_id)
-        doc = doc_ref.get()
-
-        if doc.exists:
-            return doc.to_dict()
-        return None
+        doc = self._db[collection].find_one({"id": doc_id})
+        return doc