Spaces:
Sleeping
Sleeping
File size: 6,694 Bytes
818e429 cddafb3 007891e ca4eb6d af1662b ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 af1662b ca4eb6d 292d077 ca4eb6d af1662b ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 af1662b ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 5ed2636 ca4eb6d 5ed2636 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
from time import time
from Powers import LOGGER
from threading import RLock
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 perform 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:
try:
cmds = DISABLED_CMDS[self.chat_id]["commands"]
except KeyError:
cmds = self.chat_info["commands"]
act = self.chat_info["action"]
DISABLED_CMDS[self.chat_id] = {
"command": cmds if cmds else [],
"action": act if act else "none",
}
# return bool(cmd in cmds)
return bool(cmd in cmds if cmds else [])
def add_disable(self, cmd: str):
with INSERTION_LOCK:
if not self.check_cmd_status(cmd):
return self.update(
{"_id": self.chat_id},
{
"_id": self.chat_id,
"commands": self.chat_info["commands"].append(cmd),
},
)
def remove_disabled(self, comm: str):
with INSERTION_LOCK:
if self.check_cmd_status(comm):
self.chat_info["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:
try:
cmds = DISABLED_CMDS[self.chat_id]["commands"]
except KeyError:
cmds = self.chat_info["commands"]
DISABLED_CMDS[self.chat_id] = {
"commands": cmds if cmds else [],
"action": self.chat_info["action"],
}
return cmds if cmds else []
@staticmethod
def count_disabled_all():
with INSERTION_LOCK:
collection = MongoDB(Disabling.db_name)
curr = collection.find_all()
return sum(
len(chat["commands"] if chat["commands"] else []) 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:
try:
DISABLED_CMDS[self.chat_id]["action"] = action
except KeyError:
cmds = self.chat_info["commands"]
DISABLED_CMDS[self.chat_id] = {
"commands": cmds if cmds else [],
"action": action,
}
return self.update(
{"_id": self.chat_id},
{"_id": self.chat_id, "action": action},
)
def get_action(self):
with INSERTION_LOCK:
try:
val = DISABLED_CMDS[self.chat_id]["action"]
except KeyError:
cmds = self.chat_info["commands"]
val = self.chat_info["action"]
DISABLED_CMDS[self.chat_id] = {
"commands": cmds if cmds else [],
"action": val,
}
return val if val else "none"
@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"] if i["commands"] else []) >= 1 for i in all_data
)
def rm_all_disabled(self):
with INSERTION_LOCK:
try:
DISABLED_CMDS[self.chat_id]["commands"] = []
except KeyError:
DISABLED_CMDS[self.chat_id] = {
"commands": [],
"action": self.chat_info["action"],
}
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 = {
"_id": self.chat_id,
"commands": [],
"action": "none",
}
DISABLED_CMDS[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
DISABLED_CMDS[self.chat_id] = chat_data
return chat_data
# Migrate if chat id changes!
def migrate_chat(self, new_chat_id: int):
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.pop(self.chat_id)
self.insert_one(new_data)
self.delete_one({"_id": self.chat_id})
@staticmethod
def repair_db(collection):
global DISABLED_CMDS
all_data = collection.find_all()
DISABLED_CMDS = {
i["_id"]: {
"action": i["action"] if i["action"] else "none",
"commands": i["commands"] if i["commands"] else [],
}
for i in all_data
}
keys = {
"commands": [],
"action": "none",
}
for data in all_data:
for key, val in keys.items():
try:
_ = data[key]
except KeyError:
LOGGER.warning(
f"Repairing Disabling Database - setting '{key}:{val}' for {data['_id']}",
)
collection.update({"_id": data["_id"]}, {key: val})
def __pre_req_disabling():
start = time()
LOGGER.info("Starting disabling Database Repair ...")
collection = MongoDB(Disabling.db_name)
Disabling.repair_db(collection)
LOGGER.info(f"Done in {round((time() - start), 3)}s!")
__pre_req_disabling()
|