Upload database.py

database.py

@@ -10,38 +10,14 @@ from motor import motor_asyncio
 from motor.core import AgnosticClient
 
 from logger import LOGS
-
-load_dotenv()
-MONGO_URL = os.environ["MONGO_URL"]
+from config import MONGO_URL
 
 class Database:
     def __init__(self, uri: str) -> None:
         self.client: AgnosticClient = motor_asyncio.AsyncIOMotorClient(uri)
         self.db = self.client["Akeno"]
-
-        self.afk = self.db["afk"]
-        self.antiflood = self.db["antiflood"]
-        self.autopost = self.db["autopost"]
-        self.blacklist = self.db["blacklist"]
-        self.echo = self.db["echo"]
-        self.env = self.db["env"]
-        self.filter = self.db["filter"]
-        self.forcesub = self.db["forcesub"]
-        self.gachabots = self.db["gachabots"]
-        self.cohere = self.db["cohere"]
-        self.chatbot = self.db["chatbot"]
-        self.backup_chatbot = self.db["openaichat"]
-        self.antiarabic = self.db["antiarabic"]
-        self.antinsfw = self.db["antinsfw"]
-        self.gban = self.db["gban"]
-        self.gmute = self.db["gmute"]
-        self.greetings = self.db["greetings"]
-        self.mute = self.db["mute"]
-        self.pmpermit = self.db["pmpermit"]
-        self.session = self.db["session"]
-        self.snips = self.db["snips"]
-        self.stan_users = self.db["stan_users"]
-
+        self.backup_chatbot = self.db["metai"]
+    
     async def connect(self):
         try:
             await self.client.admin.command("ping")
@@ -78,571 +54,6 @@ class Database:
     async def get_all_env(self) -> list:
         return [i async for i in self.env.find({})]
 
