Fastapi refactor update

app/repositories/impl/document_stores/mongo.py

@@ -1,36 +1,37 @@
-"""import logging
-from pymongo import MongoClient
-
-from app.repositories.abc import IDocumentStore
-
-
-class MongoDB(IDocumentStore):
-
-    def __init__(self, client: MongoClient):
-        self._client = client
-        self._logger = logging.getLogger(__name__)
-
-    def save_to_db(self, collection: str, item):
-        collection_ref = self._client[collection]
-        result = collection_ref.insert_one(item)
-        if result.inserted_id:
-            self._logger.info(f"Document added with ID: {result.inserted_id}")
-            return True, str(result.inserted_id)
-        else:
-            return False, None
-
-    def save_to_db_with_id(self, collection: str, item, doc_id: str):
-        collection_ref = self._client[collection]
-        item['_id'] = doc_id
-        result = collection_ref.replace_one({'_id': id}, item, upsert=True)
-        if result.upserted_id or result.matched_count:
-            self._logger.info(f"Document added with ID: {doc_id}")
-            return True, doc_id
-        else:
-            return False, None
-
-    def get_all(self, collection: str):
-        collection_ref = self._client[collection]
-        all_documents = list(collection_ref.find())
-        return all_documents
-"""
+import logging
+import uuid
+from typing import Optional, List, Dict
+
+from motor.motor_asyncio import AsyncIOMotorDatabase
+
+from app.repositories.abc import IDocumentStore
+
+
+class MongoDB(IDocumentStore):
+
+    def __init__(self, mongo_db: AsyncIOMotorDatabase):
+        self._mongo_db = mongo_db
+        self._logger = logging.getLogger(__name__)
+
+    async def save_to_db(self, collection: str, item, doc_id: Optional[str] = None) -> Optional[str]:
+        collection_ref = self._mongo_db[collection]
+
+        if doc_id is None:
+            doc_id = str(uuid.uuid4())
+
+        item['id'] = doc_id
+
+        result = await collection_ref.insert_one(item)
+        if result.inserted_id:
+            # returning id instead of _id
+            self._logger.info(f"Document added with ID: {doc_id}")
+            return doc_id
+
+        return None
+
+    async def get_all(self, collection: str) -> List[Dict]:
+        cursor = self._mongo_db[collection].find()
+        return [document async for document in cursor]
+
+    async def get_doc_by_id(self, collection: str, doc_id: str) -> Optional[Dict]:
+        return await self._mongo_db[collection].find_one({"id": doc_id})