Migrate to async_pymongo

This commit is contained in:
Profitroll 2023-08-14 14:52:02 +02:00
parent 5b56919b80
commit cd26990b7e
Signed by: profitroll
GPG Key ID: FA35CAB49DACD3B2
7 changed files with 23 additions and 31 deletions

View File

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

View File

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

View File

@ -1,6 +1,6 @@
"""Module that provides all database columns""" """Module that provides all database columns"""
from pymongo import MongoClient from async_pymongo import AsyncClient
from ujson import loads from ujson import loads
with open("config.json", "r", encoding="utf-8") as f: 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_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"]) 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_sent = db.get_collection("sent")
col_users = db.get_collection("users") col_users = db.get_collection("users")
col_submitted = db.get_collection("submitted") col_submitted = db.get_collection("submitted")

View File

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

View File

@ -25,7 +25,7 @@ async def callback_query_yes(app: PyroClient, clb: CallbackQuery):
user = await app.find_user(clb.from_user) user = await app.find_user(clb.from_user)
fullclb = str(clb.data).split("_") 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: try:
submission = await app.submit_media(fullclb[2]) 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) user = await app.find_user(clb.from_user)
fullclb = str(clb.data).split("_") 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 ( if (
db_entry["temp"]["uuid"] is not None db_entry["temp"]["uuid"] is not None

View File

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

View File

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