-    async def is_stan(self, client: int, user_id: int) -> bool:
-        if await self.stan_users.find_one({"client": client, "user_id": user_id}):
-            return True
-        return False
-
-    async def add_stan(self, client: int, user_id: int) -> bool:
-        if await self.is_stan(client, user_id):
-            return False
-        await self.stan_users.insert_one(
-            {"client": client, "user_id": user_id, "date": self.get_datetime()}
-        )
-        return True
-
-    async def rm_stan(self, client: int, user_id: int) -> bool:
-        if not await self.is_stan(client, user_id):
-            return False
-        await self.stan_users.delete_one({"client": client, "user_id": user_id})
-        return True
-
-    async def get_stans(self, client: int) -> list:
-        return [i async for i in self.stan_users.find({"client": client})]
-
-    async def get_all_stans(self) -> list:
-        return [i async for i in self.stan_users.find({})]
-
-    async def is_session(self, user_id: int) -> bool:
-        if await self.session.find_one({"user_id": user_id}):
-            return True
-        return False
-
-    async def update_session(self, user_id: int, session: str) -> None:
-        await self.session.update_one(
-            {"user_id": user_id},
-            {"$set": {"session": session, "date": self.get_datetime()}},
-            upsert=True,
-        )
-
-    async def rm_session(self, user_id: int) -> None:
-        await self.session.delete_one({"user_id": user_id})
-
-    async def get_session(self, user_id: int):
-        if not await self.is_session(user_id):
-            return False
-        data = await self.session.find_one({"user_id": user_id})
-        return data
-
-    async def get_all_sessions(self) -> list:
-        return [i async for i in self.session.find({})]
-
-    async def is_gbanned(self, user_id: int) -> bool:
-        if await self.gban.find_one({"user_id": user_id}):
-            return True
-        return False
-
-    async def add_gban(self, user_id: int, reason: str) -> bool:
-        if await self.is_gbanned(user_id):
-            return False
-        await self.gban.insert_one(
-            {"user_id": user_id, "reason": reason, "date": self.get_datetime()}
-        )
-        return True
-
-    async def rm_gban(self, user_id: int):
-        if not await self.is_gbanned(user_id):
-            return None
-        reason = (await self.gban.find_one({"user_id": user_id}))["reason"]
-        await self.gban.delete_one({"user_id": user_id})
-        return reason
-
-    async def get_gban(self) -> list:
-        return [i async for i in self.gban.find({})]
-
-    async def get_gban_user(self, user_id: int) -> dict | None:
-        if not await self.is_gbanned(user_id):
-            return None
-        return await self.gban.find_one({"user_id": user_id})
-
-    async def is_gmuted(self, user_id: int) -> bool:
-        if await self.gmute.find_one({"user_id": user_id}):
-            return True
-        return False
-
-    async def add_gmute(self, user_id: int, reason: str) -> bool:
-        if await self.is_gmuted(user_id):
-            return False
-        await self.gmute.insert_one(
-            {"user_id": user_id, "reason": reason, "date": self.get_datetime()}
-        )
-        return True
-
-    async def rm_gmute(self, user_id: int):
-        if not await self.is_gmuted(user_id):
-            return None
-        reason = (await self.gmute.find_one({"user_id": user_id}))["reason"]
-        await self.gmute.delete_one({"user_id": user_id})
-        return reason
-
-    async def get_gmute(self) -> list:
-        return [i async for i in self.gmute.find({})]
-
-    async def add_mute(self, client: int, user_id: int, chat_id: int, reason: str):
-        await self.mute.update_one(
-            {"client": client, "user_id": user_id, "chat_id": chat_id},
-            {"$set": {"reason": reason, "date": self.get_datetime()}},
-            upsert=True,
-        )
-
-    async def rm_mute(self, client: int, user_id: int, chat_id: int) -> str:
-        reason = (await self.get_mute(client, user_id, chat_id))["reason"]
-        await self.mute.delete_one({"client": client, "user_id": user_id, "chat_id": chat_id})
-        return reason
-
-    async def is_muted(self, client: int, user_id: int, chat_id: int) -> bool:
-        if await self.get_mute(client, user_id, chat_id):
-            return True
-        return False
-
-    async def get_mute(self, client: int, user_id: int, chat_id: int):
-        data = await self.mute.find_one({"client": client, "user_id": user_id, "chat_id": chat_id})
-        return data
-
-    async def set_afk(
-        self, user_id: int, reason: str, media: int, media_type: str
-    ) -> None:
-        await self.afk.update_one(
-            {"user_id": user_id},
-            {
-                "$set": {
-                    "reason": reason,
-                    "time": time.time(),
-                    "media": media,
-                    "media_type": media_type,
-                }
-            },
-            upsert=True,
-        )
-
-    async def get_afk(self, user_id: int):
-        data = await self.afk.find_one({"user_id": user_id})
-        return data
-
-    async def is_afk(self, user_id: int) -> bool:
-        if await self.afk.find_one({"user_id": user_id}):
-            return True
-        return False
-
-    async def rm_afk(self, user_id: int) -> None:
-        await self.afk.delete_one({"user_id": user_id})
-
-    async def set_flood(self, client_chat: tuple[int, int], settings: dict):
-        await self.antiflood.update_one(
-            {"client": client_chat[0], "chat": client_chat[1]},
-            {"$set": settings},
-            upsert=True,
-        )
-
-    async def get_flood(self, client_chat: tuple[int, int]):
-        data = await self.antiflood.find_one(
-            {"client": client_chat[0], "chat": client_chat[1]}
-        )
-        return data or {}
-
-    async def is_flood(self, client_chat: tuple[int, int]) -> bool:
-        data = await self.get_flood(client_chat)
-
-        if not data:
-            return False
-
-        if data["limit"] == 0:
-            return False
-
-        return True
-
-    async def get_all_floods(self) -> list:
-        return [i async for i in self.antiflood.find({})]
-
-    async def set_autopost(self, client: int, from_channel: int, to_channel: int):
-        await self.autopost.update_one(
-            {"client": client},
-            {
-                "$push": {
-                    "autopost": {
-                        "from_channel": from_channel,
-                        "to_channel": to_channel,
-                        "date": self.get_datetime(),
-                    }
-                }
-            },
-            upsert=True,
-        )
-
-    async def get_autopost(self, client: int, from_channel: int):
-        data = await self.autopost.find_one(
-            {
-                "client": client,
-                "autopost": {"$elemMatch": {"from_channel": from_channel}},
-            }
-        )
-        return data
-
-    async def is_autopost(
-        self, client: int, from_channel: int, to_channel: int = None
-    ) -> bool:
-        if to_channel:
-            data = await self.autopost.find_one(
-                {
-                    "client": client,
-                    "autopost": {
-                        "$elemMatch": {
-                            "from_channel": from_channel,
-                            "to_channel": to_channel,
-                        }
-                    },
-                }
-            )
-        else:
-            data = await self.autopost.find_one(
-                {
-                    "client": client,
-                    "autopost": {"$elemMatch": {"from_channel": from_channel}},
-                }
-            )
-        return True if data else False
-
-    async def rm_autopost(self, client: int, from_channel: int, to_channel: int):
-        await self.autopost.update_one(
-            {"client": client},
-            {
-                "$pull": {
-                    "autopost": {
-                        "from_channel": from_channel,
-                        "to_channel": to_channel,
-                    }
-                }
-            },
-        )
-
-    async def get_all_autoposts(self, client: int) -> list:
-        return [i async for i in self.autopost.find({"client": client})]
-
-    async def add_blacklist(self, client: int, chat: int, blacklist: str):
-        await self.blacklist.update_one(
-            {"client": client, "chat": chat},
-            {"$push": {"blacklist": blacklist}},
-            upsert=True,
-        )
-
-    async def rm_blacklist(self, client: int, chat: int, blacklist: str):
-        await self.blacklist.update_one(
-            {"client": client, "chat": chat},
-            {"$pull": {"blacklist": blacklist}},
-        )
-
-    async def is_blacklist(self, client: int, chat: int, blacklist: str) -> bool:
-        blacklists = await self.get_all_blacklists(client, chat)
-        if blacklist in blacklists:
-            return True
-        return False
-
-    async def get_all_blacklists(self, client: int, chat: int) -> list:
-        data = await self.blacklist.find_one({"client": client, "chat": chat})
-
-        if not data:
-            return []
-
-        return data["blacklist"]
-
-    async def get_blacklist_clients(self) -> list:
-        return [i async for i in self.blacklist.find({})]
-
-    async def set_echo(self, client: int, chat: int, user: int):
-        await self.echo.update_one(
-            {"client": client, "chat": chat},
-            {"$push": {"echo": user}},
-            upsert=True,
-        )
-
-    async def rm_echo(self, client: int, chat: int, user: int):
-        await self.echo.update_one(
-            {"client": client, "chat": chat},
-            {"$pull": {"echo": user}},
-        )
-
-    async def is_echo(self, client: int, chat: int, user: int) -> bool:
-        data = await self.get_all_echo(client, chat)
-        if user in data:
-            return True
-        return False
-
-    async def get_all_echo(self, client: int, chat: int) -> list:
-        data = await self.echo.find_one({"client": client, "chat": chat})
-
-        if not data:
-            return []
-
-        return data["echo"]
-
-    async def set_filter(self, client: int, chat: int, keyword: str, msgid: int):
-        await self.filter.update_one(
-            {"client": client, "chat": chat},
-            {"$push": {"filter": {"keyword": keyword, "msgid": msgid}}},
-            upsert=True,
-        )
-
-    async def rm_filter(self, client: int, chat: int, keyword: str):
-        await self.filter.update_one(
-            {"client": client, "chat": chat},
-            {"$pull": {"filter": {"keyword": keyword}}},
-        )
-
-    async def rm_all_filters(self, client: int, chat: int):
-        await self.filter.delete_one({"client": client, "chat": chat})
-
-    async def is_filter(self, client: int, chat: int, keyword: str) -> bool:
-        data = await self.get_filter(client, chat, keyword)
-        return True if data else False
-
-    async def get_filter(self, client: int, chat: int, keyword: str):
-        data = await self.filter.find_one(
-            {
-                "client": client,
-                "chat": chat,
-                "filter": {"$elemMatch": {"keyword": keyword}},
-            }
-        )
-        return data
-
-    async def get_all_filters(self, client: int, chat: int) -> list:
-        data = await self.filter.find_one({"client": client, "chat": chat})
-
-        if not data:
-            return []
-
-        return data["filter"]
-
-    async def set_snip(self, client: int, chat: int, keyword: str, msgid: int):
-        await self.snips.update_one(
-            {"client": client, "chat": chat},
-            {"$push": {"snips": {"keyword": keyword, "msgid": msgid}}},
-            upsert=True,
-        )
-
-    async def rm_snip(self, client: int, chat: int, keyword: str):
-        await self.snips.update_one(
-            {"client": client, "chat": chat},
-            {"$pull": {"snips": {"keyword": keyword}}},
-        )
-
-    async def rm_all_snips(self, client: int, chat: int):
-        await self.snips.delete_one({"client": client, "chat": chat})
-
-    async def is_snip(self, client: int, chat: int, keyword: str) -> bool:
-        data = await self.get_snip(client, chat, keyword)
-        return True if data else False
-
-    async def get_snip(self, client: int, chat: int, keyword: str):
-        data = await self.snips.find_one(
-            {
-                "client": client,
-                "chat": chat,
-                "snips": {"$elemMatch": {"keyword": keyword}},
-            }
-        )
-        return data
-
-    async def get_all_snips(self, client: int, chat: int) -> list:
-        data = await self.snips.find_one({"client": client, "chat": chat})
-
-        if not data:
-            return []
-
-        return data["snips"]
-
-    async def add_pmpermit(self, client: int, user: int):
-        await self.pmpermit.update_one(
-            {"client": client, "user": user},
-            {"$set": {"date": self.get_datetime()}},
-            upsert=True,
-        )
-
-    async def rm_pmpermit(self, client: int, user: int):
-        await self.pmpermit.delete_one({"client": client, "user": user})
-
-    async def is_pmpermit(self, client: int, user: int) -> bool:
-        data = await self.get_pmpermit(client, user)
-        return True if data else False
-
-    async def get_pmpermit(self, client: int, user: int):
-        data = await self.pmpermit.find_one({"client": client, "user": user})
-        return data
-
-    async def get_all_pmpermits(self, client: int) -> list:
-        return [i async for i in self.pmpermit.find({"client": client})]
-
-    async def set_welcome(self, client: int, chat: int, message: int):
-        await self.greetings.update_one(
-            {"client": client, "chat": chat, "welcome": True},
-            {"$set": {"message": message}},
-            upsert=True,
-        )
-
-    async def rm_welcome(self, client: int, chat: int):
-        await self.greetings.delete_one(
-            {"client": client, "chat": chat, "welcome": True}
-        )
-
-    async def is_welcome(self, client: int, chat: int) -> bool:
-        data = await self.get_welcome(client, chat)
-        return True if data else False
-
-    async def get_welcome(self, client: int, chat: int):
-        data = await self.greetings.find_one(
-            {"client": client, "chat": chat, "welcome": True}
-        )
-        return data
-
-    async def set_goodbye(self, client: int, chat: int, message: int):
-        await self.greetings.update_one(
-            {"client": client, "chat": chat, "welcome": False},
-            {"$set": {"message": message}},
-            upsert=True,
-        )
-
-    async def rm_goodbye(self, client: int, chat: int):
-        await self.greetings.delete_one(
-            {"client": client, "chat": chat, "welcome": False}
-        )
-
-    async def is_goodbye(self, client: int, chat: int) -> bool:
-        data = await self.get_goodbye(client, chat)
-        return True if data else False
-
-    async def get_goodbye(self, client: int, chat: int):
-        data = await self.greetings.find_one(
-            {"client": client, "chat": chat, "welcome": False}
-        )
-        return data
-
-    async def get_all_greetings(self, client: int) -> list:
-        return [i async for i in self.greetings.find({"client": client})]
-
-    async def add_forcesub(self, chat: int, must_join: int):
-        await self.forcesub.update_one(
-            {"chat": chat},
-            {"$push": {"must_join": must_join}},
-            upsert=True,
-        )
-
-    async def rm_forcesub(self, chat: int, must_join: int) -> int:
-        await self.forcesub.update_one(
-            {"chat": chat},
-            {"$pull": {"must_join": must_join}},
-        )
-        data = await self.forcesub.find_one({"chat": chat})
-        return len(data["must_join"])
-
-    async def rm_all_forcesub(self, in_chat: int):
-        await self.forcesub.delete_one({"chat": in_chat})
-
-    async def is_forcesub(self, chat: int, must_join: int) -> bool:
-        data = await self.get_forcesub(chat)
-        if must_join in data["must_join"]:
-            return True
-        return False
-
-    async def get_forcesub(self, in_chat: int):
-        data = await self.forcesub.find_one({"chat": in_chat})
-        return data
-
-    async def get_all_forcesubs(self) -> list:
-        return [i async for i in self.forcesub.find({})]
-
-    async def add_gachabot(
-        self, client: int, bot: tuple[int, str], catch_command: str, chat_id: int
-    ):
-        await self.gachabots.update_one(
-            {"client": client, "bot": bot[0]},
-            {
-                "$set": {
-                    "username": bot[1],
-                    "catch_command": catch_command,
-                    "chat_id": chat_id,
-                    "date": self.get_datetime(),
-                }
-            },
-            upsert=True,
-        )
-
-    async def rm_gachabot(self, client: int, bot: int, chat_id: int = None):
-        if chat_id:
-            await self.gachabots.delete_one(
-                {"client": client, "bot": bot, "chat_id": chat_id}
-            )
-        else:
-            await self.gachabots.delete_one({"client": client, "bot": bot})
-
-    async def is_gachabot(self, client: int, bot: int, chat_id: int) -> bool:
-        data = await self.get_gachabot(client, bot, chat_id)
-        return True if data else False
-
-    async def get_gachabot(self, client: int, bot: int, chat_id: int):
-        data = await self.gachabots.find_one(
-            {"client": client, "bot": bot, "chat_id": chat_id}
-        )
-
-        return data
-
-    async def get_all_gachabots(self, client: int) -> list:
-        return [i async for i in self.gachabots.find({"client": client})]
-
-    async def get_all_gachabots_id(self) -> list:
-        data = await self.gachabots.distinct("bot")
-        return data
-
-    async def _get_cohere_chat_from_db(self, user_id):
-        user_data = await self.cohere.find_one({"user_id": user_id})
-        return user_data.get("cohere_chat", []) if user_data else []
-
-    async def _update_cohere_chat_in_db(self, user_id, cohere_chat):
-        await self.cohere.update_one(
-            {"user_id": user_id},
-            {"$set": {"cohere_chat": cohere_chat}},
-            upsert=True
-        )
-
-    async def _clear_history_in_db(self, user_id):
-        unset_clear = {"cohere_chat": None}
-        return await self.cohere.update_one({"user_id": user_id}, {"$unset": unset_clear})
-
-    async def clear_database(self, user_id):
-        """Clear the cohere history for the current user."""
-        result = await self._clear_history_in_db(user_id)
-        if result.modified_count > 0:
-            return "Chat history cleared successfully."
-        else:
-            return "No chat history found to clear."
-
-    async def set_chat_setting(self, chat_id, isboolean):
-        await self.antiarabic.update_one(
-            {"chat_id": chat_id},
-            {"$set": {"arabic": isboolean}},
-            upsert=True
-        )
-
-    async def chat_antiarabic(self, chat_id):
-        user_data = await self.antiarabic.find_one({"chat_id": chat_id})
-        if user_data:
-            return user_data.get("arabic", False)
-        return False
-
-    async def add_chatbot(self, chat_id, user_id):
-        await self.chatbot.update_one(
-            {"chat_id": chat_id},
-            {"$set": {"user_id": user_id}},
-            upsert=True
-        )
-
-    async def get_chatbot(self, chat_id):
-        user_data = await self.chatbot.find_one({"chat_id": chat_id})
-        return user_data.get("user_id") if user_data else None
-
-    async def remove_chatbot(self, chat_id):
-        unset_data = {"user_id": None}
-        return await self.chatbot.update_one({"chat_id": chat_id}, {"$unset": unset_data})
-
     async def _update_openai_chat_in_db(self, user_id, chatbot_chat):
         await self.backup_chatbot.update_one(
             {"user_id": user_id},
@@ -665,18 +76,4 @@ class Database:
         else:
             return "No chat history found to clear."
 
-    async def set_chat_setting_antinsfw(self, chat_id, isboolean):
-        await self.antinsfw.update_one(
-            {"chat_id": chat_id},
-            {"$set": {"antinsfw": isboolean}},
-            upsert=True
-        )
-
-    async def chat_antinsfw(self, chat_id):
-        user_data = await self.antinsfw.find_one({"chat_id": chat_id})
-        if user_data:
-            return user_data.get("antinsfw", False)
-        return False
-
-
-db = Database(MONGO_URL)
+db = Database(MONGO_URL)