db/query/query_bot.py

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)