AgentCoreService/pqai_agent_server/models/mysql_staff_manager.py

import abc
import json
import pymysql.cursors

from typing import Dict, List
from pqai_agent.database import MySQLManager


class StaffManager(abc.ABC):

    @abc.abstractmethod
    def list_all_staffs(self, page_id: int, page_size: int) -> List[Dict]:
        pass

    @abc.abstractmethod
    def get_staff_profile(self, staff_id) -> Dict:
        pass

    @abc.abstractmethod
    def save_staff_profile(self, staff_id: str, staff_profile: Dict):
        pass


class MySQLStaffManager(StaffManager):

    def __init__(self, db_config, staff_table, user_table):
        self.db = MySQLManager(db_config)
        self.staff_table = staff_table
        self.user_table = user_table

    def save_staff_profile(self, staff_id: str, staff_profile: Dict):
        update_query = f"""
            update {self.staff_table} set agent_profile = %s where third_party_user_id = %s;
        """
        affected_rows = self.db.execute(
            update_query, (json.dumps(staff_profile), staff_id)
        )
        return affected_rows

    def get_staff_profile(self, staff_id) -> Dict:
        empty_profile = {}
        sql = f"""select agent_profile from {self.staff_table} where third_party_user_id = %s;"""
        response = self.db.select(
            sql=sql,
            cursor_type=pymysql.cursors.DictCursor,
            args=(staff_id,),
        )
        if not response:
            return empty_profile

        agent_profile = response[0]["agent_profile"]
        profile = json.loads(agent_profile)
        if not profile:
            return empty_profile
        return profile

    def list_all_staffs(self, page_id: int, page_size: int) -> Dict:
        """
        :param page_size:
        :param page_id:
        :return:
        """
        sql = f"""
            select t1.third_party_user_id as staff_id, t1.name as staff_name, t2.iconurl as avatar
            from {self.staff_table} t1 left join {self.user_table} t2
            on t1.third_party_user_id = t2.third_party_user_id
            limit %s offset %s;
        """
        staff_list = self.db.select(
            sql=sql,
            cursor_type=pymysql.cursors.DictCursor,
            args=(page_size + 1, page_size * (page_id - 1)),
        )
        if len(staff_list) > page_size:
            has_next_page = True
            next_page_id = page_id + 1
            staff_list = staff_list[:page_size]
        else:
            has_next_page = False
            next_page_id = None
        return {
            "has_next_page": has_next_page,
            "next_page": next_page_id,
            "data": staff_list,
        }