import ujson from model.base import ServiceEnum from utils.error import NotFoundError from utils.mysql import MySQL class CookiesRepository: def __init__(self, mysql: MySQL): self.mysql = mysql async def update_cookie(self, user_id: int, cookies: str, default_service: ServiceEnum): if default_service == ServiceEnum.HYPERION: query = """ UPDATE `mihoyo_cookie` SET cookie=%s WHERE user_id=%s; """ elif default_service == ServiceEnum.HOYOLAB: query = """ UPDATE `hoyoverse_cookie` SET cookie=%s WHERE user_id=%s; """ else: raise DefaultServiceNotFoundError(default_service.name) query_args = (cookies, user_id) await self.mysql.execute_and_fetchall(query, query_args) async def set_cookie(self, user_id: int, cookies: str, default_service: ServiceEnum): if default_service == ServiceEnum.HYPERION: query = """ INSERT INTO `mihoyo_cookie` (user_id,cookie) VALUES (%s,%s) ON DUPLICATE KEY UPDATE cookie=VALUES(cookie); """ elif default_service == ServiceEnum.HOYOLAB: query = """ INSERT INTO `hoyoverse_cookie` (user_id,cookie) VALUES (%s,%s) ON DUPLICATE KEY UPDATE cookie=VALUES(cookie); """ else: raise DefaultServiceNotFoundError(default_service.name) query_args = (user_id, cookies) await self.mysql.execute_and_fetchall(query, query_args) async def read_cookies(self, user_id, default_service: ServiceEnum) -> dict: if default_service == ServiceEnum.HYPERION: query = """ SELECT cookie FROM `mihoyo_cookie` WHERE user_id=%s; """ elif default_service == ServiceEnum.HOYOLAB: query = """ SELECT cookie FROM `hoyoverse_cookie` WHERE user_id=%s;; """ else: raise DefaultServiceNotFoundError(default_service.name) query_args = (user_id,) data = await self.mysql.execute_and_fetchall(query, query_args) if len(data) == 0: return {} (cookies,) = data return ujson.loads(cookies) class DefaultServiceNotFoundError(NotFoundError): entity_name: str = "ServiceEnum" entity_value_name: str = "default_service"