import threading from sqlalchemy import Boolean, Column, Integer, String, UnicodeText, distinct, func from Database.sql import BASE, SESSION from Mikobot.plugins.helper_funcs.msg_types import Types class Notes(BASE): __tablename__ = "notes" chat_id = Column(String(14), primary_key=True) name = Column(UnicodeText, primary_key=True) value = Column(UnicodeText, nullable=False) file = Column(UnicodeText) is_reply = Column(Boolean, default=False) has_buttons = Column(Boolean, default=False) msgtype = Column(Integer, default=Types.BUTTON_TEXT.value) def __init__(self, chat_id, name, value, msgtype, file=None): self.chat_id = str(chat_id) # ensure string self.name = name self.value = value self.msgtype = msgtype self.file = file def __repr__(self): return "" % self.name class Buttons(BASE): __tablename__ = "note_urls" id = Column(Integer, primary_key=True, autoincrement=True) chat_id = Column(String(14), primary_key=True) note_name = Column(UnicodeText, primary_key=True) name = Column(UnicodeText, nullable=False) url = Column(UnicodeText, nullable=False) same_line = Column(Boolean, default=False) def __init__(self, chat_id, note_name, name, url, same_line=False): self.chat_id = str(chat_id) self.note_name = note_name self.name = name self.url = url self.same_line = same_line Notes.__table__.create(checkfirst=True) Buttons.__table__.create(checkfirst=True) NOTES_INSERTION_LOCK = threading.RLock() BUTTONS_INSERTION_LOCK = threading.RLock() def add_note_to_db(chat_id, note_name, note_data, msgtype, buttons=None, file=None): if not buttons: buttons = [] with NOTES_INSERTION_LOCK: prev = SESSION.query(Notes).get((str(chat_id), note_name)) if prev: with BUTTONS_INSERTION_LOCK: prev_buttons = ( SESSION.query(Buttons) .filter( Buttons.chat_id == str(chat_id), Buttons.note_name == note_name, ) .all() ) for btn in prev_buttons: SESSION.delete(btn) SESSION.delete(prev) note = Notes( str(chat_id), note_name, note_data or "", msgtype=msgtype.value, file=file, ) SESSION.add(note) SESSION.commit() for b_name, url, same_line in buttons: add_note_button_to_db(chat_id, note_name, b_name, url, same_line) def get_note(chat_id, note_name): try: return ( SESSION.query(Notes) .filter(func.lower(Notes.name) == note_name, Notes.chat_id == str(chat_id)) .first() ) finally: SESSION.close() def rm_note(chat_id, note_name): with NOTES_INSERTION_LOCK: note = ( SESSION.query(Notes) .filter(func.lower(Notes.name) == note_name, Notes.chat_id == str(chat_id)) .first() ) if note: with BUTTONS_INSERTION_LOCK: buttons = ( SESSION.query(Buttons) .filter( Buttons.chat_id == str(chat_id), Buttons.note_name == note_name, ) .all() ) for btn in buttons: SESSION.delete(btn) SESSION.delete(note) SESSION.commit() return True else: SESSION.close() return False def get_all_chat_notes(chat_id): try: return ( SESSION.query(Notes) .filter(Notes.chat_id == str(chat_id)) .order_by(Notes.name.asc()) .all() ) finally: SESSION.close() def add_note_button_to_db(chat_id, note_name, b_name, url, same_line): with BUTTONS_INSERTION_LOCK: button = Buttons(chat_id, note_name, b_name, url, same_line) SESSION.add(button) SESSION.commit() def get_buttons(chat_id, note_name): try: return ( SESSION.query(Buttons) .filter(Buttons.chat_id == str(chat_id), Buttons.note_name == note_name) .order_by(Buttons.id) .all() ) finally: SESSION.close() def num_notes(): try: return SESSION.query(Notes).count() finally: SESSION.close() def num_chats(): try: return SESSION.query(func.count(distinct(Notes.chat_id))).scalar() finally: SESSION.close() def migrate_chat(old_chat_id, new_chat_id): with NOTES_INSERTION_LOCK: chat_notes = ( SESSION.query(Notes).filter(Notes.chat_id == str(old_chat_id)).all() ) for note in chat_notes: note.chat_id = str(new_chat_id) with BUTTONS_INSERTION_LOCK: chat_buttons = ( SESSION.query(Buttons).filter(Buttons.chat_id == str(old_chat_id)).all() ) for btn in chat_buttons: btn.chat_id = str(new_chat_id) SESSION.commit()