358 lines
14 KiB
Python
358 lines
14 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
任务队列管理模块
|
|||
|
|
支持离线处理、进度跟踪、结果导出
|
|||
|
|
使用 SQLite 数据库存储（替代原 JSON 文件）
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
from datetime import datetime
|
|||
|
|
from enum import Enum
|
|||
|
|
import logging
|
|||
|
|
from database import get_database, migrate_from_json
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TaskStatus(Enum):
|
|||
|
|
"""任务状态"""
|
|||
|
|
PENDING = "pending" # 就绪（准备好了，等待工作线程）
|
|||
|
|
PROCESSING = "processing" # 进行中
|
|||
|
|
COMPLETED = "completed" # 完成
|
|||
|
|
FAILED = "failed" # 失败
|
|||
|
|
PAUSED = "paused" # 暂停（将在指定时间后自动恢复）
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TaskQueue:
|
|||
|
|
"""任务队列管理器（使用 SQLite 数据库）"""
|
|||
|
|
|
|||
|
|
def __init__(self, queue_file="data/task_queue.json", results_dir="data/results"):
|
|||
|
|
self.results_dir = results_dir
|
|||
|
|
self.lock = threading.Lock()
|
|||
|
|
self.db = get_database()
|
|||
|
|
self._ensure_dirs()
|
|||
|
|
|
|||
|
|
# 从旧 JSON 文件迁移数据（只执行一次）
|
|||
|
|
if os.path.exists(queue_file):
|
|||
|
|
migrate_from_json(queue_file)
|
|||
|
|
|
|||
|
|
def _ensure_dirs(self):
|
|||
|
|
"""确保必要的目录存在"""
|
|||
|
|
os.makedirs(self.results_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
def add_task(self, url, months=6, use_proxy=False, proxy_api_url=None, username=None, articles_only=True):
|
|||
|
|
"""添加新任务到队列
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
url: 百家号URL
|
|||
|
|
months: 获取月数
|
|||
|
|
use_proxy: 是否使用代理
|
|||
|
|
proxy_api_url: 代理API地址
|
|||
|
|
username: 用户名
|
|||
|
|
articles_only: 是否仅爬取文章（跳过视频）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
task_id: 任务ID
|
|||
|
|
"""
|
|||
|
|
with self.lock:
|
|||
|
|
task_id = f"task_{int(time.time() * 1000)}"
|
|||
|
|
created_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|||
|
|
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
cursor.execute('''
|
|||
|
|
INSERT INTO tasks (
|
|||
|
|
task_id, url, months, use_proxy, proxy_api_url,
|
|||
|
|
username, status, created_at, progress, current_step,
|
|||
|
|
total_articles, processed_articles, articles_only
|
|||
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|||
|
|
''', (
|
|||
|
|
task_id, url, months, 1 if use_proxy else 0, proxy_api_url,
|
|||
|
|
username, TaskStatus.PENDING.value, created_at, 0, "等待处理",
|
|||
|
|
0, 0, 1 if articles_only else 0
|
|||
|
|
))
|
|||
|
|
conn.commit()
|
|||
|
|
|
|||
|
|
logger.info(f"添加任务: {task_id}")
|
|||
|
|
return task_id
|
|||
|
|
|
|||
|
|
def get_task(self, task_id):
|
|||
|
|
"""获取任务信息"""
|
|||
|
|
with self.lock:
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
cursor.execute(
|
|||
|
|
"SELECT * FROM tasks WHERE task_id = ?",
|
|||
|
|
(task_id,)
|
|||
|
|
)
|
|||
|
|
row = cursor.fetchone()
|
|||
|
|
|
|||
|
|
if row:
|
|||
|
|
task = dict(row)
|
|||
|
|
# 将 use_proxy 从整数转换为布尔值
|
|||
|
|
task['use_proxy'] = bool(task['use_proxy'])
|
|||
|
|
# 将 articles_only 从整数转换为布尔值
|
|||
|
|
task['articles_only'] = bool(task.get('articles_only', 1))
|
|||
|
|
return task
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def get_all_tasks(self, username=None):
|
|||
|
|
"""获取所有任务（可按用户过滤）"""
|
|||
|
|
with self.lock:
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
if username:
|
|||
|
|
cursor.execute(
|
|||
|
|
"SELECT * FROM tasks WHERE username = ? ORDER BY created_at DESC",
|
|||
|
|
(username,)
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
cursor.execute("SELECT * FROM tasks ORDER BY created_at DESC")
|
|||
|
|
|
|||
|
|
rows = cursor.fetchall()
|
|||
|
|
tasks = []
|
|||
|
|
for row in rows:
|
|||
|
|
task = dict(row)
|
|||
|
|
# 将 use_proxy 从整数转换为布尔值
|
|||
|
|
task['use_proxy'] = bool(task['use_proxy'])
|
|||
|
|
# 将 articles_only 从整数转换为布尔值
|
|||
|
|
task['articles_only'] = bool(task.get('articles_only', 1))
|
|||
|
|
tasks.append(task)
|
|||
|
|
return tasks
|
|||
|
|
|
|||
|
|
def get_pending_task(self):
|
|||
|
|
"""获取下一个待处理的任务（包括检查暂停任务是否可恢复）"""
|
|||
|
|
with self.lock:
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
# 首先检查是否有暂停任务需要恢复
|
|||
|
|
from datetime import datetime, timedelta
|
|||
|
|
current_time = datetime.now()
|
|||
|
|
|
|||
|
|
cursor.execute(
|
|||
|
|
"SELECT * FROM tasks WHERE status = ? ORDER BY paused_at ASC",
|
|||
|
|
(TaskStatus.PAUSED.value,)
|
|||
|
|
)
|
|||
|
|
paused_tasks = cursor.fetchall()
|
|||
|
|
|
|||
|
|
for row in paused_tasks:
|
|||
|
|
paused_at_str = row['paused_at']
|
|||
|
|
if paused_at_str:
|
|||
|
|
paused_at = datetime.strptime(paused_at_str, '%Y-%m-%d %H:%M:%S')
|
|||
|
|
# 检查是否已经暂停超过10分钟
|
|||
|
|
if current_time - paused_at >= timedelta(minutes=10):
|
|||
|
|
task_id = row['task_id']
|
|||
|
|
# 恢复任务为待处理状态（保留 last_page 和 last_ctime）
|
|||
|
|
cursor.execute("""
|
|||
|
|
UPDATE tasks SET
|
|||
|
|
status = ?,
|
|||
|
|
current_step = ?,
|
|||
|
|
retry_count = 0,
|
|||
|
|
paused_at = NULL
|
|||
|
|
WHERE task_id = ?
|
|||
|
|
""", (TaskStatus.PENDING.value, "等待处理（从断点继续）", task_id))
|
|||
|
|
conn.commit()
|
|||
|
|
logger.info(f"任务 {task_id} 已从暂停状态恢复，将从第{row.get('last_page', 1)}页继续")
|
|||
|
|
|
|||
|
|
# 获取待处理任务
|
|||
|
|
cursor.execute(
|
|||
|
|
"SELECT * FROM tasks WHERE status = ? ORDER BY created_at ASC LIMIT 1",
|
|||
|
|
(TaskStatus.PENDING.value,)
|
|||
|
|
)
|
|||
|
|
row = cursor.fetchone()
|
|||
|
|
|
|||
|
|
if row:
|
|||
|
|
task = dict(row)
|
|||
|
|
# 将 use_proxy 从整数转换为布尔值
|
|||
|
|
task['use_proxy'] = bool(task['use_proxy'])
|
|||
|
|
# 将 articles_only 从整数转换为布尔值
|
|||
|
|
task['articles_only'] = bool(task.get('articles_only', 1))
|
|||
|
|
return task
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def update_task_status(self, task_id, status, **kwargs):
|
|||
|
|
"""更新任务状态
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
task_id: 任务ID
|
|||
|
|
status: 新状态
|
|||
|
|
**kwargs: 其他要更新的字段
|
|||
|
|
"""
|
|||
|
|
with self.lock:
|
|||
|
|
status_value = status.value if isinstance(status, TaskStatus) else status
|
|||
|
|
|
|||
|
|
# 准备更新字段
|
|||
|
|
update_fields = {"status": status_value}
|
|||
|
|
|
|||
|
|
# 更新时间戳
|
|||
|
|
if status == TaskStatus.PROCESSING:
|
|||
|
|
update_fields["started_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|||
|
|
elif status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
|
|||
|
|
update_fields["completed_at"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|||
|
|
|
|||
|
|
# 合并其他字段
|
|||
|
|
update_fields.update(kwargs)
|
|||
|
|
|
|||
|
|
# 构建 SQL 更新语句
|
|||
|
|
set_clause = ", ".join([f"{key} = ?" for key in update_fields.keys()])
|
|||
|
|
values = list(update_fields.values()) + [task_id]
|
|||
|
|
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
cursor.execute(
|
|||
|
|
f"UPDATE tasks SET {set_clause} WHERE task_id = ?",
|
|||
|
|
values
|
|||
|
|
)
|
|||
|
|
conn.commit()
|
|||
|
|
return cursor.rowcount > 0
|
|||
|
|
|
|||
|
|
def update_task_progress(self, task_id, progress, current_step=None, processed_articles=None):
|
|||
|
|
"""更新任务进度
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
task_id: 任务ID
|
|||
|
|
progress: 进度百分比 (0-100)
|
|||
|
|
current_step: 当前步骤描述
|
|||
|
|
processed_articles: 已处理文章数
|
|||
|
|
"""
|
|||
|
|
with self.lock:
|
|||
|
|
update_fields = {
|
|||
|
|
"progress": min(100, max(0, progress))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if current_step is not None:
|
|||
|
|
update_fields["current_step"] = current_step
|
|||
|
|
if processed_articles is not None:
|
|||
|
|
update_fields["processed_articles"] = processed_articles
|
|||
|
|
|
|||
|
|
set_clause = ", ".join([f"{key} = ?" for key in update_fields.keys()])
|
|||
|
|
values = list(update_fields.values()) + [task_id]
|
|||
|
|
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
cursor.execute(
|
|||
|
|
f"UPDATE tasks SET {set_clause} WHERE task_id = ?",
|
|||
|
|
values
|
|||
|
|
)
|
|||
|
|
conn.commit()
|
|||
|
|
return cursor.rowcount > 0
|
|||
|
|
|
|||
|
|
def get_queue_stats(self, username=None):
|
|||
|
|
"""获取队列统计信息"""
|
|||
|
|
with self.lock:
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
# 基础查询
|
|||
|
|
if username:
|
|||
|
|
base_query = "SELECT status, COUNT(*) as count FROM tasks WHERE username = ? GROUP BY status"
|
|||
|
|
cursor.execute(base_query, (username,))
|
|||
|
|
else:
|
|||
|
|
base_query = "SELECT status, COUNT(*) as count FROM tasks GROUP BY status"
|
|||
|
|
cursor.execute(base_query)
|
|||
|
|
|
|||
|
|
# 统计各状态数量
|
|||
|
|
status_counts = {row["status"]: row["count"] for row in cursor.fetchall()}
|
|||
|
|
|
|||
|
|
# 获取总数
|
|||
|
|
if username:
|
|||
|
|
cursor.execute("SELECT COUNT(*) as total FROM tasks WHERE username = ?", (username,))
|
|||
|
|
else:
|
|||
|
|
cursor.execute("SELECT COUNT(*) as total FROM tasks")
|
|||
|
|
total = cursor.fetchone()["total"]
|
|||
|
|
|
|||
|
|
stats = {
|
|||
|
|
"total": total,
|
|||
|
|
"pending": status_counts.get(TaskStatus.PENDING.value, 0),
|
|||
|
|
"processing": status_counts.get(TaskStatus.PROCESSING.value, 0),
|
|||
|
|
"completed": status_counts.get(TaskStatus.COMPLETED.value, 0),
|
|||
|
|
"failed": status_counts.get(TaskStatus.FAILED.value, 0),
|
|||
|
|
"paused": status_counts.get(TaskStatus.PAUSED.value, 0)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return stats
|
|||
|
|
|
|||
|
|
def delete_task(self, task_id):
|
|||
|
|
"""删除任务（先自动终止再删除）"""
|
|||
|
|
with self.lock:
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
# 检查任务是否存在
|
|||
|
|
cursor.execute("SELECT status FROM tasks WHERE task_id = ?", (task_id,))
|
|||
|
|
row = cursor.fetchone()
|
|||
|
|
|
|||
|
|
if not row:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 如果任务还在运行，先终止
|
|||
|
|
if row["status"] in [TaskStatus.PENDING.value, TaskStatus.PROCESSING.value]:
|
|||
|
|
cursor.execute('''
|
|||
|
|
UPDATE tasks SET
|
|||
|
|
status = ?,
|
|||
|
|
error = ?,
|
|||
|
|
current_step = ?,
|
|||
|
|
completed_at = ?
|
|||
|
|
WHERE task_id = ?
|
|||
|
|
''', (
|
|||
|
|
TaskStatus.FAILED.value,
|
|||
|
|
"任务已被用户删除",
|
|||
|
|
"任务已终止",
|
|||
|
|
datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|||
|
|
task_id
|
|||
|
|
))
|
|||
|
|
conn.commit()
|
|||
|
|
logger.info(f"终止任务: {task_id}")
|
|||
|
|
|
|||
|
|
# 然后从数据库中删除
|
|||
|
|
cursor.execute("DELETE FROM tasks WHERE task_id = ?", (task_id,))
|
|||
|
|
conn.commit()
|
|||
|
|
|
|||
|
|
if cursor.rowcount > 0:
|
|||
|
|
logger.info(f"删除任务: {task_id}")
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def cancel_task(self, task_id):
|
|||
|
|
"""终止任务（将等待中或处理中任务标记为失败）"""
|
|||
|
|
with self.lock:
|
|||
|
|
with self.db.get_connection() as conn:
|
|||
|
|
cursor = conn.cursor()
|
|||
|
|
|
|||
|
|
# 检查任务状态
|
|||
|
|
cursor.execute("SELECT status FROM tasks WHERE task_id = ?", (task_id,))
|
|||
|
|
row = cursor.fetchone()
|
|||
|
|
|
|||
|
|
if row and row["status"] in [TaskStatus.PENDING.value, TaskStatus.PROCESSING.value]:
|
|||
|
|
cursor.execute('''
|
|||
|
|
UPDATE tasks SET
|
|||
|
|
status = ?,
|
|||
|
|
error = ?,
|
|||
|
|
current_step = ?,
|
|||
|
|
completed_at = ?
|
|||
|
|
WHERE task_id = ?
|
|||
|
|
''', (
|
|||
|
|
TaskStatus.FAILED.value,
|
|||
|
|
"任务已被用户终止",
|
|||
|
|
"任务已终止",
|
|||
|
|
datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|||
|
|
task_id
|
|||
|
|
))
|
|||
|
|
conn.commit()
|
|||
|
|
return cursor.rowcount > 0
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 全局队列实例
|
|||
|
|
_task_queue = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_task_queue():
|
|||
|
|
"""获取全局任务队列实例"""
|
|||
|
|
global _task_queue
|
|||
|
|
if _task_queue is None:
|
|||
|
|
_task_queue = TaskQueue()
|
|||
|
|
return _task_queue
|