399 lines
12 KiB
Python
Executable File
399 lines
12 KiB
Python
Executable File
"""
|
|
会话存储模块 - 提供持久化存储功能
|
|
|
|
功能:
|
|
1. 用户会话历史保存到 SQLite 数据库
|
|
2. 活跃任务状态保存
|
|
3. 服务重启后能恢复会话
|
|
"""
|
|
|
|
import sqlite3
|
|
import json
|
|
import os
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
# 默认数据库路径
|
|
DEFAULT_DB_PATH = str(Path(__file__).resolve().parent / "data" / "session_storage.db")
|
|
|
|
|
|
class SessionStorage:
|
|
"""会话存储管理器"""
|
|
|
|
def __init__(self, db_path: str = DEFAULT_DB_PATH):
|
|
self.db_path = db_path
|
|
self._ensure_db_dir()
|
|
self._init_tables()
|
|
|
|
def _ensure_db_dir(self):
|
|
"""确保数据库目录存在"""
|
|
db_dir = os.path.dirname(self.db_path)
|
|
if db_dir and not os.path.exists(db_dir):
|
|
os.makedirs(db_dir, exist_ok=True)
|
|
|
|
def _get_connection(self) -> sqlite3.Connection:
|
|
"""获取数据库连接"""
|
|
conn = sqlite3.connect(self.db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
def _init_tables(self):
|
|
"""初始化数据库表"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
# 用户会话历史表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS user_sessions (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id TEXT NOT NULL,
|
|
task_id TEXT NOT NULL,
|
|
job_id TEXT NOT NULL,
|
|
prefix TEXT NOT NULL,
|
|
stage TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
timestamp TEXT NOT NULL,
|
|
final_members INTEGER,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 用户活跃任务表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS user_active_jobs (
|
|
user_id TEXT PRIMARY KEY,
|
|
job_id TEXT NOT NULL,
|
|
prefix TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
stages TEXT, -- JSON 格式存储
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 运行中任务表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS running_tasks (
|
|
task_id TEXT PRIMARY KEY,
|
|
job_id TEXT NOT NULL,
|
|
prefix TEXT NOT NULL,
|
|
stage TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
timestamp TEXT NOT NULL,
|
|
user_id TEXT NOT NULL,
|
|
final_members INTEGER,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 创建索引
|
|
cursor.execute('''
|
|
CREATE INDEX IF NOT EXISTS idx_user_sessions_user_id
|
|
ON user_sessions(user_id)
|
|
''')
|
|
cursor.execute('''
|
|
CREATE INDEX IF NOT EXISTS idx_user_sessions_timestamp
|
|
ON user_sessions(timestamp)
|
|
''')
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# ==================== 用户会话历史 ====================
|
|
|
|
def get_user_history(self, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
|
|
"""获取用户历史任务"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
SELECT task_id, job_id, prefix, stage, status, timestamp, final_members
|
|
FROM user_sessions
|
|
WHERE user_id = ?
|
|
ORDER BY timestamp DESC
|
|
LIMIT ?
|
|
''', (user_id, limit))
|
|
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
return [dict(row) for row in reversed(rows)] # 反转回时间正序
|
|
|
|
def add_user_task(self, user_id: str, task_info: Dict[str, Any]):
|
|
"""添加任务到用户历史"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
INSERT INTO user_sessions
|
|
(user_id, task_id, job_id, prefix, stage, status, timestamp, final_members)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
''', (
|
|
user_id,
|
|
task_info.get('task_id', ''),
|
|
task_info.get('job_id', ''),
|
|
task_info.get('prefix', ''),
|
|
task_info.get('stage', ''),
|
|
task_info.get('status', ''),
|
|
task_info.get('timestamp', datetime.now().strftime('%Y-%m-%d %H:%M:%S')),
|
|
task_info.get('final_members')
|
|
))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def update_task_status(self, user_id: str, task_id: str, status: str,
|
|
final_members: Optional[int] = None):
|
|
"""更新任务状态"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
if final_members is not None:
|
|
cursor.execute('''
|
|
UPDATE user_sessions
|
|
SET status = ?, final_members = ?
|
|
WHERE user_id = ? AND task_id = ?
|
|
''', (status, final_members, user_id, task_id))
|
|
else:
|
|
cursor.execute('''
|
|
UPDATE user_sessions
|
|
SET status = ?
|
|
WHERE user_id = ? AND task_id = ?
|
|
''', (status, user_id, task_id))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_latest_task(self, user_id: str) -> Optional[Dict[str, Any]]:
|
|
"""获取用户最近一个任务"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
SELECT task_id, job_id, prefix, stage, status, timestamp, final_members
|
|
FROM user_sessions
|
|
WHERE user_id = ?
|
|
ORDER BY timestamp DESC
|
|
LIMIT 1
|
|
''', (user_id,))
|
|
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
return dict(row) if row else None
|
|
|
|
# ==================== 用户活跃任务 ====================
|
|
|
|
def get_user_active_job(self, user_id: str) -> Optional[Dict[str, Any]]:
|
|
"""获取用户当前活跃 job"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
SELECT user_id, job_id, prefix, created_at, stages
|
|
FROM user_active_jobs
|
|
WHERE user_id = ?
|
|
''', (user_id,))
|
|
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
if row:
|
|
job = dict(row)
|
|
job['stages'] = json.loads(job['stages'] or '{}')
|
|
return job
|
|
return None
|
|
|
|
def set_user_active_job(self, user_id: str, job_info: Dict[str, Any]):
|
|
"""设置用户当前活跃 job"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
stages_json = json.dumps(job_info.get('stages', {}))
|
|
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO user_active_jobs
|
|
(user_id, job_id, prefix, created_at, stages, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
|
|
''', (
|
|
user_id,
|
|
job_info.get('job_id', ''),
|
|
job_info.get('prefix', ''),
|
|
job_info.get('created_at', datetime.now().strftime('%Y-%m-%d %H:%M:%S')),
|
|
stages_json
|
|
))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def update_job_stage(self, user_id: str, stage: str, status: str):
|
|
"""更新 job 的阶段状态"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
# 先获取当前 stages
|
|
cursor.execute('''
|
|
SELECT stages FROM user_active_jobs WHERE user_id = ?
|
|
''', (user_id,))
|
|
|
|
row = cursor.fetchone()
|
|
if row:
|
|
stages = json.loads(row['stages'] or '{}')
|
|
stages[stage] = status
|
|
|
|
cursor.execute('''
|
|
UPDATE user_active_jobs
|
|
SET stages = ?, updated_at = CURRENT_TIMESTAMP
|
|
WHERE user_id = ?
|
|
''', (json.dumps(stages), user_id))
|
|
|
|
conn.commit()
|
|
|
|
conn.close()
|
|
|
|
def clear_user_active_job(self, user_id: str):
|
|
"""清除用户活跃 job"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
DELETE FROM user_active_jobs WHERE user_id = ?
|
|
''', (user_id,))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# ==================== 运行中任务 ====================
|
|
|
|
def get_running_tasks(self) -> Dict[str, Dict[str, Any]]:
|
|
"""获取所有运行中任务"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
SELECT task_id, job_id, prefix, stage, status, timestamp, user_id, final_members
|
|
FROM running_tasks
|
|
''')
|
|
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
return {row['task_id']: dict(row) for row in rows}
|
|
|
|
def add_running_task(self, task_id: str, task_info: Dict[str, Any]):
|
|
"""添加运行中任务"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO running_tasks
|
|
(task_id, job_id, prefix, stage, status, timestamp, user_id, final_members)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
''', (
|
|
task_id,
|
|
task_info.get('job_id', ''),
|
|
task_info.get('prefix', ''),
|
|
task_info.get('stage', ''),
|
|
task_info.get('status', ''),
|
|
task_info.get('timestamp', datetime.now().strftime('%Y-%m-%d %H:%M:%S')),
|
|
task_info.get('user_id', ''),
|
|
task_info.get('final_members')
|
|
))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def update_running_task(self, task_id: str, status: str,
|
|
final_members: Optional[int] = None):
|
|
"""更新运行中任务状态"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
if final_members is not None:
|
|
cursor.execute('''
|
|
UPDATE running_tasks
|
|
SET status = ?, final_members = ?, updated_at = CURRENT_TIMESTAMP
|
|
WHERE task_id = ?
|
|
''', (status, final_members, task_id))
|
|
else:
|
|
cursor.execute('''
|
|
UPDATE running_tasks
|
|
SET status = ?, updated_at = CURRENT_TIMESTAMP
|
|
WHERE task_id = ?
|
|
''', (status, task_id))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def remove_running_task(self, task_id: str):
|
|
"""移除运行中任务"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
DELETE FROM running_tasks WHERE task_id = ?
|
|
''', (task_id,))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# ==================== 统计和清理 ====================
|
|
|
|
def get_user_stats(self, user_id: str) -> Dict[str, int]:
|
|
"""获取用户任务统计"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
SELECT status, COUNT(*) as count
|
|
FROM user_sessions
|
|
WHERE user_id = ?
|
|
GROUP BY status
|
|
''', (user_id,))
|
|
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
stats = {'running': 0, 'completed': 0, 'failed': 0}
|
|
for row in rows:
|
|
stats[row['status']] = row['count']
|
|
return stats
|
|
|
|
def cleanup_old_sessions(self, days: int = 30):
|
|
"""清理旧会话数据"""
|
|
conn = self._get_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute('''
|
|
DELETE FROM user_sessions
|
|
WHERE created_at < datetime('now', '-{} days')
|
|
'''.format(days))
|
|
|
|
deleted = cursor.rowcount
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
return deleted
|
|
|
|
def close(self):
|
|
"""关闭存储（SQLite 不需要显式关闭，但保留接口一致性）"""
|
|
pass
|
|
|
|
|
|
# 全局存储实例
|
|
_storage_instance: Optional[SessionStorage] = None
|
|
|
|
|
|
def get_storage() -> SessionStorage:
|
|
"""获取全局存储实例"""
|
|
global _storage_instance
|
|
if _storage_instance is None:
|
|
_storage_instance = SessionStorage()
|
|
return _storage_instance
|
|
|
|
|
|
def init_storage(db_path: str = DEFAULT_DB_PATH) -> SessionStorage:
|
|
"""初始化存储并返回实例"""
|
|
global _storage_instance
|
|
_storage_instance = SessionStorage(db_path)
|
|
return _storage_instance |