""" MinerU Tianshu - SQLite Task Database Manager 天枢任务数据库管理器 负责任务的持久化存储、状态管理和原子性操作 """ import sqlite3 import json import uuid from contextlib import contextmanager from typing import Optional, List, Dict from pathlib import Path class TaskDB: """任务数据库管理类""" def __init__(self, db_path='mineru_tianshu.db'): self.db_path = db_path self._init_db() def _get_conn(self): """获取数据库连接（每次创建新连接，避免 pickle 问题）""" conn = sqlite3.connect( self.db_path, check_same_thread=False, timeout=30.0 ) conn.row_factory = sqlite3.Row return conn @contextmanager def get_cursor(self): """上下文管理器，自动提交和错误处理""" conn = self._get_conn() cursor = conn.cursor() try: yield cursor conn.commit() except Exception as e: conn.rollback() raise e finally: conn.close() # 关闭连接 def _init_db(self): """初始化数据库表""" with self.get_cursor() as cursor: cursor.execute(''' CREATE TABLE IF NOT EXISTS tasks ( task_id TEXT PRIMARY KEY, file_name TEXT NOT NULL, file_path TEXT, status TEXT DEFAULT 'pending', priority INTEGER DEFAULT 0, backend TEXT DEFAULT 'pipeline', options TEXT, result_path TEXT, error_message TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, started_at TIMESTAMP, completed_at TIMESTAMP, worker_id TEXT, retry_count INTEGER DEFAULT 0 ) ''') # 创建索引加速查询 cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON tasks(status)') cursor.execute('CREATE INDEX IF NOT EXISTS idx_priority ON tasks(priority DESC)') cursor.execute('CREATE INDEX IF NOT EXISTS idx_created_at ON tasks(created_at)') cursor.execute('CREATE INDEX IF NOT EXISTS idx_worker_id ON tasks(worker_id)') def create_task(self, file_name: str, file_path: str, backend: str = 'pipeline', options: dict = None, priority: int = 0) -> str: """ 创建新任务 Args: file_name: 文件名 file_path: 文件路径 backend: 处理后端 (pipeline/vlm-transformers/vlm-vllm-engine) options: 处理选项 (dict) priority: 优先级，数字越大越优先 Returns: task_id: 任务ID """ task_id = str(uuid.uuid4()) with self.get_cursor() as cursor: cursor.execute(''' INSERT INTO tasks (task_id, file_name, file_path, backend, options, priority) VALUES (?, ?, ?, ?, ?, ?) ''', (task_id, file_name, file_path, backend, json.dumps(options or {}), priority)) return task_id def get_next_task(self, worker_id: str) -> Optional[Dict]: """ 获取下一个待处理任务（原子操作，防止并发冲突） Args: worker_id: Worker ID Returns: task: 任务字典，如果没有任务返回 None """ with self.get_cursor() as cursor: # 使用事务确保原子性 cursor.execute('BEGIN IMMEDIATE') # 按优先级和创建时间获取任务 cursor.execute(''' SELECT * FROM tasks WHERE status = 'pending' ORDER BY priority DESC, created_at ASC LIMIT 1 ''') task = cursor.fetchone() if task: # 立即标记为 processing cursor.execute(''' UPDATE tasks SET status = 'processing', started_at = CURRENT_TIMESTAMP, worker_id = ? WHERE task_id = ? ''', (worker_id, task['task_id'])) return dict(task) return None def update_task_status(self, task_id: str, status: str, result_path: str = None, error_message: str = None): """ 更新任务状态 Args: task_id: 任务ID status: 新状态 (pending/processing/completed/failed/cancelled) result_path: 结果路径（可选） error_message: 错误信息（可选） """ with self.get_cursor() as cursor: updates = ['status = ?'] params = [status] if status == 'completed': updates.append('completed_at = CURRENT_TIMESTAMP') if result_path: updates.append('result_path = ?') params.append(result_path) if status == 'failed' and error_message: updates.append('error_message = ?') params.append(error_message) updates.append('completed_at = CURRENT_TIMESTAMP') params.append(task_id) cursor.execute(f''' UPDATE tasks SET {', '.join(updates)} WHERE task_id = ? ''', params) def get_task(self, task_id: str) -> Optional[Dict]: """ 查询任务详情 Args: task_id: 任务ID Returns: task: 任务字典，如果不存在返回 None """ with self.get_cursor() as cursor: cursor.execute('SELECT * FROM tasks WHERE task_id = ?', (task_id,)) task = cursor.fetchone() return dict(task) if task else None def get_queue_stats(self) -> Dict[str, int]: """ 获取队列统计信息 Returns: stats: 各状态的任务数量 """ with self.get_cursor() as cursor: cursor.execute(''' SELECT status, COUNT(*) as count FROM tasks GROUP BY status ''') stats = {row['status']: row['count'] for row in cursor.fetchall()} return stats def get_tasks_by_status(self, status: str, limit: int = 100) -> List[Dict]: """ 根据状态获取任务列表 Args: status: 任务状态 limit: 返回数量限制 Returns: tasks: 任务列表 """ with self.get_cursor() as cursor: cursor.execute(''' SELECT * FROM tasks WHERE status = ? ORDER BY created_at DESC LIMIT ? ''', (status, limit)) return [dict(row) for row in cursor.fetchall()] def cleanup_old_tasks(self, days: int = 7): """ 清理旧任务记录 Args: days: 保留最近N天的任务 """ with self.get_cursor() as cursor: cursor.execute(''' DELETE FROM tasks WHERE completed_at < datetime('now', '-' || ? || ' days') AND status IN ('completed', 'failed') ''', (days,)) deleted_count = cursor.rowcount return deleted_count def reset_stale_tasks(self, timeout_minutes: int = 60): """ 重置超时的 processing 任务为 pending Args: timeout_minutes: 超时时间（分钟） """ with self.get_cursor() as cursor: cursor.execute(''' UPDATE tasks SET status = 'pending', worker_id = NULL, retry_count = retry_count + 1 WHERE status = 'processing' AND started_at < datetime('now', '-' || ? || ' minutes') ''', (timeout_minutes,)) reset_count = cursor.rowcount return reset_count if __name__ == '__main__': # 测试代码 db = TaskDB('test_tianshu.db') # 创建测试任务 task_id = db.create_task( file_name='test.pdf', file_path='/tmp/test.pdf', backend='pipeline', options={'lang': 'ch', 'formula_enable': True}, priority=1 ) print(f"Created task: {task_id}") # 查询任务 task = db.get_task(task_id) print(f"Task details: {task}") # 获取统计 stats = db.get_queue_stats() print(f"Queue stats: {stats}") # 清理测试数据库 Path('test_tianshu.db').unlink(missing_ok=True) print("Test completed!")