from threading import RLock from Powers import LOGGER from Powers.database import MongoDB INSERTION_LOCK = RLock() DISABLED_CMDS = {} class Disabling(MongoDB): """Class to manage database for Disabling for chats.""" # Database name to connect to to preform operations db_name = "disabled" def __init__(self, chat_id: int) -> None: super().__init__(self.db_name) self.chat_id = chat_id self.chat_info = self.__ensure_in_db() def check_cmd_status(self, cmd: str): with INSERTION_LOCK: # cmds = self.chat_info["commands"] cmds = DISABLED_CMDS[self.chat_id]["commands"] # return bool(cmd in cmds) return bool(cmd in cmds) def add_disable(self, cmd: str): with INSERTION_LOCK: if not self.check_cmd_status(cmd): # DISABLED_CMDS[self.chat_id]["commands"].append(cmd) return self.update( {"_id": self.chat_id}, { "_id": self.chat_id, "commands": self.chat_info["commands"] + [cmd], }, ) def remove_disabled(self, comm: str): with INSERTION_LOCK: if self.check_cmd_status(comm): self.chat_info["commands"].remove(comm) DISABLED_CMDS[self.chat_id]["commands"].remove(comm) return self.update( {"_id": self.chat_id}, { "_id": self.chat_id, "commands": self.chat_info["commands"], }, ) def get_disabled(self): with INSERTION_LOCK: global DISABLED_CMDS try: cmds = DISABLED_CMDS[self.chat_id]["commands"] except KeyError: cmds = self.chat_info["commands"] DISABLED_CMDS[self.chat_id]["commands"] = cmds return cmds @staticmethod def count_disabled_all(): with INSERTION_LOCK: collection = MongoDB(Disabling.db_name) curr = collection.find_all() return sum(len(chat["commands"]) for chat in curr) @staticmethod def count_disabling_chats(): with INSERTION_LOCK: collection = MongoDB(Disabling.db_name) curr = collection.find_all() return sum(1 for chat in curr if chat["commands"]) def set_action(self, action: str): with INSERTION_LOCK: global DISABLED_CMDS DISABLED_CMDS[self.chat_id]["action"] = action return self.update( {"_id": self.chat_id}, {"_id": self.chat_id, "action": action}, ) def get_action(self): with INSERTION_LOCK: global DISABLED_CMDS try: action = DISABLED_CMDS[self.chat_id]["action"] except KeyError: action = self.chat_info["action"] DISABLED_CMDS[self.chat_id]["action"] = action return action @staticmethod def count_action_dis_all(action: str): with INSERTION_LOCK: collection = MongoDB(Disabling.db_name) all_data = collection.find_all({"action": action}) return sum(len(i["commands"]) >= 1 for i in all_data) def rm_all_disabled(self): with INSERTION_LOCK: DISABLED_CMDS[self.chat_id]["commands"] = [] return self.update( {"_id": self.chat_id}, {"commands": []}, ) def __ensure_in_db(self): try: chat_data = DISABLED_CMDS[self.chat_id] except KeyError: chat_data = self.find_one({"_id": self.chat_id}) if not chat_data: new_data = new_data = { "_id": self.chat_id, "commands": [], "action": "none", } self.insert_one(new_data) LOGGER.info(f"Initialized Disabling Document for chat {self.chat_id}") return new_data return chat_data # Migrate if chat id changes! def migrate_chat(self, new_chat_id: int): global DISABLED_CMDS # global only when we are modifying the value old_chat_db = self.find_one({"_id": self.chat_id}) new_data = old_chat_db.update({"_id": new_chat_id}) DISABLED_CMDS[new_chat_id] = DISABLED_CMDS[self.chat_id] del DISABLED_CMDS[self.chat_id] self.insert_one(new_data) self.delete_one({"_id": self.chat_id}) def __load_disable_cache(): global DISABLED_CMDS collection = MongoDB(Disabling.db_name) all_data = collection.find_all() DISABLED_CMDS = { i["_id"]: {"action": i["action"], "commands": i["commands"]} for i in all_data } __load_disable_cache()