from db.models import Bot, Session_Publisher from db.models import Session as SessionModel from db.query.base_query import BaseQuery class BotQuery(BaseQuery): def __init__(self, user): super().__init__(user) def add_bot(self, db, session_id, bot, metadata_id, type="bot_one"): # Create the new session if type == "bot_one": new_bot_one = Session_Publisher( id=session_id, user_id=self.user.get("id"), bot_name=bot, metadata_id=metadata_id, ) else: new_bot_one = SessionModel( id=session_id, user_id=self.user.get("id"), bot_id=bot, metadata_id=metadata_id, ) return self.add(db, new_bot_one) def get_session_ids_bot(self, db, metadata_id, type="bot_one"): model = Session_Publisher columns = [ Session_Publisher.id, Session_Publisher.bot_name, Session_Publisher.updated_at, ] filter_conditions = [ Session_Publisher.user_id == self.user.get("id"), Session_Publisher.metadata_id == metadata_id, ] sessions = self.get_columns( db, model=model, columns=columns, filter_conditions=filter_conditions, multiple=True, ) return sessions def update_session_bot(self, db, bot_id, name): update_data = {"name": name} self.update(db, Bot, bot_id, update_data) def delete_session_bot(self, db, bot_id): self.delete(db, Bot, bot_id) def get_bot(self, db, bot_id): return self.get(db, Bot, bot_id) # def get_all_bots(self, db): # return self.get_all(db, Bot)