362 lines
11 KiB
Python
362 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Queue manager for handling file processing queues.
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
from typing import Dict, List, Optional, Any
|
||
from huey import Huey
|
||
from huey.api import Task
|
||
from datetime import datetime, timedelta
|
||
|
||
from .config import huey
|
||
from .tasks import process_file_async, process_multiple_files_async, process_zip_file_async, cleanup_processed_files
|
||
|
||
|
||
class QueueManager:
|
||
"""队列管理器，用于管理文件处理任务"""
|
||
|
||
def __init__(self):
|
||
self.huey = huey
|
||
print(f"队列管理器已初始化，使用数据库: {os.path.join(os.path.dirname(__file__), '..', 'queue_data', 'huey.db')}")
|
||
|
||
def enqueue_file(
|
||
self,
|
||
project_id: str,
|
||
file_path: str,
|
||
original_filename: str = None,
|
||
delay: int = 0
|
||
) -> str:
|
||
"""
|
||
将文件加入处理队列
|
||
|
||
Args:
|
||
project_id: 项目ID
|
||
file_path: 文件路径
|
||
original_filename: 原始文件名
|
||
delay: 延迟执行时间（秒）
|
||
|
||
Returns:
|
||
任务ID
|
||
"""
|
||
if delay > 0:
|
||
task = process_file_async.schedule(
|
||
args=(project_id, file_path, original_filename),
|
||
delay=timedelta(seconds=delay)
|
||
)
|
||
else:
|
||
task = process_file_async(project_id, file_path, original_filename)
|
||
|
||
print(f"文件已加入队列: {file_path}, 任务ID: {task.id}")
|
||
return task.id
|
||
|
||
def enqueue_multiple_files(
|
||
self,
|
||
project_id: str,
|
||
file_paths: List[str],
|
||
original_filenames: List[str] = None,
|
||
delay: int = 0
|
||
) -> List[str]:
|
||
"""
|
||
将多个文件加入处理队列
|
||
|
||
Args:
|
||
project_id: 项目ID
|
||
file_paths: 文件路径列表
|
||
original_filenames: 原始文件名列表
|
||
delay: 延迟执行时间（秒）
|
||
|
||
Returns:
|
||
任务ID列表
|
||
"""
|
||
if delay > 0:
|
||
task = process_multiple_files_async.schedule(
|
||
args=(project_id, file_paths, original_filenames),
|
||
delay=timedelta(seconds=delay)
|
||
)
|
||
else:
|
||
task = process_multiple_files_async(project_id, file_paths, original_filenames)
|
||
|
||
print(f"批量文件已加入队列: {len(file_paths)} 个文件, 任务ID: {task.id}")
|
||
return [task.id]
|
||
|
||
def enqueue_zip_file(
|
||
self,
|
||
project_id: str,
|
||
zip_path: str,
|
||
extract_to: str = None,
|
||
delay: int = 0
|
||
) -> str:
|
||
"""
|
||
将zip文件加入处理队列
|
||
|
||
Args:
|
||
project_id: 项目ID
|
||
zip_path: zip文件路径
|
||
extract_to: 解压目标目录
|
||
delay: 延迟执行时间（秒）
|
||
|
||
Returns:
|
||
任务ID
|
||
"""
|
||
if delay > 0:
|
||
task = process_zip_file_async.schedule(
|
||
args=(project_id, zip_path, extract_to),
|
||
delay=timedelta(seconds=delay)
|
||
)
|
||
else:
|
||
task = process_zip_file_async(project_id, zip_path, extract_to)
|
||
|
||
print(f"zip文件已加入队列: {zip_path}, 任务ID: {task.id}")
|
||
return task.id
|
||
|
||
def get_task_status(self, task_id: str) -> Dict[str, Any]:
|
||
"""
|
||
获取任务状态
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
任务状态信息
|
||
"""
|
||
try:
|
||
# 尝试从结果存储中获取任务结果
|
||
try:
|
||
# 使用huey的内置方法检查结果
|
||
if hasattr(self.huey, 'result') and self.huey.result:
|
||
result = self.huey.result(task_id)
|
||
if result is not None:
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "complete",
|
||
"result": result
|
||
}
|
||
except Exception:
|
||
pass
|
||
|
||
# 检查任务是否在待处理队列中
|
||
try:
|
||
pending_tasks = list(self.huey.pending())
|
||
for task in pending_tasks:
|
||
if hasattr(task, 'id') and task.id == task_id:
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "pending"
|
||
}
|
||
except Exception:
|
||
pass
|
||
|
||
# 检查任务是否在定时队列中
|
||
try:
|
||
scheduled_tasks = list(self.huey.scheduled())
|
||
for task in scheduled_tasks:
|
||
if hasattr(task, 'id') and task.id == task_id:
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "scheduled"
|
||
}
|
||
except Exception:
|
||
pass
|
||
|
||
# 如果都找不到，可能任务不存在或已完成但结果已清理
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "unknown",
|
||
"message": "任务状态未知，可能已完成或不存在"
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"task_id": task_id,
|
||
"status": "error",
|
||
"message": f"获取任务状态失败: {str(e)}"
|
||
}
|
||
|
||
def get_queue_stats(self) -> Dict[str, Any]:
|
||
"""
|
||
获取队列统计信息
|
||
|
||
Returns:
|
||
队列统计信息
|
||
"""
|
||
try:
|
||
# 使用简化的统计方法
|
||
stats = {
|
||
"total_tasks": 0,
|
||
"pending_tasks": 0,
|
||
"running_tasks": 0,
|
||
"completed_tasks": 0,
|
||
"error_tasks": 0,
|
||
"scheduled_tasks": 0,
|
||
"recent_tasks": [],
|
||
"queue_database": os.path.join(os.path.dirname(__file__), '..', 'queue_data', 'huey.db')
|
||
}
|
||
|
||
# 尝试获取待处理任务数量
|
||
try:
|
||
pending_tasks = list(self.huey.pending())
|
||
stats["pending_tasks"] = len(pending_tasks)
|
||
stats["total_tasks"] += len(pending_tasks)
|
||
except Exception as e:
|
||
print(f"获取pending任务失败: {e}")
|
||
|
||
# 尝试获取定时任务数量
|
||
try:
|
||
scheduled_tasks = list(self.huey.scheduled())
|
||
stats["scheduled_tasks"] = len(scheduled_tasks)
|
||
stats["total_tasks"] += len(scheduled_tasks)
|
||
except Exception as e:
|
||
print(f"获取scheduled任务失败: {e}")
|
||
|
||
return stats
|
||
|
||
except Exception as e:
|
||
return {
|
||
"error": f"获取队列统计信息失败: {str(e)}"
|
||
}
|
||
|
||
def cancel_task(self, task_id: str) -> bool:
|
||
"""
|
||
取消任务
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
是否成功取消
|
||
"""
|
||
try:
|
||
task = self.huey.get_task(task_id)
|
||
if task and task.get_status() in ["pending", "scheduled"]:
|
||
# Huey不直接支持取消任务，但可以从队列中移除
|
||
# 这里返回False表示不能直接取消
|
||
return False
|
||
else:
|
||
return False
|
||
except Exception as e:
|
||
print(f"取消任务失败: {str(e)}")
|
||
return False
|
||
|
||
def cleanup_old_tasks(self, older_than_days: int = 7) -> Dict[str, Any]:
|
||
"""
|
||
清理旧的任务记录
|
||
|
||
Args:
|
||
older_than_days: 清理多少天前的任务记录
|
||
|
||
Returns:
|
||
清理结果
|
||
"""
|
||
try:
|
||
# 简化的清理方法 - 清空整个队列
|
||
self.huey.flush()
|
||
|
||
return {
|
||
"status": "success",
|
||
"message": f"已清空队列（简化清理）",
|
||
"older_than_days": older_than_days,
|
||
"cleaned_count": "unknown (queue flushed)"
|
||
}
|
||
|
||
except Exception as e:
|
||
return {
|
||
"status": "error",
|
||
"message": f"清理任务记录失败: {str(e)}"
|
||
}
|
||
|
||
def enqueue_cleanup_task(
|
||
self,
|
||
project_id: str,
|
||
older_than_days: int = 30,
|
||
delay: int = 0
|
||
) -> str:
|
||
"""
|
||
将清理任务加入队列
|
||
|
||
Args:
|
||
project_id: 项目ID
|
||
older_than_days: 清理多少天前的文件
|
||
delay: 延迟执行时间（秒）
|
||
|
||
Returns:
|
||
任务ID
|
||
"""
|
||
if delay > 0:
|
||
task = cleanup_processed_files.schedule(
|
||
args=(project_id, older_than_days),
|
||
delay=timedelta(seconds=delay)
|
||
)
|
||
else:
|
||
task = cleanup_processed_files(project_id, older_than_days)
|
||
|
||
print(f"清理任务已加入队列: 项目 {project_id}, 任务ID: {task.id}")
|
||
return task.id
|
||
|
||
def list_pending_tasks(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||
"""
|
||
列出待处理的任务
|
||
|
||
Args:
|
||
limit: 返回的最大任务数
|
||
|
||
Returns:
|
||
待处理任务列表
|
||
"""
|
||
try:
|
||
pending_tasks = []
|
||
|
||
# 获取pending任务
|
||
try:
|
||
tasks = list(self.huey.pending())
|
||
for i, task in enumerate(tasks[:limit]):
|
||
if hasattr(task, 'id'):
|
||
pending_tasks.append({
|
||
"task_id": task.id,
|
||
"status": "pending",
|
||
})
|
||
except Exception as e:
|
||
print(f"获取pending任务失败: {e}")
|
||
|
||
# 获取scheduled任务
|
||
try:
|
||
tasks = list(self.huey.scheduled())
|
||
for i, task in enumerate(tasks[:limit - len(pending_tasks)]):
|
||
if hasattr(task, 'id'):
|
||
pending_tasks.append({
|
||
"task_id": task.id,
|
||
"status": "scheduled",
|
||
})
|
||
except Exception as e:
|
||
print(f"获取scheduled任务失败: {e}")
|
||
|
||
return pending_tasks
|
||
|
||
except Exception as e:
|
||
print(f"获取待处理任务失败: {str(e)}")
|
||
return []
|
||
|
||
def get_task_result(self, task_id: str) -> Optional[Any]:
|
||
"""
|
||
获取任务结果
|
||
|
||
Args:
|
||
task_id: 任务ID
|
||
|
||
Returns:
|
||
任务结果，如果任务未完成则返回None
|
||
"""
|
||
try:
|
||
task = self.huey.get_task(task_id)
|
||
if task and task.get_status() in ["complete", "finished"]:
|
||
return task.get()
|
||
return None
|
||
except Exception as e:
|
||
print(f"获取任务结果失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
# 全局队列管理器实例
|
||
queue_manager = QueueManager() |