|
from async_pymongo import AsyncClient |
|
|
|
from Mikobot import DB_NAME, MONGO_DB_URI |
|
|
|
|
|
class UsersData: |
|
def __init__(self, uri, database_name): |
|
self._client = AsyncClient(uri) |
|
self.db = self._client[database_name] |
|
self.col = self.db["userlist"] |
|
self.grp = self.db["groups"] |
|
|
|
@staticmethod |
|
def new_user(id, name): |
|
return dict( |
|
id=id, |
|
name=name, |
|
ban_status=dict( |
|
is_banned=False, |
|
ban_reason="", |
|
), |
|
) |
|
|
|
@staticmethod |
|
def new_group(id, title): |
|
return dict( |
|
id=id, |
|
title=title, |
|
chat_status=dict( |
|
is_disabled=False, |
|
reason="", |
|
), |
|
) |
|
|
|
async def add_user(self, id, name): |
|
user = self.new_user(id, name) |
|
await self.col.insert_one(user) |
|
|
|
async def is_user_exist(self, id): |
|
user = await self.col.find_one({"id": int(id)}) |
|
return bool(user) |
|
|
|
async def total_users_count(self): |
|
return await self.col.count_documents({}) |
|
|
|
async def remove_ban(self, id): |
|
return await self.col.delete_one({"_id": id}) |
|
|
|
async def ban_user(self, user_id, ban_reason="No Reason"): |
|
return await self.col.insert_one({"_id": user_id, "reason": ban_reason}) |
|
|
|
async def get_ban_status(self, id): |
|
user = await self.col.find_one({"_id": int(id)}) |
|
if user: |
|
return True, user |
|
return False, None |
|
|
|
async def get_all_users(self): |
|
return self.col.find({}) |
|
|
|
async def delete_user(self, user_id): |
|
await self.col.delete_many({"id": int(user_id)}) |
|
|
|
async def is_chat_exist(self, id): |
|
user = await self.grp.find_one({"id": int(id)}) |
|
return bool(user) |
|
|
|
async def get_banned(self): |
|
users = self.col.find({"ban_status.is_banned": True}) |
|
chats = self.grp.find({"chat_status.is_disabled": True}) |
|
b_chats = [chat["id"] async for chat in chats] |
|
b_users = [user["id"] async for user in users] |
|
return b_users, b_chats |
|
|
|
async def add_chat(self, chat, title): |
|
chat = self.new_group(chat, title) |
|
await self.grp.insert_one(chat) |
|
|
|
async def get_chat(self, chat): |
|
chat = await self.grp.find_one({"id": int(chat)}) |
|
return chat.get("chat_status") if chat else False |
|
|
|
async def re_enable_chat(self, id): |
|
chat_status = dict( |
|
is_disabled=False, |
|
reason="", |
|
) |
|
await self.grp.update_one( |
|
{"id": int(id)}, {"$set": {"chat_status": chat_status}} |
|
) |
|
|
|
async def disable_chat(self, chat, reason="No Reason"): |
|
chat_status = dict( |
|
is_disabled=True, |
|
reason=reason, |
|
) |
|
await self.grp.update_one( |
|
{"id": int(chat)}, {"$set": {"chat_status": chat_status}} |
|
) |
|
|
|
async def total_chat_count(self): |
|
return await self.grp.count_documents({}) |
|
|
|
async def get_all_chats(self): |
|
return self.grp.find({}) |
|
|
|
async def get_db_size(self): |
|
return (await self.db.command("dbstats"))["dataSize"] |
|
|
|
|
|
db = UsersData(MONGO_DB_URI, DB_NAME) |
|
|
