def set_nsfw(chat_id):
with INSERTION_LOCK:
nsfwchat = SESSION.query(NSFWChats).get(str(chat_id))
if not nsfwchat:
nsfwchat = NSFWChats(str(chat_id))
SESSION.add(nsfwchat)
SESSION.commit()
def get_user_num_chats(user_id):
try:
return (
SESSION.query(ChatMembers).filter(ChatMembers.user == int(user_id)).count()
)
finally:
SESSION.close()
def get_userid_by_name(username):
try:
return (
SESSION.query(Users)
.filter(func.lower(Users.username) == username.lower())
.all()
)
finally:
SESSION.close()
def is_nsfw(chat_id):
try:
chat = SESSION.query(NSFWChats).get(str(chat_id))
if chat:
return True
else:
return False
finally:
SESSION.close()
def del_user(user_id):
with INSERTION_LOCK:
curr = SESSION.query(Users).get(user_id)
if curr:
SESSION.delete(curr)
SESSION.commit()
return True
ChatMembers.query.filter(ChatMembers.user == user_id).delete()
SESSION.commit()
SESSION.close()
return False
def rem_chat(chat_id):
with INSERTION_LOCK:
chat = SESSION.query(Chats).get(str(chat_id))
if chat:
SESSION.delete(chat)
SESSION.commit()
else:
SESSION.close()
def migrate_chat(old_chat_id, new_chat_id):
with INSERTION_LOCK:
chat = SESSION.query(Chats).get(str(old_chat_id))
if chat:
chat.chat_id = str(new_chat_id)
SESSION.add(chat)
SESSION.flush()
chat_members = (
SESSION.query(ChatMembers)
.filter(ChatMembers.chat == str(old_chat_id))
.all()
)
for member in chat_members:
member.chat = str(new_chat_id)
SESSION.add(member)
SESSION.commit()
def update_user(user_id, username, chat_id=None, chat_name=None):
with INSERTION_LOCK:
user = SESSION.query(Users).get(user_id)
if not user:
user = Users(user_id, username)
SESSION.add(user)
SESSION.flush()
else:
user.username = username
if not chat_id or not chat_name:
SESSION.commit()
return
chat = SESSION.query(Chats).get(str(chat_id))
if not chat:
chat = Chats(str(chat_id), chat_name)
SESSION.add(chat)
SESSION.flush()
else:
chat.chat_name = chat_name
member = (
SESSION.query(ChatMembers)
.filter(ChatMembers.chat == chat.chat_id, ChatMembers.user == user.user_id)
.first()
)
if not member:
chat_member = ChatMembers(chat.chat_id, user.user_id)
SESSION.add(chat_member)
SESSION.commit()
def ensure_bot_in_db():
with INSERTION_LOCK:
bot = Users(dispatcher.bot.id, dispatcher.bot.username)
SESSION.merge(bot)
SESSION.commit()
def num_users():
try:
return SESSION.query(Users).count()
finally:
SESSION.close()
def num_chats():
try:
return SESSION.query(Chats).count()
finally:
SESSION.close()
def get_all_chats():
try:
return SESSION.query(Chats).all()
finally:
SESSION.close()
def get_chat_members(chat_id):
try:
return SESSION.query(ChatMembers).filter(ChatMembers.chat == str(chat_id)).all()
finally:
SESSION.close()
def get_name_by_userid(user_id):
try:
return SESSION.query(Users).get(Users.user_id == int(user_id)).first()
finally:
SESSION.close()
def get_all_nsfw_chats():
try:
return SESSION.query(NSFWChats.chat_id).all()
finally:
SESSION.close()
def rem_nsfw(chat_id):
with INSERTION_LOCK:
nsfwchat = SESSION.query(NSFWChats).get(str(chat_id))
if nsfwchat:
SESSION.delete(nsfwchat)
SESSION.commit()