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!")
|
||
|