548 lines
19 KiB
Python
548 lines
19 KiB
Python
"""
|
||
数据库管理器
|
||
使用 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 {}
|