feat: MySQL datasource dynamically adds user and dynamically reloads push targets support

This commit is contained in:
LWR 2023-02-16 21:38:22 +08:00
parent 5e882fa1c1
commit 063608a3ad
3 changed files with 156 additions and 63 deletions

View File

@ -91,6 +91,24 @@ class DataSource(metaclass=abc.ABCMeta):
raise DataSourceException(f"不存在的 UID: {uid}")
return up
def get_bot(self, qq: int) -> Bot:
"""
根据 QQ 获取 Bot 实例
Args:
qq: 需要获取 Bot QQ
Returns:
Bot 实例
Raises:
DataSourceException: QQ 不存在
"""
bot = next((b for b in self.bots if b.qq == qq), None)
if bot is None:
raise DataSourceException(f"不存在的 QQ: {qq}")
return bot
def get_ups_by_target(self, target_id: int, target_type: PushType) -> List[Up]:
"""
根据推送目标号码和推送目标类型获取 Up 实例列表
@ -272,30 +290,16 @@ class MySQLDataSource(DataSource):
except pymysql.err.Error as ex:
raise DataSourceException(f"从 MySQL 中读取配置时发生了错误 {ex}")
async def load(self):
async def __load_targets(self, uid: int) -> List[PushTarget]:
"""
MySQL 中初始化配置
MySQL 中读取指定 UID 的推送配置
Args:
uid: 要读取配置的 UID
Returns:
推送目标列表
"""
if self.bots:
return
logger.info("已选用 MySQL 作为 Bot 数据源")
logger.info("开始从 MySQL 中初始化 Bot 配置")
if not self.__pool:
await self.__connect()
users = await self.__query("SELECT * FROM `bot`")
bots = set(map(lambda u: u.get("bot"), users))
for bot in bots:
bot_users = list(filter(lambda u: u.get("bot") == bot, users))
ups = []
for now_user in bot_users:
uid = now_user.get("uid")
live_on = await self.__query(
"SELECT g.`uid`, g.`uname`, g.`room_id`, `key`, `type`, `num`, `enabled`, `message` "
"FROM `groups` AS `g` LEFT JOIN `live_on` AS `l` "
@ -360,9 +364,82 @@ class MySQLDataSource(DataSource):
)
)
return targets
async def load(self):
"""
MySQL 中初始化配置
"""
if self.bots:
return
logger.info("已选用 MySQL 作为 Bot 数据源")
logger.info("开始从 MySQL 中初始化 Bot 配置")
if not self.__pool:
await self.__connect()
users = await self.__query("SELECT * FROM `bot`")
bots = set(map(lambda u: u.get("bot"), users))
for bot in bots:
bot_users = list(filter(lambda u: u.get("bot") == bot, users))
ups = []
for now_user in bot_users:
uid = now_user.get("uid")
targets = await self.__load_targets(uid)
ups.append(Up(uid=uid, targets=targets))
self.bots.append(Bot(qq=bot, ups=ups))
super().format_data()
logger.success(f"成功从 MySQL 中导入了 {len(self.get_up_list())} 个 UP 主")
async def reload_targets(self, up: Union[int, Up]):
"""
重新从 MySQL 中读取特定 Up 的推送配置
Args:
up: 需要重载配置的 Up 实例或其 UID
"""
if isinstance(up, int):
try:
up = self.get_up(up)
except DataSourceException:
logger.warning(f"重载配置时出现异常, UID: {up} 不存在")
return
logger.info(f"开始从 MySQL 中重载 {up.uname} (UID: {up.uid}, 房间号: {up.room_id}) 的推送配置")
if not self.__pool:
await self.__connect()
up.targets = await self.__load_targets(up.uid)
super().format_data()
logger.success(f"已成功重载 {up.uname} (UID: {up.uid}, 房间号: {up.room_id}) 的推送配置")
async def load_new(self, uid: int):
"""
MySQL 中追加读取指定 UID 的用户
Args:
uid: 需要追加读取配置的 UID
"""
user = await self.__query(f"SELECT * FROM `bot` WHERE uid = {uid}")
if len(user) == 0:
logger.error(f"载入 UID: {uid} 的推送配置失败, UID 不存在")
raise DataSourceException(f"载入 UID: {uid} 的推送配置失败, UID 不存在")
bot = user[0].get("bot")
targets = await self.__load_targets(uid)
up = Up(uid=uid, targets=targets)
self.get_bot(bot).ups.append(up)
super().format_data()
logger.success(f"已成功载入 UID: {uid} 的推送配置")
await up.connect()

View File

@ -155,6 +155,14 @@ class Up(BaseModel):
else:
logger.success(f"已成功连接到 {self.uname} 的直播间 {self.room_id}")
if not await redis.exists_live_status(self.room_id):
room_info = await self.__live_room.get_room_play_info()
status = room_info["live_status"]
await redis.set_live_status(self.room_id, status)
if status == 1:
start_time = room_info["live_time"]
await redis.set_live_start_time(self.room_id, start_time)
self.__is_reconnect = True
@self.__room.on("LIVE")

View File

@ -207,6 +207,10 @@ async def zunionstore(dest: str, source: Union[str, List[str]]):
# 直播间状态0未开播1正在直播2轮播
async def exists_live_status(room_id: int) -> bool:
return await hexists("LiveStatus", room_id)
async def get_live_status(room_id: int) -> int:
return await hgeti("LiveStatus", room_id)
@ -217,6 +221,10 @@ async def set_live_status(room_id: int, status: int):
# 直播开始时间
async def exists_live_start_time(room_id: int) -> bool:
return await hexists("StartTime", room_id)
async def get_live_start_time(room_id: int) -> int:
return await hgeti("StartTime", room_id)