feature: 数据库连接与SQL集中管理，提高代码可读性
This commit is contained in:
liuwei
+39
-100
@@ -4,14 +4,17 @@ import mysql.connector.pooling
|
||||
import tomllib
|
||||
import pytz
|
||||
import redis
|
||||
from typing import Optional, Tuple
|
||||
from typing import Optional
|
||||
|
||||
from wcferry import Wcf, WxMsg
|
||||
|
||||
from message_util import MessageUtil
|
||||
from robot_cmd.robot_command import GroupBotManager, Feature, PermissionStatus
|
||||
from db.connection import DBConnectionManager
|
||||
from db.sign_in import SignInDB
|
||||
from db.sign_in_redis import SignInRedisDB
|
||||
|
||||
# 创建表的SQL语句
|
||||
# 创建表的SQL语句保留在这里，用于初始化表
|
||||
CREATE_TABLE_SQL = """
|
||||
CREATE TABLE IF NOT EXISTS t_sign_record (
|
||||
id BIGINT AUTO_INCREMENT PRIMARY KEY,
|
||||
@@ -45,8 +48,16 @@ class SignInSystem:
|
||||
self.gbm = gbm
|
||||
self.message_util = message_util
|
||||
self.all_contacts = all_contacts
|
||||
self.db_pool = db_pool
|
||||
self.redis_pool = redis_pool
|
||||
|
||||
# 初始化数据库连接管理器
|
||||
self.db_manager = DBConnectionManager()
|
||||
self.db_manager.mysql_pool = db_pool
|
||||
self.db_manager.redis_pool = redis_pool
|
||||
|
||||
# 初始化数据库操作类
|
||||
self.sign_in_db = SignInDB(self.db_manager)
|
||||
self.sign_in_redis = SignInRedisDB(self.db_manager)
|
||||
|
||||
self.command = self.config['command']
|
||||
self.min_point = self.config['min-point']
|
||||
self.max_point = self.config['max-point']
|
||||
@@ -56,49 +67,16 @@ class SignInSystem:
|
||||
self.timezone = 'Asia/Shanghai'
|
||||
|
||||
# 从 Redis 初始化签到数据
|
||||
self.today_signin_count = self._load_signin_count_from_redis()
|
||||
with self._get_redis_connection() as redis_client:
|
||||
last_reset_date_str = redis_client.get('group:sign_in:last_reset_date')
|
||||
if last_reset_date_str:
|
||||
self.last_reset_date = datetime.strptime(last_reset_date_str, '%Y-%m-%d').date()
|
||||
else:
|
||||
self.last_reset_date = datetime.now(tz=pytz.timezone(self.timezone)).date()
|
||||
self._save_last_reset_date_to_redis()
|
||||
self.today_signin_count = self.sign_in_redis.load_signin_count()
|
||||
last_reset_date = self.sign_in_redis.get_last_reset_date()
|
||||
if last_reset_date:
|
||||
self.last_reset_date = last_reset_date
|
||||
else:
|
||||
self.last_reset_date = datetime.now(tz=pytz.timezone(self.timezone)).date()
|
||||
self.sign_in_redis.save_last_reset_date(self.last_reset_date)
|
||||
|
||||
self.LOG.info(f"[签到] 组件初始化完成 {self.command_format}")
|
||||
|
||||
def _get_db_connection(self):
|
||||
"""从连接池获取数据库连接"""
|
||||
return self.db_pool.get_connection()
|
||||
|
||||
def _get_redis_connection(self):
|
||||
"""从连接池获取 Redis 连接"""
|
||||
return redis.Redis(connection_pool=self.redis_pool)
|
||||
|
||||
def _load_signin_count_from_redis(self) -> dict:
|
||||
"""从 Redis 加载签到人数数据"""
|
||||
signin_count = {}
|
||||
with self._get_redis_connection() as redis_client:
|
||||
keys = redis_client.keys('group:sign_in:*')
|
||||
for key in keys:
|
||||
if key == 'group:sign_in:last_reset_date':
|
||||
continue
|
||||
group_id = key.replace('group:sign_in:', '')
|
||||
count = redis_client.get(key)
|
||||
if count is not None:
|
||||
signin_count[group_id] = int(count)
|
||||
return signin_count
|
||||
|
||||
def _save_signin_count_to_redis(self, group_id: str, count: int):
|
||||
"""保存签到人数到 Redis"""
|
||||
with self._get_redis_connection() as redis_client:
|
||||
redis_client.set(f'group:sign_in:{group_id}', count)
|
||||
|
||||
def _save_last_reset_date_to_redis(self):
|
||||
"""保存最后重置日期到 Redis"""
|
||||
with self._get_redis_connection() as redis_client:
|
||||
redis_client.set('group:sign_in:last_reset_date', self.last_reset_date.strftime('%Y-%m-%d'))
|
||||
|
||||
@property
|
||||
def command_format(self):
|
||||
return ','.join(self.command)
|
||||
@@ -109,23 +87,16 @@ class SignInSystem:
|
||||
|
||||
def initialize_table(self):
|
||||
"""初始化数据库表"""
|
||||
with self._get_db_connection() as conn:
|
||||
with conn.cursor(dictionary=True) as cursor: # 使用 dictionary=True 返回字典格式
|
||||
cursor.execute(CREATE_TABLE_SQL)
|
||||
conn.commit()
|
||||
self.sign_in_db.execute_update(CREATE_TABLE_SQL)
|
||||
|
||||
def reset_today_count_if_needed(self):
|
||||
"""检查并重置每日签到计数"""
|
||||
current_date = datetime.now(tz=pytz.timezone(self.timezone)).date()
|
||||
if current_date != self.last_reset_date:
|
||||
self.today_signin_count.clear()
|
||||
with self._get_redis_connection() as redis_client:
|
||||
keys = redis_client.keys('group:sign_in:*')
|
||||
for key in keys:
|
||||
if key != 'group:sign_in:last_reset_date':
|
||||
redis_client.delete(key)
|
||||
self.sign_in_redis.reset_daily_counts()
|
||||
self.last_reset_date = current_date
|
||||
self._save_last_reset_date_to_redis()
|
||||
self.sign_in_redis.save_last_reset_date(self.last_reset_date)
|
||||
self.LOG.info(f"[签到] 已重置每日签到计数，日期更新为 {current_date}")
|
||||
|
||||
def get_today_signin_count(self, group_id: str) -> int:
|
||||
@@ -135,15 +106,7 @@ class SignInSystem:
|
||||
|
||||
def get_user_record(self, wx_id: str, group_id: str) -> Optional[dict]:
|
||||
"""获取用户签到记录"""
|
||||
with self._get_db_connection() as conn:
|
||||
with conn.cursor(dictionary=True) as cursor:
|
||||
query = """
|
||||
SELECT wx_id, group_id, wx_nick_name, points, sign_stat, signin_streak
|
||||
FROM t_sign_record
|
||||
WHERE wx_id = %s AND group_id = %s
|
||||
"""
|
||||
cursor.execute(query, (wx_id, group_id))
|
||||
return cursor.fetchone()
|
||||
return self.sign_in_db.get_user_record(wx_id, group_id)
|
||||
|
||||
def calculate_points(self, streak: int) -> int:
|
||||
"""根据连续签到天数计算积分"""
|
||||
@@ -214,46 +177,22 @@ class SignInSystem:
|
||||
|
||||
today_signin_rank = self.get_today_signin_count(message.roomid) + 1
|
||||
self.today_signin_count[message.roomid] = today_signin_rank
|
||||
self._save_signin_count_to_redis(message.roomid, today_signin_rank)
|
||||
self.sign_in_redis.save_signin_count(message.roomid, today_signin_rank)
|
||||
|
||||
points_to_add = self.calculate_points(streak)
|
||||
|
||||
with self._get_db_connection() as conn:
|
||||
with conn.cursor(dictionary=True) as cursor:
|
||||
if user_record:
|
||||
update_sql = """
|
||||
UPDATE t_sign_record
|
||||
SET wx_nick_name = %s, points = points + %s,
|
||||
sign_stat = %s, signin_streak = %s,
|
||||
update_time = %s
|
||||
WHERE wx_id = %s AND group_id = %s
|
||||
"""
|
||||
cursor.execute(update_sql, (
|
||||
wx_nick_name, points_to_add, current_time, streak,
|
||||
current_time, message.sender, message.roomid
|
||||
))
|
||||
else:
|
||||
insert_sql = """
|
||||
INSERT INTO t_sign_record
|
||||
(wx_id, group_id, wx_nick_name, points, sign_stat, signin_streak)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
"""
|
||||
cursor.execute(insert_sql, (
|
||||
message.sender, message.roomid, wx_nick_name, points_to_add, current_time, streak
|
||||
))
|
||||
conn.commit()
|
||||
# output = ("\n"
|
||||
# f"-----Bot-----\n"
|
||||
# f"@{wx_nick_name} 签到成功！\n"
|
||||
# f"签到成功！你领到了 {points_to_add} 个积分！✅\n"
|
||||
# f"你是今天第 {today_signin_rank} 个签到的！🎉\n")
|
||||
#
|
||||
# if streak_broken and old_streak > 0: # 只有在真的断签且之前有签到记录时才显示
|
||||
# output += f"你断开了 {old_streak} 天的连续签到！[心碎]"
|
||||
# elif streak > 1:
|
||||
# output += f"你连续签到了 {streak} 天！"
|
||||
# if streak > 1 and not streak_broken:
|
||||
# output += "[爱心]"
|
||||
# 使用数据库操作类更新或创建签到记录
|
||||
if user_record:
|
||||
self.sign_in_db.update_sign_record(
|
||||
message.sender, message.roomid, wx_nick_name,
|
||||
points_to_add, current_time, streak
|
||||
)
|
||||
else:
|
||||
self.sign_in_db.create_sign_record(
|
||||
message.sender, message.roomid, wx_nick_name,
|
||||
points_to_add, current_time, streak
|
||||
)
|
||||
|
||||
output = f"签到成功，加[{points_to_add}]分，第[{today_signin_rank}]个！"
|
||||
|
||||
if streak_broken and old_streak > 0: # 只有在真的断签且之前有签到记录时才显示
|
||||
|
||||
Reference in New Issue
Block a user