Spaces:
Sleeping
Sleeping
File size: 6,624 Bytes
11ae35a 818e429 11ae35a cddafb3 ca4eb6d 292d077 ca4eb6d 292d077 89ad488 292d077 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 292d077 89ad488 292d077 89ad488 ca4eb6d 292d077 89ad488 af1662b ca4eb6d 89ad488 6cef7ec ca4eb6d 292d077 89ad488 292d077 ca4eb6d af1662b ca4eb6d 292d077 ca4eb6d 292d077 89ad488 292d077 89ad488 ca4eb6d 292d077 89ad488 af1662b ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d 51e1c40 ca4eb6d 292d077 ca4eb6d 292d077 ca4eb6d aef1161 51e1c40 aef1161 292d077 89ad488 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 200 201 202 203 204 |
from threading import RLock
from time import time
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 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 or [],
"action": act or "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 or [],
"action": self.chat_info["action"],
}
return cmds or []
@staticmethod
def count_disabled_all():
with INSERTION_LOCK:
collection = MongoDB(Disabling.db_name)
curr = collection.find_all()
return sum(
len(chat["commands"] or []) for chat in curr
)
@staticmethod
def count_disabling_chats():
with INSERTION_LOCK:
collection = MongoDB(Disabling.db_name)
curr = collection.find_all()
return sum(bool(chat["commands"])
for chat in curr)
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 or [],
"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 or [],
"action": val,
}
return val or "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"] or []) >= 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)
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})
def clean_disable(self):
with INSERTION_LOCK:
return 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"] or "none",
"commands": i["commands"] or [],
}
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()
|