437 lines
16 KiB
Python
437 lines
16 KiB
Python
|
"""
|
|||
|
|
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 问题）
|
|||
|
|
|
|||
|
|
并发安全说明：
|
|||
|
|
- 使用 check_same_thread=False 是安全的，因为：
|
|||
|
|
1. 每次调用都创建新连接，不跨线程共享
|
|||
|
|
2. 连接使用完立即关闭（在 get_cursor 上下文管理器中）
|
|||
|
|
3. 不使用连接池，避免线程间共享同一连接
|
|||
|
|
- timeout=30.0 防止死锁，如果锁等待超过30秒会抛出异常
|
|||
|
|
"""
|
|||
|
|
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, max_retries: int = 3) -> Optional[Dict]:
|
|||
|
|
"""
|
|||
|
|
获取下一个待处理任务（原子操作，防止并发冲突）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
worker_id: Worker ID
|
|||
|
|
max_retries: 当任务被其他 worker 抢走时的最大重试次数（默认3次）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
task: 任务字典，如果没有任务返回 None
|
|||
|
|
|
|||
|
|
并发安全说明：
|
|||
|
|
1. 使用 BEGIN IMMEDIATE 立即获取写锁
|
|||
|
|
2. UPDATE 时检查 status = 'pending' 防止重复拉取
|
|||
|
|
3. 检查 rowcount 确保更新成功
|
|||
|
|
4. 如果任务被抢走，立即重试而不是返回 None（避免不必要的等待）
|
|||
|
|
"""
|
|||
|
|
for attempt in range(max_retries):
|
|||
|
|
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，并确保状态仍是 pending
|
|||
|
|
cursor.execute('''
|
|||
|
|
UPDATE tasks
|
|||
|
|
SET status = 'processing',
|
|||
|
|
started_at = CURRENT_TIMESTAMP,
|
|||
|
|
worker_id = ?
|
|||
|
|
WHERE task_id = ? AND status = 'pending'
|
|||
|
|
''', (worker_id, task['task_id']))
|
|||
|
|
|
|||
|
|
# 检查是否更新成功（防止被其他 worker 抢走）
|
|||
|
|
if cursor.rowcount == 0:
|
|||
|
|
# 任务被其他进程抢走了，立即重试
|
|||
|
|
# 因为队列中可能还有其他待处理任务
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
return dict(task)
|
|||
|
|
else:
|
|||
|
|
# 队列中没有待处理任务，返回 None
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 重试次数用尽，仍未获取到任务（高并发场景）
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def _build_update_clauses(self, status: str, result_path: str = None,
|
|||
|
|
error_message: str = None, worker_id: str = None,
|
|||
|
|
task_id: str = None):
|
|||
|
|
"""
|
|||
|
|
构建 UPDATE 和 WHERE 子句的辅助方法
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
status: 新状态
|
|||
|
|
result_path: 结果路径（可选）
|
|||
|
|
error_message: 错误信息（可选）
|
|||
|
|
worker_id: Worker ID（可选）
|
|||
|
|
task_id: 任务ID（可选）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
tuple: (update_clauses, update_params, where_clauses, where_params)
|
|||
|
|
"""
|
|||
|
|
update_clauses = ['status = ?']
|
|||
|
|
update_params = [status]
|
|||
|
|
where_clauses = []
|
|||
|
|
where_params = []
|
|||
|
|
|
|||
|
|
# 添加 task_id 条件（如果提供）
|
|||
|
|
if task_id:
|
|||
|
|
where_clauses.append('task_id = ?')
|
|||
|
|
where_params.append(task_id)
|
|||
|
|
|
|||
|
|
# 处理 completed 状态
|
|||
|
|
if status == 'completed':
|
|||
|
|
update_clauses.append('completed_at = CURRENT_TIMESTAMP')
|
|||
|
|
if result_path:
|
|||
|
|
update_clauses.append('result_path = ?')
|
|||
|
|
update_params.append(result_path)
|
|||
|
|
# 只更新正在处理的任务
|
|||
|
|
where_clauses.append("status = 'processing'")
|
|||
|
|
if worker_id:
|
|||
|
|
where_clauses.append('worker_id = ?')
|
|||
|
|
where_params.append(worker_id)
|
|||
|
|
|
|||
|
|
# 处理 failed 状态
|
|||
|
|
elif status == 'failed':
|
|||
|
|
update_clauses.append('completed_at = CURRENT_TIMESTAMP')
|
|||
|
|
if error_message:
|
|||
|
|
update_clauses.append('error_message = ?')
|
|||
|
|
update_params.append(error_message)
|
|||
|
|
# 只更新正在处理的任务
|
|||
|
|
where_clauses.append("status = 'processing'")
|
|||
|
|
if worker_id:
|
|||
|
|
where_clauses.append('worker_id = ?')
|
|||
|
|
where_params.append(worker_id)
|
|||
|
|
|
|||
|
|
return update_clauses, update_params, where_clauses, where_params
|
|||
|
|
|
|||
|
|
def update_task_status(self, task_id: str, status: str,
|
|||
|
|
result_path: str = None, error_message: str = None,
|
|||
|
|
worker_id: str = None):
|
|||
|
|
"""
|
|||
|
|
更新任务状态
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
task_id: 任务ID
|
|||
|
|
status: 新状态 (pending/processing/completed/failed/cancelled)
|
|||
|
|
result_path: 结果路径（可选）
|
|||
|
|
error_message: 错误信息（可选）
|
|||
|
|
worker_id: Worker ID（可选，用于并发检查）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 更新是否成功
|
|||
|
|
|
|||
|
|
并发安全说明：
|
|||
|
|
1. 更新为 completed/failed 时会检查状态是 processing
|
|||
|
|
2. 如果提供 worker_id，会检查任务是否属于该 worker
|
|||
|
|
3. 返回 False 表示任务被其他进程修改了
|
|||
|
|
"""
|
|||
|
|
with self.get_cursor() as cursor:
|
|||
|
|
# 使用辅助方法构建 UPDATE 和 WHERE 子句
|
|||
|
|
update_clauses, update_params, where_clauses, where_params = \
|
|||
|
|
self._build_update_clauses(status, result_path, error_message, worker_id, task_id)
|
|||
|
|
|
|||
|
|
# 合并参数：先 UPDATE 部分，再 WHERE 部分
|
|||
|
|
all_params = update_params + where_params
|
|||
|
|
|
|||
|
|
sql = f'''
|
|||
|
|
UPDATE tasks
|
|||
|
|
SET {', '.join(update_clauses)}
|
|||
|
|
WHERE {' AND '.join(where_clauses)}
|
|||
|
|
'''
|
|||
|
|
|
|||
|
|
cursor.execute(sql, all_params)
|
|||
|
|
|
|||
|
|
# 检查更新是否成功
|
|||
|
|
success = cursor.rowcount > 0
|
|||
|
|
|
|||
|
|
# 调试日志（仅在失败时）
|
|||
|
|
if not success and status in ['completed', 'failed']:
|
|||
|
|
from loguru import logger
|
|||
|
|
logger.debug(
|
|||
|
|
f"Status update failed: task_id={task_id}, status={status}, "
|
|||
|
|
f"worker_id={worker_id}, SQL: {sql}, params: {all_params}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return success
|
|||
|
|
|
|||
|
|
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_task_files(self, days: int = 7):
|
|||
|
|
"""
|
|||
|
|
清理旧任务的结果文件（保留数据库记录）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
days: 清理多少天前的任务文件
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
int: 删除的文件目录数
|
|||
|
|
|
|||
|
|
注意：
|
|||
|
|
- 只删除结果文件，保留数据库记录
|
|||
|
|
- 数据库中的 result_path 字段会被清空
|
|||
|
|
- 用户仍可查询任务状态和历史记录
|
|||
|
|
"""
|
|||
|
|
from pathlib import Path
|
|||
|
|
import shutil
|
|||
|
|
|
|||
|
|
with self.get_cursor() as cursor:
|
|||
|
|
# 查询要清理文件的任务
|
|||
|
|
cursor.execute('''
|
|||
|
|
SELECT task_id, result_path FROM tasks
|
|||
|
|
WHERE completed_at < datetime('now', '-' || ? || ' days')
|
|||
|
|
AND status IN ('completed', 'failed')
|
|||
|
|
AND result_path IS NOT NULL
|
|||
|
|
''', (days,))
|
|||
|
|
|
|||
|
|
old_tasks = cursor.fetchall()
|
|||
|
|
file_count = 0
|
|||
|
|
|
|||
|
|
# 删除结果文件
|
|||
|
|
for task in old_tasks:
|
|||
|
|
if task['result_path']:
|
|||
|
|
result_path = Path(task['result_path'])
|
|||
|
|
if result_path.exists() and result_path.is_dir():
|
|||
|
|
try:
|
|||
|
|
shutil.rmtree(result_path)
|
|||
|
|
file_count += 1
|
|||
|
|
|
|||
|
|
# 清空数据库中的 result_path，表示文件已被清理
|
|||
|
|
cursor.execute('''
|
|||
|
|
UPDATE tasks
|
|||
|
|
SET result_path = NULL
|
|||
|
|
WHERE task_id = ?
|
|||
|
|
''', (task['task_id'],))
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
from loguru import logger
|
|||
|
|
logger.warning(f"Failed to delete result files for task {task['task_id']}: {e}")
|
|||
|
|
|
|||
|
|
return file_count
|
|||
|
|
|
|||
|
|
def cleanup_old_task_records(self, days: int = 30):
|
|||
|
|
"""
|
|||
|
|
清理极旧的任务记录（可选功能）
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
days: 删除多少天前的任务记录
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
int: 删除的记录数
|
|||
|
|
|
|||
|
|
注意：
|
|||
|
|
- 这个方法会永久删除数据库记录
|
|||
|
|
- 建议设置较长的保留期（如30-90天）
|
|||
|
|
- 一般情况下不需要调用此方法
|
|||
|
|
"""
|
|||
|
|
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!")
|
|||
|
|
|