411 lines
15 KiB
Python
411 lines
15 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
天网任务数据模型
|
||
"""
|
||
import sqlite3
|
||
import json
|
||
import logging
|
||
from datetime import datetime, timedelta
|
||
from typing import Optional, Dict, List
|
||
from pathlib import Path
|
||
from enum import Enum
|
||
import os
|
||
from db.chat_log_db import _get_pooled_conn, _return_conn
|
||
|
||
logger = logging.getLogger(__name__)
|
||
_DB_TYPE = os.getenv("DB_TYPE", "sqlite").lower()
|
||
_MYSQL_HOST = os.getenv("MYSQL_HOST", "127.0.0.1")
|
||
_MYSQL_PORT = int(os.getenv("MYSQL_PORT", "3306"))
|
||
_MYSQL_USER = os.getenv("MYSQL_USER", "root")
|
||
_MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "")
|
||
_MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "ai_cs")
|
||
|
||
def _is_mysql() -> bool:
|
||
return _DB_TYPE in ("mysql", "mariadb")
|
||
|
||
def _sql(query: str) -> str:
|
||
return query.replace("?", "%s") if _is_mysql() else query
|
||
|
||
def _now_str() -> str:
|
||
if _is_mysql():
|
||
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
return datetime.now().isoformat()
|
||
|
||
class TaskStatus(Enum):
|
||
"""任务状态"""
|
||
PENDING = "pending" # 待触发
|
||
WAITING = "waiting" # 等待触发
|
||
RUNNING = "running" # 执行中
|
||
COMPLETED = "completed" # 已完成
|
||
FAILED = "failed" # 失败
|
||
CANCELLED = "cancelled" # 已取消
|
||
|
||
class TaskPriority(Enum):
|
||
"""任务优先级"""
|
||
NORMAL = "normal"
|
||
HIGH = "high"
|
||
URGENT = "urgent"
|
||
|
||
|
||
class _PooledMySQLConn:
|
||
"""包装 pymysql 连接，close 时归还到共享连接池。"""
|
||
|
||
def __init__(self, conn):
|
||
self._conn = conn
|
||
|
||
def __getattr__(self, name):
|
||
return getattr(self._conn, name)
|
||
|
||
def close(self):
|
||
_return_conn(self._conn)
|
||
|
||
class TaskManager:
|
||
"""任务管理器 - SQLite 存储"""
|
||
|
||
def __init__(self, db_path: str = None):
|
||
if db_path is None:
|
||
db_path = Path(__file__).parent / "tasks.db"
|
||
|
||
self.db_path = db_path
|
||
self._init_db()
|
||
logger.info(f"任务管理器初始化完成，数据库：{self.db_path}")
|
||
|
||
def _init_db(self):
|
||
"""初始化数据库"""
|
||
if _is_mysql():
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
cursor.execute('''
|
||
CREATE TABLE IF NOT EXISTS tasks (
|
||
task_id VARCHAR(128) PRIMARY KEY,
|
||
specified_customer_id VARCHAR(128),
|
||
specified_customer_name VARCHAR(255),
|
||
type VARCHAR(64) NOT NULL,
|
||
customer_name VARCHAR(255),
|
||
customer_id VARCHAR(128),
|
||
trigger_type VARCHAR(64),
|
||
trigger_keyword VARCHAR(255),
|
||
trigger_keywords TEXT,
|
||
action_type VARCHAR(64),
|
||
action_file_url TEXT,
|
||
action_message TEXT,
|
||
priority VARCHAR(16) DEFAULT 'normal',
|
||
timeout_hours INT DEFAULT 24,
|
||
status VARCHAR(32) DEFAULT 'pending',
|
||
retry_count INT DEFAULT 0,
|
||
max_retry INT DEFAULT 3,
|
||
created_at DATETIME,
|
||
created_by VARCHAR(255),
|
||
triggered_at DATETIME,
|
||
completed_at DATETIME,
|
||
error_message TEXT,
|
||
result TEXT
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
|
||
''')
|
||
cursor.execute("SHOW INDEX FROM tasks")
|
||
exists = {str(r.get("Key_name", "")) for r in cursor.fetchall()}
|
||
if "idx_status" not in exists:
|
||
cursor.execute('CREATE INDEX idx_status ON tasks(status)')
|
||
if "idx_customer" not in exists:
|
||
cursor.execute('CREATE INDEX idx_customer ON tasks(customer_id)')
|
||
if "idx_created" not in exists:
|
||
cursor.execute('CREATE INDEX idx_created ON tasks(created_at)')
|
||
conn.commit()
|
||
conn.close()
|
||
else:
|
||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||
conn = sqlite3.connect(self.db_path)
|
||
cursor = conn.cursor()
|
||
cursor.execute('''
|
||
CREATE TABLE IF NOT EXISTS tasks (
|
||
task_id TEXT PRIMARY KEY,
|
||
specified_customer_id TEXT,
|
||
specified_customer_name TEXT,
|
||
type TEXT NOT NULL,
|
||
customer_name TEXT,
|
||
customer_id TEXT,
|
||
trigger_type TEXT,
|
||
trigger_keyword TEXT,
|
||
trigger_keywords TEXT,
|
||
action_type TEXT,
|
||
action_file_url TEXT,
|
||
action_message TEXT,
|
||
priority TEXT DEFAULT 'normal',
|
||
timeout_hours INTEGER DEFAULT 24,
|
||
status TEXT DEFAULT 'pending',
|
||
retry_count INTEGER DEFAULT 0,
|
||
max_retry INTEGER DEFAULT 3,
|
||
created_at TEXT,
|
||
created_by TEXT,
|
||
triggered_at TEXT,
|
||
completed_at TEXT,
|
||
error_message TEXT,
|
||
result TEXT
|
||
)
|
||
''')
|
||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON tasks(status)')
|
||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_customer ON tasks(customer_id)')
|
||
cursor.execute('CREATE INDEX IF NOT EXISTS idx_created ON tasks(created_at)')
|
||
conn.commit()
|
||
conn.close()
|
||
logger.info("数据库表初始化完成")
|
||
|
||
def _get_conn(self):
|
||
"""获取数据库连接"""
|
||
if _is_mysql():
|
||
return _PooledMySQLConn(_get_pooled_conn())
|
||
conn = sqlite3.connect(self.db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
return conn
|
||
|
||
def add_task(self, task: dict) -> bool:
|
||
"""添加任务"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
# 处理 keywords 数组
|
||
trigger_keywords = task.get('trigger', {}).get('keywords', [])
|
||
if isinstance(trigger_keywords, list):
|
||
trigger_keywords = json.dumps(trigger_keywords)
|
||
|
||
insert_sql = '''
|
||
INSERT OR REPLACE INTO tasks (
|
||
task_id, specified_customer_id, specified_customer_name,
|
||
type, customer_name, customer_id,
|
||
trigger_type, trigger_keyword, trigger_keywords,
|
||
action_type, action_file_url, action_message,
|
||
priority, timeout_hours, status, retry_count, max_retry,
|
||
created_at, created_by
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
'''
|
||
if _is_mysql():
|
||
insert_sql = '''
|
||
REPLACE INTO tasks (
|
||
task_id, specified_customer_id, specified_customer_name,
|
||
type, customer_name, customer_id,
|
||
trigger_type, trigger_keyword, trigger_keywords,
|
||
action_type, action_file_url, action_message,
|
||
priority, timeout_hours, status, retry_count, max_retry,
|
||
created_at, created_by
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
'''
|
||
cursor.execute(_sql(insert_sql), (
|
||
task.get('task_id'),
|
||
task.get('customer', {}).get('id'),
|
||
task.get('customer', {}).get('name'),
|
||
task.get('type'),
|
||
task.get('customer', {}).get('name'),
|
||
task.get('customer', {}).get('id'),
|
||
task.get('trigger', {}).get('type'),
|
||
task.get('trigger', {}).get('keyword'),
|
||
trigger_keywords,
|
||
task.get('action', {}).get('type'),
|
||
task.get('action', {}).get('file_url'),
|
||
task.get('action', {}).get('message'),
|
||
task.get('priority', 'normal'),
|
||
task.get('timeout_hours', 24),
|
||
task.get('status', 'pending'),
|
||
task.get('retry_count', 0),
|
||
task.get('max_retry', 3),
|
||
task.get('created_at', _now_str()),
|
||
task.get('created_by')
|
||
))
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
logger.info(f"任务添加成功：{task.get('task_id')}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"添加任务失败：{e}")
|
||
return False
|
||
|
||
def get_task(self, task_id: str) -> Optional[dict]:
|
||
"""查询任务"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute(_sql('SELECT * FROM tasks WHERE task_id = ?'), (task_id,))
|
||
row = cursor.fetchone()
|
||
conn.close()
|
||
|
||
if row:
|
||
return dict(row)
|
||
return None
|
||
|
||
except Exception as e:
|
||
logger.error(f"查询任务失败：{e}")
|
||
return None
|
||
|
||
def update_task_status(self, task_id: str, status: TaskStatus, error_message: str = None, result: dict = None):
|
||
"""更新任务状态"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
placeholder = "%s" if _is_mysql() else "?"
|
||
updates = [f'status = {placeholder}']
|
||
params = [status.value]
|
||
|
||
if status == TaskStatus.RUNNING:
|
||
updates.append(f'triggered_at = {placeholder}')
|
||
params.append(_now_str())
|
||
|
||
if status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]:
|
||
updates.append(f'completed_at = {placeholder}')
|
||
params.append(_now_str())
|
||
|
||
if error_message:
|
||
updates.append(f'error_message = {placeholder}')
|
||
params.append(error_message)
|
||
|
||
if result:
|
||
updates.append(f'result = {placeholder}')
|
||
params.append(json.dumps(result))
|
||
|
||
params.append(task_id)
|
||
|
||
cursor.execute(_sql(f'''
|
||
UPDATE tasks
|
||
SET {', '.join(updates)}
|
||
WHERE task_id = ?
|
||
'''), params)
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
logger.info(f"任务状态更新：{task_id} -> {status.value}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"更新任务状态失败：{e}")
|
||
|
||
def increment_retry(self, task_id: str) -> int:
|
||
"""增加重试次数，返回当前重试次数"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute(_sql('''
|
||
UPDATE tasks
|
||
SET retry_count = retry_count + 1
|
||
WHERE task_id = ?
|
||
'''), (task_id,))
|
||
|
||
cursor.execute(_sql('SELECT retry_count, max_retry FROM tasks WHERE task_id = ?'), (task_id,))
|
||
row = cursor.fetchone()
|
||
conn.close()
|
||
|
||
if row:
|
||
return row['retry_count']
|
||
return 0
|
||
|
||
except Exception as e:
|
||
logger.error(f"增加重试次数失败：{e}")
|
||
return 999 # 超过最大重试次数
|
||
|
||
def get_pending_tasks(self, customer_id: str = None) -> List[dict]:
|
||
"""获取待触发的任务"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
if customer_id:
|
||
cursor.execute(_sql('''
|
||
SELECT * FROM tasks
|
||
WHERE status = 'pending'
|
||
AND customer_id = ?
|
||
ORDER BY
|
||
CASE priority
|
||
WHEN 'urgent' THEN 1
|
||
WHEN 'high' THEN 2
|
||
ELSE 3
|
||
END,
|
||
created_at
|
||
'''), (customer_id,))
|
||
else:
|
||
cursor.execute('''
|
||
SELECT * FROM tasks
|
||
WHERE status = 'pending'
|
||
ORDER BY
|
||
CASE priority
|
||
WHEN 'urgent' THEN 1
|
||
WHEN 'high' THEN 2
|
||
ELSE 3
|
||
END,
|
||
created_at
|
||
''')
|
||
|
||
rows = cursor.fetchall()
|
||
conn.close()
|
||
|
||
return [dict(row) for row in rows]
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取待触发任务失败：{e}")
|
||
return []
|
||
|
||
def get_timeout_tasks(self) -> List[dict]:
|
||
"""获取超时任务"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
if _is_mysql():
|
||
cursor.execute('''
|
||
SELECT * FROM tasks
|
||
WHERE status = 'pending'
|
||
AND created_at < DATE_SUB(NOW(), INTERVAL timeout_hours HOUR)
|
||
''')
|
||
else:
|
||
cursor.execute('''
|
||
SELECT * FROM tasks
|
||
WHERE status = 'pending'
|
||
AND datetime(created_at, '+' || timeout_hours || ' hours') < datetime('now')
|
||
''')
|
||
|
||
rows = cursor.fetchall()
|
||
conn.close()
|
||
|
||
return [dict(row) for row in rows]
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取超时任务失败：{e}")
|
||
return []
|
||
|
||
def cancel_task(self, task_id: str, reason: str = None):
|
||
"""取消任务"""
|
||
self.update_task_status(task_id, TaskStatus.CANCELLED, error_message=reason)
|
||
logger.info(f"任务已取消：{task_id}")
|
||
|
||
def get_statistics(self) -> dict:
|
||
"""获取任务统计"""
|
||
try:
|
||
conn = self._get_conn()
|
||
cursor = conn.cursor()
|
||
|
||
stats = {}
|
||
for status in TaskStatus:
|
||
cursor.execute(_sql('SELECT COUNT(*) FROM tasks WHERE status = ?'), (status.value,))
|
||
stats[status.value] = cursor.fetchone()[0]
|
||
|
||
conn.close()
|
||
return stats
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取统计失败：{e}")
|
||
return {}
|
||
|
||
|
||
# 单例
|
||
_task_manager: Optional[TaskManager] = None
|
||
|
||
def get_task_manager() -> TaskManager:
|
||
"""获取任务管理器单例"""
|
||
global _task_manager
|
||
if _task_manager is None:
|
||
_task_manager = TaskManager()
|
||
return _task_manager
|