Firestore to Mongodb

modules/batch_users/service.py

@@ -9,8 +9,7 @@ import pandas as pd
 from typing import Dict
 
 import shortuuid
-from google.cloud.firestore_v1 import Client
-from google.cloud.firestore_v1.base_query import FieldFilter
+from pymongo.database import Database
 
 from modules.batch_users.batch_users import BatchUsersDTO, UserDTO
 from modules.helper.file_helper import FileHelper
@@ -32,8 +31,8 @@ class BatchUsers:
         "speaking": 0,
     }
 
-    def __init__(self, firestore: Client):
-        self._db = firestore
+    def __init__(self, mongo: Database):
+        self._db: Database = mongo
         self._logger = getLogger(__name__)
 
     def batch_users(self, request_data: Dict):
@@ -141,7 +140,7 @@ class BatchUsers:
     def _insert_new_user(self, user: UserDTO):
         new_user = {
             **user.dict(exclude={
-                'id', 'passport_id', 'groupName', 'expiryDate',
+                'passport_id', 'groupName', 'expiryDate',
                 'corporate', 'passwordHash', 'passwordSalt'
             }),
             'bio': "",
@@ -155,11 +154,12 @@ class BatchUsers:
             'registrationDate': datetime.now(),
             'subscriptionExpirationDate': user.expiryDate
         }
-        self._db.collection('users').document(str(user.id)).set(new_user)
+        self._db.users.insert_one(new_user)
 
     def _create_code(self, user: UserDTO, maker_id: str) -> str:
         code = shortuuid.ShortUUID().random(length=6)
-        self._db.collection('codes').document(code).set({
+        self._db.codes.insert_one({
+            'id': code,
             'code': code,
             'creator': maker_id,
             'expiryDate': user.expiryDate,
@@ -198,31 +198,36 @@ class BatchUsers:
             }
         ]
         for group in default_groups:
-            self._db.collection('groups').document(group['id']).set(group)
+            self._db.groups.insert_one(group)
 
     def _assign_corporate_to_user(self, user: UserDTO, code: str):
         user_id = str(user.id)
-        corporate_users = self._db.collection('users').where(
-            filter=FieldFilter('email', '==', user.corporate)
-        ).limit(1).get()
-        if len(corporate_users) > 0:
-            corporate_user = corporate_users[0]
-            self._db.collection('codes').document(code).set({'creator': corporate_user.id}, merge=True)
-
+        corporate_user = self._db.users.find_one(
+            {"email": user.corporate}
+        )
+        if corporate_user:
+            self._db.codes.update_one(
+                {"id": code},
+                {"$set": {"creator": corporate_user.id}},
+                upsert=True
+            )
             group_type = "Students" if user.type == "student" else "Teachers"
 
-            groups = self._db.collection('groups').where(
-                filter=FieldFilter('admin', '==', corporate_user.id)
-            ).where(
-                filter=FieldFilter('name', '==', group_type)
-            ).limit(1).get()
+            group = self._db.groups.find_one(
+                {
+                    "admin": corporate_user.id,
+                    "name": group_type
+                }
+            )
 
-            if len(groups) > 0:
-                group = groups[0]
-                participants = group.get('participants')
+            if group:
+                participants = group['participants']
                 if user_id not in participants:
                     participants.append(user_id)
-                    group.reference.update({'participants': participants})
+                    self._db.groups.update_one(
+                        {"id": group.id},
+                        {"$set": {"participants": participants}}
+                    )
                     
             else:
                 group = {
@@ -233,18 +238,19 @@ class BatchUsers:
                     'disableEditing': True,
                 }
                 
-                self._db.collection('groups').document(group['id']).set(group)
+                self._db.groups.insert_one(group)
 
     def _assign_user_to_group_by_name(self, user: UserDTO, maker_id: str):
         user_id = str(user.id)
 
-        groups = self._db.collection('groups').where(
-            filter=FieldFilter('admin', '==', maker_id)
-        ).where(
-            filter=FieldFilter('name', '==', user.groupName.strip())
-        ).limit(1).get()
+        group = self._db.groups.find_one(
+            {
+                "admin": maker_id,
+                "name": user.group_name.strip()
+            }
+        )
 
-        if len(groups) == 0:
+        if group:
             new_group = {
                 'id': str(uuid.uuid4()),
                 'admin': maker_id,
@@ -252,10 +258,12 @@ class BatchUsers:
                 'participants': [user_id],
                 'disableEditing': False,
             }
-            self._db.collection('groups').document(new_group['id']).set(new_group)
+            self._db.groups.insert_one(new_group)
         else:
-            group = groups[0]
-            participants = group.get('participants')
+            participants = group.participants
             if user_id not in participants:
                 participants.append(user_id)
-                group.reference.update({'participants': participants})
+                self._db.groups.update_one(
+                    {"id": group.id},
+                    {"$set": {"participants": participants}}
+                )