File size: 2,601 Bytes
c7dfe8b |
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 |
import threading
from sqlalchemy import Column, String, UnicodeText, distinct, func
from Database.sql import BASE, SESSION
class Disable(BASE):
__tablename__ = "disabled_commands"
chat_id = Column(String(14), primary_key=True)
command = Column(UnicodeText, primary_key=True)
def __init__(self, chat_id, command):
self.chat_id = chat_id
self.command = command
def __repr__(self):
return "Disabled cmd {} in {}".format(self.command, self.chat_id)
Disable.__table__.create(checkfirst=True)
DISABLE_INSERTION_LOCK = threading.RLock()
DISABLED = {}
def disable_command(chat_id, disable):
with DISABLE_INSERTION_LOCK:
disabled = SESSION.query(Disable).get((str(chat_id), disable))
if not disabled:
DISABLED.setdefault(str(chat_id), set()).add(disable)
disabled = Disable(str(chat_id), disable)
SESSION.add(disabled)
SESSION.commit()
return True
SESSION.close()
return False
def enable_command(chat_id, enable):
with DISABLE_INSERTION_LOCK:
disabled = SESSION.query(Disable).get((str(chat_id), enable))
if disabled:
if enable in DISABLED.get(str(chat_id)): # sanity check
DISABLED.setdefault(str(chat_id), set()).remove(enable)
SESSION.delete(disabled)
SESSION.commit()
return True
SESSION.close()
return False
def is_command_disabled(chat_id, cmd):
return str(cmd).lower() in DISABLED.get(str(chat_id), set())
def get_all_disabled(chat_id):
return DISABLED.get(str(chat_id), set())
def num_chats():
try:
return SESSION.query(func.count(distinct(Disable.chat_id))).scalar()
finally:
SESSION.close()
def num_disabled():
try:
return SESSION.query(Disable).count()
finally:
SESSION.close()
def migrate_chat(old_chat_id, new_chat_id):
with DISABLE_INSERTION_LOCK:
chats = SESSION.query(Disable).filter(Disable.chat_id == str(old_chat_id)).all()
for chat in chats:
chat.chat_id = str(new_chat_id)
SESSION.add(chat)
if str(old_chat_id) in DISABLED:
DISABLED[str(new_chat_id)] = DISABLED.get(str(old_chat_id), set())
SESSION.commit()
def __load_disabled_commands():
global DISABLED
try:
all_chats = SESSION.query(Disable).all()
for chat in all_chats:
DISABLED.setdefault(chat.chat_id, set()).add(chat.command)
finally:
SESSION.close()
__load_disabled_commands()
|