Spaces:
Sleeping
Sleeping
File size: 4,341 Bytes
11ae35a ca4eb6d 11ae35a cddafb3 ca4eb6d 89ad488 ca4eb6d 6cef7ec ca4eb6d 6cef7ec ca4eb6d 6cef7ec ca4eb6d |
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 |
from threading import RLock
from time import time
from Powers import LOGGER
from Powers.database import MongoDB
INSERTION_LOCK = RLock()
class Chats(MongoDB):
"""Class to manage users for bot."""
# Database name to connect to to preform operations
db_name = "chats"
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 user_is_in_chat(self, user_id: int):
return user_id in set(self.chat_info["users"])
def update_chat(self, chat_name: str, user_id: int):
with INSERTION_LOCK:
if chat_name == self.chat_info["chat_name"] and self.user_is_in_chat(
user_id,
):
return True
if chat_name != self.chat_info["chat_name"] and self.user_is_in_chat(
user_id,
):
return self.update(
{"_id": self.chat_id},
{"chat_name": chat_name},
)
if chat_name == self.chat_info["chat_name"] and not self.user_is_in_chat(
user_id,
):
self.chat_info["users"].append(user_id)
return self.update(
{"_id": self.chat_id},
{"users": self.chat_info["users"]},
)
users_old = self.chat_info["users"]
users_old.append(user_id)
users = list(set(users_old))
return self.update(
{"_id": self.chat_id},
{
"_id": self.chat_id,
"chat_name": chat_name,
"users": users,
},
)
def count_chat_users(self):
with INSERTION_LOCK:
return len(self.chat_info["users"]) or 0
def chat_members(self):
with INSERTION_LOCK:
return self.chat_info["users"]
@staticmethod
def remove_chat(chat_id: int):
with INSERTION_LOCK:
collection = MongoDB(Chats.db_name)
collection.delete_one({"_id": chat_id})
@staticmethod
def count_chats():
with INSERTION_LOCK:
collection = MongoDB(Chats.db_name)
return collection.count() or 0
@staticmethod
def list_chats_by_id():
with INSERTION_LOCK:
collection = MongoDB(Chats.db_name)
chats = collection.find_all()
chat_list = {i["_id"] for i in chats}
return list(chat_list)
@staticmethod
def list_chats_full():
with INSERTION_LOCK:
collection = MongoDB(Chats.db_name)
return collection.find_all()
@staticmethod
def get_chat_info(chat_id: int):
with INSERTION_LOCK:
collection = MongoDB(Chats.db_name)
return collection.find_one({"_id": chat_id})
def load_from_db(self):
with INSERTION_LOCK:
return self.find_all()
def __ensure_in_db(self):
chat_data = self.find_one({"_id": self.chat_id})
if not chat_data:
new_data = {"_id": self.chat_id, "chat_name": "", "users": []}
self.insert_one(new_data)
return new_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})
self.insert_one(new_data)
self.delete_one({"_id": self.chat_id})
@staticmethod
def repair_db(collection):
all_data = collection.find_all()
keys = {"chat_name": "", "users": []}
for data in all_data:
for key, val in keys.items():
try:
_ = data[key]
except KeyError:
LOGGER.warning(
f"Repairing Chats Database - setting '{key}:{val}' for {data['_id']}",
)
collection.update({"_id": data["_id"]}, {key: val})
def __pre_req_chats():
start = time()
LOGGER.info("Starting Chats Database Repair...")
collection = MongoDB(Chats.db_name)
Chats.repair_db(collection)
LOGGER.info(f"Done in {round((time() - start), 3)}s!")
|