Files
ai_mip/db_manager.py

548 lines
19 KiB
Python
Raw Normal View History

2026-01-16 22:06:46 +08:00
"""
数据库管理器
使用 MySQL 数据库
"""
import json
import random
from datetime import datetime
from typing import List, Dict, Optional, Union
from pathlib import Path
from loguru import logger
from config import Config
try:
import pymysql
MYSQL_AVAILABLE = True
except ImportError:
MYSQL_AVAILABLE = False
logger.error("pymysql 未安装,请安装: pip install pymysql")
raise ImportError("使用 MySQL 需要安装 pymysql: pip install pymysql")
class DatabaseManager:
"""数据库管理器,使用 MySQL"""
def __init__(self, db_path: str = None):
"""
初始化数据库连接
Args:
db_path: 忽略仅为了兼容性
"""
if not MYSQL_AVAILABLE:
raise ImportError("使用 MySQL 需要安装 pymysql: pip install pymysql")
self.db_config = {
'host': Config.MYSQL_HOST,
'port': Config.MYSQL_PORT,
'user': Config.MYSQL_USER,
'password': Config.MYSQL_PASSWORD,
'database': Config.MYSQL_DATABASE,
'charset': 'utf8mb4'
}
logger.info(f"MySQL数据库初始化: {Config.MYSQL_HOST}:{Config.MYSQL_PORT}/{Config.MYSQL_DATABASE}")
def get_connection(self) -> 'pymysql.Connection':
"""获取MySQL数据库连接"""
conn = pymysql.connect(**self.db_config)
return conn
def _dict_from_row(self, row) -> Dict:
"""将数据库行转换为字典"""
if row is None:
return None
return dict(row) if isinstance(row, dict) else row
def _get_placeholder(self) -> str:
"""获取SQL占位符MySQL使用 %s"""
return '%s'
def _execute_query(self, conn, sql: str, params: tuple = None):
"""执行SQL查询使用DictCursor"""
cursor = conn.cursor(pymysql.cursors.DictCursor)
if params:
cursor.execute(sql, params)
else:
cursor.execute(sql)
return cursor
class SiteManager(DatabaseManager):
"""站点管理"""
def add_site(self, site_url: str, site_name: str = None,
site_dimension: str = None, frequency: int = None,
time_start: str = None, time_end: str = None,
interval_minutes: int = None) -> Optional[int]:
"""
添加新站点
Args:
site_url: 网站URL
site_name: 网站名称
site_dimension: 网站维度标签
frequency: 频次
time_start: 开始时间
time_end: 结束时间
interval_minutes: 执行间隔分钟
Returns:
站点ID失败返回None
"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = conn.cursor(pymysql.cursors.DictCursor)
# 生成随机目标点击次数(兼容原有逻辑)
target_clicks = random.randint(
getattr(Config, 'MIN_CLICK_COUNT', 1),
getattr(Config, 'MAX_CLICK_COUNT', 10)
)
sql = f"""
INSERT INTO ai_mip_site (
site_url, site_name, status, frequency,
time_start, time_end, interval_minutes,
site_dimension, created_by
) VALUES ({ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph})
"""
cursor.execute(sql, (
site_url,
site_name or site_url,
'active',
frequency or 1,
time_start or '09:00:00',
time_end or '21:00:00',
interval_minutes or 60,
site_dimension,
'system'
))
site_id = cursor.lastrowid
conn.commit()
conn.close()
logger.info(f"成功添加站点: {site_url} (ID: {site_id})")
return site_id
except pymysql.IntegrityError:
logger.warning(f"站点URL已存在: {site_url}")
return None
except Exception as e:
logger.error(f"添加站点失败: {str(e)}")
return None
def get_site_by_url(self, site_url: str) -> Optional[Dict]:
"""根据URL获取站点信息"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = self._execute_query(conn, f"SELECT * FROM ai_mip_site WHERE site_url = {ph}", (site_url,))
row = cursor.fetchone()
conn.close()
return self._dict_from_row(row) if row else None
except Exception as e:
logger.error(f"查询站点失败: {str(e)}")
return None
def get_site_by_id(self, site_id: int) -> Optional[Dict]:
"""根据ID获取站点信息"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = self._execute_query(conn, f"SELECT * FROM ai_mip_site WHERE id = {ph}", (site_id,))
row = cursor.fetchone()
conn.close()
return self._dict_from_row(row) if row else None
except Exception as e:
logger.error(f"查询站点失败: {str(e)}")
return None
def get_active_sites(self) -> List[Dict]:
"""获取所有活跃站点"""
try:
conn = self.get_connection()
cursor = self._execute_query(conn, "SELECT * FROM ai_mip_site WHERE status = 'active' ORDER BY created_at DESC")
rows = cursor.fetchall()
conn.close()
return [self._dict_from_row(row) for row in rows]
except Exception as e:
logger.error(f"查询活跃站点失败: {str(e)}")
return []
def get_all_sites(self) -> List[Dict]:
"""获取所有站点"""
try:
conn = self.get_connection()
cursor = self._execute_query(conn, "SELECT * FROM ai_mip_site ORDER BY created_at DESC")
rows = cursor.fetchall()
conn.close()
return [self._dict_from_row(row) for row in rows]
except Exception as e:
logger.error(f"查询所有站点失败: {str(e)}")
return []
def update_site_status(self, site_id: int, status: str) -> bool:
"""更新站点状态"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = self._execute_query(
conn,
f"UPDATE ai_mip_site SET status = {ph}, updated_by = {ph} WHERE id = {ph}",
(status, 'system', site_id)
)
conn.commit()
conn.close()
logger.info(f"更新站点状态: ID={site_id}, status={status}")
return True
except Exception as e:
logger.error(f"更新站点状态失败: {str(e)}")
return False
def increment_click_count(self, site_id: int, count: int = 1) -> bool:
"""增加点击次数"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = self._execute_query(
conn,
f"UPDATE ai_mip_site SET click_count = click_count + {ph} WHERE id = {ph}",
(count, site_id)
)
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"更新点击次数失败: {str(e)}")
return False
def increment_reply_count(self, site_id: int, count: int = 1) -> bool:
"""增加回复次数"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = self._execute_query(
conn,
f"UPDATE ai_mip_site SET reply_count = reply_count + {ph} WHERE id = {ph}",
(count, site_id)
)
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"更新回复次数失败: {str(e)}")
return False
def delete_site(self, site_id: int) -> bool:
"""删除站点"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = self._execute_query(conn, f"DELETE FROM ai_mip_site WHERE id = {ph}", (site_id,))
conn.commit()
conn.close()
logger.info(f"已删除站点: ID={site_id}")
return True
except Exception as e:
logger.error(f"删除站点失败: {str(e)}")
return False
class ClickManager(DatabaseManager):
"""点击记录管理"""
def record_click(self, site_id: int, site_url: str,
user_ip: str = None, device_type: str = 'pc',
task_id: str = None) -> Optional[int]:
"""
记录一次点击
Args:
site_id: 站点ID
site_url: 站点URL
user_ip: 用户IP代理IP
device_type: 设备类型
task_id: 任务ID
Returns:
点击记录ID
"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = conn.cursor(pymysql.cursors.DictCursor)
sql = f"""
INSERT INTO ai_mip_click (
site_id, site_url, click_time, user_ip,
device_type, task_id, operator
) VALUES ({ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph})
"""
cursor.execute(sql, (
site_id,
site_url,
datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
user_ip,
device_type,
task_id or f'TASK_{datetime.now().strftime("%Y%m%d%H%M%S")}',
'RPA_SYSTEM'
))
click_id = cursor.lastrowid
conn.commit()
conn.close()
# 更新站点点击次数
site_mgr = SiteManager()
site_mgr.increment_click_count(site_id)
logger.info(f"记录点击: site_id={site_id}, click_id={click_id}")
return click_id
except Exception as e:
logger.error(f"记录点击失败: {str(e)}")
return None
def get_clicks_by_site(self, site_id: int, limit: int = 100) -> List[Dict]:
"""获取站点的点击记录"""
try:
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT * FROM ai_mip_click
WHERE site_id = ?
ORDER BY click_time DESC
LIMIT ?
""", (site_id, limit))
rows = cursor.fetchall()
conn.close()
return [self._dict_from_row(row) for row in rows]
except Exception as e:
logger.error(f"查询点击记录失败: {str(e)}")
return []
def get_click_count_by_site(self, site_id: int) -> int:
"""获取站点的总点击次数"""
try:
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) as count FROM ai_mip_click WHERE site_id = ?", (site_id,))
row = cursor.fetchone()
conn.close()
return row['count'] if row else 0
except Exception as e:
logger.error(f"查询点击次数失败: {str(e)}")
return 0
class InteractionManager(DatabaseManager):
"""互动记录管理"""
def record_interaction(self, site_id: int, click_id: int = None,
task_id: str = None, interaction_type: str = 'reply',
reply_content: str = None, is_successful: bool = False,
response_received: bool = False, response_content: str = None,
proxy_ip: str = None, fingerprint_id: str = None,
error_message: str = None) -> Optional[int]:
"""
记录一次互动
Args:
site_id: 站点ID
click_id: 关联的点击记录ID
task_id: 任务ID
interaction_type: 互动类型reply/comment等
reply_content: 回复内容
is_successful: 是否成功
response_received: 是否收到回复
response_content: 对方回复内容
proxy_ip: 使用的代理IP
fingerprint_id: 浏览器指纹ID
error_message: 错误信息
Returns:
互动记录ID
"""
try:
conn = self.get_connection()
ph = self._get_placeholder()
cursor = conn.cursor(pymysql.cursors.DictCursor)
now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
sql = f"""
INSERT INTO ai_mip_interaction (
site_id, click_id, task_id, interaction_type,
interaction_time, interaction_status, reply_content,
execution_mode, browser_type, proxy_ip, fingerprint_id,
response_received, response_content, is_successful,
error_message, operator
) VALUES ({ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph}, {ph})
"""
cursor.execute(sql, (
site_id,
click_id,
task_id or f'TASK_{datetime.now().strftime("%Y%m%d%H%M%S")}',
interaction_type,
now,
'success' if is_successful else 'failed',
reply_content,
'auto',
'playwright',
proxy_ip,
fingerprint_id,
1 if response_received else 0,
response_content,
1 if is_successful else 0,
error_message,
'RPA_SYSTEM'
))
interaction_id = cursor.lastrowid
conn.commit()
conn.close()
# 如果收到回复,更新站点回复次数
if response_received:
site_mgr = SiteManager()
site_mgr.increment_reply_count(site_id)
logger.info(f"记录互动: site_id={site_id}, interaction_id={interaction_id}, success={is_successful}")
return interaction_id
except Exception as e:
logger.error(f"记录互动失败: {str(e)}")
return None
def get_interactions_by_site(self, site_id: int, limit: int = 100) -> List[Dict]:
"""获取站点的互动记录"""
try:
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT * FROM ai_mip_interaction
WHERE site_id = ?
ORDER BY interaction_time DESC
LIMIT ?
""", (site_id, limit))
rows = cursor.fetchall()
conn.close()
return [self._dict_from_row(row) for row in rows]
except Exception as e:
logger.error(f"查询互动记录失败: {str(e)}")
return []
def get_successful_interactions_count(self, site_id: int) -> int:
"""获取站点的成功互动次数"""
try:
conn = self.get_connection()
cursor = conn.cursor()
cursor.execute("""
SELECT COUNT(*) as count FROM ai_mip_interaction
WHERE site_id = ? AND is_successful = 1
""", (site_id,))
row = cursor.fetchone()
conn.close()
return row['count'] if row else 0
except Exception as e:
logger.error(f"查询成功互动次数失败: {str(e)}")
return 0
class StatisticsManager(DatabaseManager):
"""统计数据管理"""
def get_statistics(self) -> Dict:
"""
获取全局统计数据
Returns:
统计数据字典
"""
try:
conn = self.get_connection()
# 站点统计
cursor = self._execute_query(conn, "SELECT COUNT(*) as total FROM ai_mip_site")
total_sites = cursor.fetchone()['total']
cursor = self._execute_query(conn, "SELECT COUNT(*) as total FROM ai_mip_site WHERE status = 'active'")
active_sites = cursor.fetchone()['total']
# 点击统计
cursor = self._execute_query(conn, "SELECT COUNT(*) as total FROM ai_mip_click")
total_clicks = cursor.fetchone()['total']
# 互动统计
cursor = self._execute_query(conn, "SELECT COUNT(*) as total FROM ai_mip_interaction WHERE response_received = 1")
total_replies = cursor.fetchone()['total']
cursor = self._execute_query(conn, "SELECT COUNT(*) as total FROM ai_mip_interaction WHERE is_successful = 1")
successful_interactions = cursor.fetchone()['total']
conn.close()
reply_rate = (total_replies / total_clicks * 100) if total_clicks > 0 else 0
success_rate = (successful_interactions / total_clicks * 100) if total_clicks > 0 else 0
return {
'total_sites': total_sites,
'active_sites': active_sites,
'total_clicks': total_clicks,
'total_replies': total_replies,
'successful_interactions': successful_interactions,
'reply_rate': f"{reply_rate:.2f}%",
'success_rate': f"{success_rate:.2f}%"
}
except Exception as e:
logger.error(f"获取统计数据失败: {str(e)}")
return {}
def get_site_statistics(self, site_id: int) -> Dict:
"""
获取单个站点的统计数据
Args:
site_id: 站点ID
Returns:
站点统计数据
"""
try:
site_mgr = SiteManager(self.db_path)
click_mgr = ClickManager(self.db_path)
interaction_mgr = InteractionManager(self.db_path)
site = site_mgr.get_site_by_id(site_id)
if not site:
return {}
click_count = click_mgr.get_click_count_by_site(site_id)
success_count = interaction_mgr.get_successful_interactions_count(site_id)
return {
'site_url': site['site_url'],
'site_name': site['site_name'],
'status': site['status'],
'click_count': click_count,
'reply_count': site['reply_count'],
'successful_interactions': success_count,
'reply_rate': f"{(site['reply_count'] / click_count * 100) if click_count > 0 else 0:.2f}%"
}
except Exception as e:
logger.error(f"获取站点统计失败: {str(e)}")
return {}