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 get_current_welcome_settings(chat_id):
try:
return SESSION.query(Welcome).filter(Welcome.chat_id == chat_id).one()
except:
return None
finally:
SESSION.close()
def add_note_to_db(chat_id, note_name, note_message_id):
with NOTES_INSERTION_LOCK:
prev = SESSION.query(Notes).get((str(chat_id), note_name))
if prev:
SESSION.delete(prev)
note = Notes(str(chat_id), note_name, note_message_id)
SESSION.add(note)
SESSION.commit()
def __load_flood_settings():
global CHAT_FLOOD
try:
all_chats = SESSION.query(FloodControl).all()
CHAT_FLOOD = {chat.chat_id: (None, DEF_COUNT, chat.limit) for chat in all_chats}
finally:
SESSION.close()
return CHAT_FLOOD
def migrate_chat(old_chat_id, new_chat_id):
with INSERTION_LOCK:
flood = SESSION.query(FloodControl).get(str(old_chat_id))
if flood:
CHAT_FLOOD[str(new_chat_id)] = CHAT_FLOOD.get(str(old_chat_id), DEF_OBJ)
flood.chat_id = str(new_chat_id)
SESSION.commit()
SESSION.close()
def get_credential(chat_id):
with INSERTION_LOCK:
saved_cred = SESSION.query(gDriveCreds).get(chat_id)
creds = None
# The gDrive table stores the user's access and refresh tokens, and is
# created automatically when the authorization flow
# completes for the first time.
if saved_cred is not None:
creds = pickle.loads(saved_cred.credential_string)
return creds
def set_credential(chat_id, credential_string):
with INSERTION_LOCK:
saved_cred = SESSION.query(gDriveCreds).get(chat_id)
if not saved_cred:
saved_cred = gDriveCreds(chat_id)
saved_cred.credential_string = pickle.dumps(credential_string)
SESSION.add(saved_cred)
SESSION.commit()
def rm_note(chat_id, note_name):
with NOTES_INSERTION_LOCK:
note = SESSION.query(Notes).get((str(chat_id), note_name))
if note:
SESSION.delete(note)
SESSION.commit()
return True
else:
SESSION.close()
return False
def add_welcome_setting(chat_id, should_clean_welcome, previous_welcome, f_mesg_id):
adder = SESSION.query(Welcome).get(chat_id)
if adder:
adder.should_clean_welcome = should_clean_welcome
adder.previous_welcome = previous_welcome
adder.f_mesg_id = f_mesg_id
else:
adder = Welcome(chat_id, should_clean_welcome, previous_welcome, f_mesg_id)
SESSION.add(adder)
SESSION.commit()
def set_flood(chat_id, amount):
with INSERTION_LOCK:
flood = SESSION.query(FloodControl).get(str(chat_id))
if not flood:
flood = FloodControl(str(chat_id))
flood.user_id = None
flood.limit = amount
CHAT_FLOOD[str(chat_id)] = (None, DEF_COUNT, amount)
SESSION.add(flood)
SESSION.commit()
def clear_credential(chat_id):
with INSERTION_LOCK:
saved_cred = SESSION.query(gDriveCreds).get(chat_id)
if saved_cred:
SESSION.delete(saved_cred)
SESSION.commit()
def update_previous_welcome(chat_id, previous_welcome):
row = SESSION.query(Welcome).get(chat_id)
row.previous_welcome = previous_welcome
# commit the changes to the DB
SESSION.commit()
def rm_welcome_setting(chat_id):
rem = SESSION.query(Welcome).get(chat_id)
if rem:
SESSION.delete(rem)
SESSION.commit()
def num_chats():
try:
return SESSION.query(func.count(distinct(Notes.chat_id))).scalar()
finally:
SESSION.close()
def num_notes():
try:
return SESSION.query(Notes).count()
finally:
SESSION.close()
def get_note(chat_id, note_name):
try:
return SESSION.query(Notes).get((str(chat_id), note_name))
finally:
SESSION.close()