Migrate to async_pymongo

classes/pyroclient.py

@@ -141,7 +141,7 @@ class PyroClient(PyroClient):
     async def submit_media(
         self, id: str
     ) -> Tuple[Union[Message, None], Union[str, None]]:
-        db_entry = col_submitted.find_one({"_id": ObjectId(id)})
+        db_entry = await col_submitted.find_one({"_id": ObjectId(id)})
         submission = None
 
         if db_entry is None:
@@ -226,7 +226,7 @@ class PyroClient(PyroClient):
                     )
             raise SubmissionDuplicatesError(str(filepath), duplicates)
 
-        col_submitted.find_one_and_update(
+        await col_submitted.find_one_and_update(
             {"_id": ObjectId(id)}, {"$set": {"done": True}}
         )
 
@@ -258,12 +258,12 @@ class PyroClient(PyroClient):
             * `PyroUser`: PyroUser object
         """
         if (
-            col_users.find_one(
+            await col_users.find_one(
                 {"id": user.id if isinstance(user, User) else user}
             )  # type: ignore
             is None
         ):
-            col_users.insert_one(
+            await col_users.insert_one(
                 {
                     "id": user.id if isinstance(user, User) else user,
                     "locale": user.language_code if isinstance(user, User) else None,
@@ -273,7 +273,7 @@ class PyroClient(PyroClient):
                 }
             )  # type: ignore
 
-        db_record = col_users.find_one(
+        db_record = await col_users.find_one(
             {"id": user.id if isinstance(user, User) else user}
         )  # type: ignore
 

classes/pyrouser.py

@@ -20,19 +20,19 @@ class PyroUser:
     cooldown: datetime
     subscription: dict
 
-    async def update_locale(self, locale: str):
-        col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})
+    async def update_locale(self, locale: str) -> None:
+        await col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})
 
-    async def update_cooldown(self, time: datetime = datetime.now()):
-        col_users.update_one({"_id": self._id}, {"$set": {"cooldown": time}})
+    async def update_cooldown(self, time: datetime = datetime.now()) -> None:
+        await col_users.update_one({"_id": self._id}, {"$set": {"cooldown": time}})
 
     async def block(self) -> None:
         """Ban user from using command and submitting content."""
-        col_users.update_one({"_id": self._id}, {"$set": {"banned": True}})
+        await col_users.update_one({"_id": self._id}, {"$set": {"banned": True}})
 
     async def unblock(self) -> None:
         """Allow user to use command and submit posts again."""
-        col_users.update_one({"_id": self._id}, {"$set": {"banned": False}})
+        await col_users.update_one({"_id": self._id}, {"$set": {"banned": False}})
 
     async def is_limited(self, app: Union[PyroClient, None] = None) -> bool:
         """Check if user is on a cooldown after submitting something.
@@ -41,11 +41,9 @@ class PyroUser:
             `bool`: Must be `True` if on the cooldown and `False` if not
         """
         admins = (
-            app.admins
-            if app is not None
-            else (
-                await config_get("admins", "bot") + [await config_get("owner", "bot")]
-            )
+            await config_get("admins", "bot") + [await config_get("owner", "bot")]
+            if app is None
+            else app.admins
         )
 
         return (datetime.now() - self.cooldown).total_seconds() < (

modules/database.py

@@ -1,6 +1,6 @@
 """Module that provides all database columns"""
 
-from pymongo import MongoClient
+from async_pymongo import AsyncClient
 from ujson import loads
 
 with open("config.json", "r", encoding="utf-8") as f:
@@ -20,15 +20,9 @@ else:
         db_config["host"], db_config["port"], db_config["name"]
     )
 
-db_client = MongoClient(con_string)
+db_client = AsyncClient(con_string)
 db = db_client.get_database(name=db_config["name"])
 
-collections = db.list_collection_names()
-
-for collection in ["sent", "users", "submitted"]:
-    if collection not in collections:
-        db.create_collection(collection)
-
 col_sent = db.get_collection("sent")
 col_users = db.get_collection("users")
 col_submitted = db.get_collection("submitted")

modules/sender.py

@@ -179,7 +179,7 @@ async def send_content(app: PyroClient, http_session: ClientSession) -> None:
 
         del response
 
-        submitted = col_submitted.find_one({"temp.file": media.filename})
+        submitted = await col_submitted.find_one({"temp.file": media.filename})
 
         if submitted is not None and submitted["caption"] is not None:
             caption = submitted["caption"].strip()
@@ -229,7 +229,7 @@ async def send_content(app: PyroClient, http_session: ClientSession) -> None:
             # rmtree(path.join(app.config['locations']['tmp'], tmp_dir), ignore_errors=True)
             return
 
-        col_sent.insert_one(
+        await col_sent.insert_one(
             {
                 "date": datetime.now(),
                 "image": media.id,

plugins/callbacks/submission.py

@@ -25,7 +25,7 @@ async def callback_query_yes(app: PyroClient, clb: CallbackQuery):
     user = await app.find_user(clb.from_user)
     fullclb = str(clb.data).split("_")
 
-    db_entry = col_submitted.find_one({"_id": ObjectId(fullclb[2])})
+    db_entry = await col_submitted.find_one({"_id": ObjectId(fullclb[2])})
 
     try:
         submission = await app.submit_media(fullclb[2])
@@ -127,7 +127,7 @@ async def callback_query_no(app: PyroClient, clb: CallbackQuery):
     user = await app.find_user(clb.from_user)
     fullclb = str(clb.data).split("_")
 
-    db_entry = col_submitted.find_one_and_delete({"_id": ObjectId(fullclb[2])})
+    db_entry = await col_submitted.find_one_and_delete({"_id": ObjectId(fullclb[2])})
 
     if (
         db_entry["temp"]["uuid"] is not None

plugins/handlers/submission.py

@@ -152,7 +152,7 @@ async def get_submission(app: PyroClient, msg: Message):
                 + sep,
             )
 
-            inserted = col_submitted.insert_one(
+            inserted = await col_submitted.insert_one(
                 {
                     "user": msg.from_user.id,
                     "date": datetime.now(),
@@ -165,7 +165,7 @@ async def get_submission(app: PyroClient, msg: Message):
             )
 
         else:
-            inserted = col_submitted.insert_one(
+            inserted = await col_submitted.insert_one(
                 {
                     "user": msg.from_user.id,
                     "date": datetime.now(),

requirements.txt

@@ -4,12 +4,12 @@ convopyro==0.5
 pillow~=10.0.0
 psutil~=5.9.4
 pykeyboard==0.1.5
-pymongo~=4.4.0
 pyrogram==2.0.106
 python_dateutil==2.8.2
 pytimeparse~=1.1.8
 tgcrypto==1.2.5
 uvloop==0.17.0
 --extra-index-url https://git.end-play.xyz/api/packages/profitroll/pypi/simple
+async_pymongo==0.1.4
 libbot[speed,pyrogram]==2.0.1
 photosapi_client==0.5.0