649 lines
28 KiB
Python
649 lines
28 KiB
Python
import hashlib
|
||
import json
|
||
import secrets
|
||
import string
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from config import Config
|
||
from crawler_client import (
|
||
CrawlerError,
|
||
cancel_collection,
|
||
query_collection,
|
||
resume_collection,
|
||
start_collection,
|
||
)
|
||
from database import connection, json_dumps, json_loads, now_iso
|
||
from extract_service import transform_result
|
||
|
||
|
||
MODE_NAMES = {"mild": "轻度", "regular": "常规", "deep": "深度"}
|
||
TERMINAL_STATUSES = {"completed", "failed", "cancelled", "partial"}
|
||
# 远程 overview/sources 初始化较重，同一业务进程内串行提交，
|
||
# 避免多个新任务同时触发问一问基础采集导致网关超时。
|
||
COLLECTION_START_LOCK = threading.Lock()
|
||
|
||
|
||
class TaskError(RuntimeError):
|
||
def __init__(self, code: str, message: str, status_code: int = 400, details: dict | None = None):
|
||
super().__init__(message)
|
||
self.code = code
|
||
self.message = message
|
||
self.status_code = status_code
|
||
self.details = details or {}
|
||
|
||
|
||
def _public_task_id() -> str:
|
||
day = datetime.now().strftime("%Y%m%d")
|
||
chars = string.ascii_uppercase + string.digits
|
||
return f"HC-{day}-{''.join(secrets.choice(chars) for _ in range(5))}"
|
||
|
||
|
||
def _file_id() -> str:
|
||
return f"file_{secrets.token_hex(8)}"
|
||
|
||
|
||
def validate_create(data: dict[str, Any]) -> tuple[str, str, list[str], str, str]:
|
||
mode = str(data.get("mode") or "")
|
||
if mode not in MODE_NAMES:
|
||
raise TaskError("INVALID_MODE", "mode 必须是 mild、regular 或 deep", 422)
|
||
platform = str(data.get("platform") or "")
|
||
if platform != "xiaohongshu_wenyiwen":
|
||
raise TaskError("INVALID_PLATFORM", "第一期仅支持 xiaohongshu_wenyiwen", 422)
|
||
raw_keywords = data.get("keywords")
|
||
if not isinstance(raw_keywords, list):
|
||
raise TaskError("KEYWORD_REQUIRED", "keywords 必须是数组", 422)
|
||
keywords = list(dict.fromkeys(str(item).strip() for item in raw_keywords if str(item).strip()))
|
||
if not keywords:
|
||
raise TaskError("KEYWORD_REQUIRED", "至少需要一个搜索词", 422)
|
||
if mode == "mild" and len(keywords) > 50:
|
||
raise TaskError("KEYWORD_LIMIT_EXCEEDED", "轻度任务最多支持 50 个搜索词", 422)
|
||
if mode in {"regular", "deep"} and len(keywords) != 1:
|
||
raise TaskError("SINGLE_KEYWORD_REQUIRED", "常规和深度任务只能包含一个搜索词", 422)
|
||
if any(len(keyword) > 200 for keyword in keywords):
|
||
raise TaskError("INVALID_KEYWORD", "单个搜索词不能超过 200 个字符", 422)
|
||
name = str(data.get("name") or "").strip() or f"{MODE_NAMES[mode]}｜问一问｜{keywords[0]}"
|
||
if len(name) > 100:
|
||
raise TaskError("INVALID_TASK_NAME", "任务名称不能超过 100 个字符", 422)
|
||
created_by_id = str(data.get("created_by_id") or "usr_local")
|
||
created_by_name = str(data.get("created_by_name") or "本地用户")
|
||
return mode, platform, keywords, name, created_by_id, created_by_name
|
||
|
||
|
||
def create_task(data: dict[str, Any], retry_of_task_id: int | None = None) -> dict[str, Any]:
|
||
mode, platform, keywords, name, creator_id, creator_name = validate_create(data)
|
||
now = now_iso()
|
||
with connection() as conn:
|
||
while True:
|
||
public_id = _public_task_id()
|
||
if not conn.execute("SELECT 1 FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone():
|
||
break
|
||
cursor = conn.execute(
|
||
"""INSERT INTO research_tasks
|
||
(public_task_id,name,platform,mode,status,execution_stage,progress_percent,
|
||
created_by_id,created_by_name,retry_of_task_id,created_at,updated_at)
|
||
VALUES (?,?,?,?, 'waiting','queued',0,?,?,?,?,?)""",
|
||
(public_id, name, platform, mode, creator_id, creator_name, retry_of_task_id, now, now),
|
||
)
|
||
task_pk = cursor.lastrowid
|
||
conn.executemany(
|
||
"""INSERT INTO task_keywords
|
||
(task_id,keyword_order,keyword,normalized_keyword,created_at)
|
||
VALUES (?,?,?,?,?)""",
|
||
[(task_pk, index, keyword, keyword.casefold(), now) for index, keyword in enumerate(keywords)],
|
||
)
|
||
_event(conn, task_pk, "created", "queued", 0, "user", creator_id, {"keywords": keywords})
|
||
task = get_task(public_id)
|
||
threading.Thread(target=_execute_task, args=(public_id,), daemon=True, name=f"holy-crab-{public_id}").start()
|
||
return task
|
||
|
||
|
||
def _event(conn, task_pk: int, event_type: str, stage: str, progress: int, actor_type: str,
|
||
actor_id: str | None = None, payload: dict | None = None) -> None:
|
||
conn.execute(
|
||
"""INSERT INTO task_events
|
||
(task_id,event_type,stage,progress_percent,actor_type,actor_id,payload_json,occurred_at)
|
||
VALUES (?,?,?,?,?,?,?,?)""",
|
||
(task_pk, event_type, stage, progress, actor_type, actor_id, json_dumps(payload) if payload else None, now_iso()),
|
||
)
|
||
|
||
|
||
def _keywords(conn, task_pk: int) -> list[str]:
|
||
rows = conn.execute(
|
||
"SELECT keyword FROM task_keywords WHERE task_id=? ORDER BY keyword_order", (task_pk,)
|
||
).fetchall()
|
||
return [row["keyword"] for row in rows]
|
||
|
||
|
||
def _update(public_id: str, *, status: str | None = None, stage: str | None = None,
|
||
progress: int | float | None = None, crawler_task_id: str | None = None,
|
||
result: dict | None = None, error_code: str | None = None,
|
||
error_message: str | None = None, completed: bool = False) -> None:
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row:
|
||
return
|
||
# cancelled 任务仍允许写入最终部分结果，但禁止恢复为其他状态。
|
||
if row["status"] == "cancelled" and result is None:
|
||
return
|
||
next_status = "cancelled" if row["status"] == "cancelled" else (status or row["status"])
|
||
next_progress = max(row["progress_percent"], min(100, int(float(progress or 0))))
|
||
now = now_iso()
|
||
conn.execute(
|
||
"""UPDATE research_tasks SET status=?,execution_stage=?,progress_percent=?,
|
||
crawler_task_id=COALESCE(?,crawler_task_id),
|
||
result_json=COALESCE(?,result_json),error_code=?,error_message=?,
|
||
started_at=COALESCE(started_at,?),
|
||
completed_at=CASE WHEN ? THEN ? ELSE completed_at END,
|
||
updated_at=?,version=version+1 WHERE id=?""",
|
||
(
|
||
next_status, stage or row["execution_stage"], next_progress,
|
||
crawler_task_id, json_dumps(result) if result is not None else None,
|
||
error_code, error_message, now, completed, now, now, row["id"],
|
||
),
|
||
)
|
||
_event(conn, row["id"], "status_changed", stage or row["execution_stage"], next_progress,
|
||
"worker", payload={"status": next_status, "error_code": error_code})
|
||
|
||
|
||
def _execute_task(public_id: str) -> None:
|
||
try:
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row or row["status"] == "cancelled":
|
||
return
|
||
keywords = _keywords(conn, row["id"])
|
||
mode = row["mode"]
|
||
_update(public_id, status="running", stage="initializing", progress=5)
|
||
with COLLECTION_START_LOCK:
|
||
result = start_collection(mode, keywords)
|
||
crawler_task_id = str(result.get("task_id") or "")
|
||
if not crawler_task_id:
|
||
raise CrawlerError("爬虫服务未返回 task_id")
|
||
with connection() as conn:
|
||
local = conn.execute(
|
||
"SELECT status FROM research_tasks WHERE public_task_id=?", (public_id,)
|
||
).fetchone()
|
||
if local and local["status"] == "cancelled":
|
||
cancel_collection(crawler_task_id)
|
||
return
|
||
|
||
crawler_status = str(result.get("status") or "")
|
||
is_async = mode == "deep" or crawler_status in {"running", "waiting"}
|
||
if not is_async:
|
||
_complete_task(public_id, result, crawler_task_id=crawler_task_id)
|
||
return
|
||
|
||
_update(
|
||
public_id, status="running", stage="collecting_search",
|
||
progress=max(10, float(result.get("progress") or 10)),
|
||
crawler_task_id=crawler_task_id, result=result,
|
||
)
|
||
while True:
|
||
time.sleep(Config.CRAWLER_POLL_INTERVAL_SECONDS)
|
||
with connection() as conn:
|
||
local = conn.execute(
|
||
"SELECT status FROM research_tasks WHERE public_task_id=?", (public_id,)
|
||
).fetchone()
|
||
if not local or local["status"] == "cancelled":
|
||
return
|
||
queried = query_collection(crawler_task_id)
|
||
status = str(queried.get("status") or "running")
|
||
if status in {"completed", "partial"}:
|
||
_complete_task(
|
||
public_id, queried, status=status,
|
||
crawler_task_id=crawler_task_id,
|
||
)
|
||
return
|
||
if status == "failed":
|
||
raise CrawlerError(str(queried.get("error_message") or "爬虫任务执行失败"))
|
||
_update(
|
||
public_id, status="running",
|
||
stage=str(queried.get("current_stage") or "collecting_search"),
|
||
progress=float(queried.get("progress") or 10),
|
||
crawler_task_id=crawler_task_id, result=queried,
|
||
)
|
||
except Exception as exc:
|
||
_update(
|
||
public_id, status="failed", stage="failed", progress=0,
|
||
error_code="COLLECTOR_FAILED", error_message=str(exc), completed=True,
|
||
)
|
||
|
||
|
||
def _complete_task(
|
||
public_id: str,
|
||
result: dict[str, Any],
|
||
status: str = "completed",
|
||
crawler_task_id: str | None = None,
|
||
) -> None:
|
||
# 对原始爬虫结果执行数据清洗，生成结构化的 extracted 字段
|
||
result = transform_result(result)
|
||
_write_result_files(public_id, result)
|
||
_update(
|
||
public_id, status=status, stage="completed", progress=100,
|
||
crawler_task_id=crawler_task_id, result=result, completed=True,
|
||
error_code=("PARTIAL_RESULT" if status == "partial" else None),
|
||
error_message=("部分数据采集失败" if status == "partial" else None),
|
||
)
|
||
|
||
|
||
def _write_result_files(public_id: str, result: dict[str, Any]) -> None:
|
||
folder = Config.FILE_STORAGE_PATH / public_id
|
||
folder.mkdir(parents=True, exist_ok=True)
|
||
json_path = folder / "xhs_wen_result.json"
|
||
summary_path = folder / "research_summary.md"
|
||
json_bytes = json.dumps(result, ensure_ascii=False, indent=2).encode("utf-8")
|
||
json_path.write_bytes(json_bytes)
|
||
summary = (
|
||
f"# 问一问任务数据摘要\n\n"
|
||
f"- 任务 ID：{public_id}\n"
|
||
f"- 生成时间：{now_iso()}\n"
|
||
f"- 数据文件：xhs_wen_result.json\n"
|
||
f"- SHA-256：{hashlib.sha256(json_bytes).hexdigest()}\n"
|
||
)
|
||
summary_path.write_text(summary, encoding="utf-8")
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT id FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row:
|
||
return
|
||
for path, stage, category, fmt, count in [
|
||
(json_path, "cleaned", "search_result", "json", _result_count(result)),
|
||
(summary_path, "cleaned", "metadata", "md", 1),
|
||
]:
|
||
content = path.read_bytes()
|
||
existing = conn.execute(
|
||
"SELECT public_file_id FROM task_files WHERE task_id=? AND name=?",
|
||
(row["id"], path.name),
|
||
).fetchone()
|
||
public_file_id = existing["public_file_id"] if existing else _file_id()
|
||
conn.execute(
|
||
"""INSERT OR REPLACE INTO task_files
|
||
(public_file_id,task_id,name,stage,category,format,record_count,size_bytes,
|
||
storage_key,sha256,schema_version,generated_at,available)
|
||
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)""",
|
||
(
|
||
public_file_id, row["id"], path.name, stage, category, fmt,
|
||
count, len(content), str(path), hashlib.sha256(content).hexdigest(),
|
||
"1.0", now_iso(), 1,
|
||
),
|
||
)
|
||
|
||
|
||
def _result_count(result: dict[str, Any]) -> int:
|
||
data = result.get("data")
|
||
if isinstance(data, list):
|
||
return len(data)
|
||
items = result.get("items")
|
||
return len(items) if isinstance(items, list) else 1
|
||
|
||
|
||
def _task_dict(conn, row, include_result: bool = True) -> dict[str, Any]:
|
||
keywords = _keywords(conn, row["id"])
|
||
saved_result = json_loads(row["result_json"]) if row["result_json"] else None
|
||
# 兼容修复前生成的同步任务：爬虫任务 ID 曾只保存在 result.task_id。
|
||
collector_task_id = row["crawler_task_id"]
|
||
if not collector_task_id and isinstance(saved_result, dict):
|
||
collector_task_id = saved_result.get("task_id")
|
||
files = conn.execute(
|
||
"SELECT * FROM task_files WHERE task_id=? AND available=1 ORDER BY id", (row["id"],)
|
||
).fetchall()
|
||
data = {
|
||
"task_id": row["public_task_id"],
|
||
"collector_task_id": collector_task_id,
|
||
"name": row["name"],
|
||
"mode": row["mode"],
|
||
"platform": {"code": row["platform"], "name": "问一问"},
|
||
"keywords": keywords,
|
||
"status": row["status"],
|
||
"execution_stage": row["execution_stage"],
|
||
"progress_percent": row["progress_percent"],
|
||
"created_by": {"id": row["created_by_id"], "name": row["created_by_name"]},
|
||
"created_at": row["created_at"],
|
||
"updated_at": row["updated_at"],
|
||
"started_at": row["started_at"],
|
||
"completed_at": row["completed_at"],
|
||
"retry_of_task_id": _retry_public_id(conn, row["retry_of_task_id"]),
|
||
"cancel_requested": bool(row["cancel_requested_at"]),
|
||
"error": (
|
||
{"code": row["error_code"], "message": row["error_message"]}
|
||
if row["error_code"] or row["error_message"] else None
|
||
),
|
||
"usable_by_agent": row["status"] == "completed",
|
||
"files": [_file_dict(file) for file in files],
|
||
}
|
||
if include_result and saved_result is not None:
|
||
# 兜底：对历史任务（未经过 transform_result 的旧数据）
|
||
# 也执行清洗转换，确保前端 / AI 始终拿到带 extracted 的完整结果
|
||
data["result"] = transform_result(saved_result)
|
||
return data
|
||
|
||
|
||
def _retry_public_id(conn, retry_pk: int | None) -> str | None:
|
||
if not retry_pk:
|
||
return None
|
||
row = conn.execute("SELECT public_task_id FROM research_tasks WHERE id=?", (retry_pk,)).fetchone()
|
||
return row["public_task_id"] if row else None
|
||
|
||
|
||
def _file_dict(row) -> dict[str, Any]:
|
||
return {
|
||
"file_id": row["public_file_id"], "name": row["name"], "format": row["format"],
|
||
"stage": row["stage"], "stage_name": "清洗数据" if row["stage"] == "cleaned" else "原始数据",
|
||
"category": row["category"], "category_name": {
|
||
"search_result": "搜索结果", "entry": "收录内容",
|
||
"comment": "评论", "metadata": "任务元数据",
|
||
}.get(row["category"], row["category"]),
|
||
"record_count": row["record_count"], "size_bytes": row["size_bytes"],
|
||
"sha256": row["sha256"], "schema_version": row["schema_version"],
|
||
"generated_at": row["generated_at"], "previewable": row["size_bytes"] <= 2 * 1024 * 1024,
|
||
}
|
||
|
||
|
||
def get_task(public_id: str) -> dict[str, Any]:
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row:
|
||
raise TaskError("TASK_NOT_FOUND", "任务不存在", 404)
|
||
return _task_dict(conn, row)
|
||
|
||
|
||
def get_task_extracted(public_id: str) -> dict[str, Any]:
|
||
"""直接返回任务的清洗解析结果，方便 curl / Agent 读取。
|
||
|
||
普通任务详情接口需要服务前端展示，因此会包装任务状态、文件列表和原始结果。
|
||
Agent 做分析时更适合直接读取 result.extracted，这里把 extracted 提升到
|
||
data 顶层返回。
|
||
"""
|
||
task = get_task(public_id)
|
||
result = task.get("result") if isinstance(task, dict) else None
|
||
extracted = result.get("extracted") if isinstance(result, dict) else None
|
||
if not isinstance(extracted, dict):
|
||
raise TaskError("EXTRACTED_NOT_READY", "任务清洗数据尚未生成", 404)
|
||
extracted = dict(extracted)
|
||
metadata = dict(extracted.get("提取元数据") or {})
|
||
metadata.setdefault("来源任务ID", public_id)
|
||
metadata.setdefault("核心数据路径", "data.result.extracted")
|
||
metadata.setdefault("采集任务ID", task.get("collector_task_id"))
|
||
metadata.setdefault("任务状态", task.get("status"))
|
||
metadata.setdefault("关键词", task.get("keywords"))
|
||
extracted["提取元数据"] = metadata
|
||
return extracted
|
||
|
||
|
||
def list_tasks(q: str = "", status: str = "", page: int = 1, page_size: int = 10) -> dict[str, Any]:
|
||
if page < 1 or page_size not in {10, 20, 50}:
|
||
raise TaskError("INVALID_PAGINATION", "page 必须大于 0，page_size 只允许 10、20、50", 422)
|
||
clauses, params = [], []
|
||
if status:
|
||
clauses.append("t.status=?")
|
||
params.append(status)
|
||
if q:
|
||
clauses.append(
|
||
"(t.name LIKE ? OR t.public_task_id LIKE ? OR EXISTS "
|
||
"(SELECT 1 FROM task_keywords k WHERE k.task_id=t.id AND k.keyword LIKE ?))"
|
||
)
|
||
like = f"%{q}%"
|
||
params.extend([like, like, like])
|
||
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
|
||
with connection() as conn:
|
||
total = conn.execute(f"SELECT COUNT(*) AS count FROM research_tasks t {where}", params).fetchone()["count"]
|
||
rows = conn.execute(
|
||
f"SELECT t.* FROM research_tasks t {where} ORDER BY t.created_at DESC LIMIT ? OFFSET ?",
|
||
[*params, page_size, (page - 1) * page_size],
|
||
).fetchall()
|
||
return {
|
||
"items": [_task_dict(conn, row, include_result=False) for row in rows],
|
||
"page": page, "page_size": page_size, "total": total,
|
||
"total_pages": max(1, (total + page_size - 1) // page_size),
|
||
}
|
||
|
||
|
||
def cancel_task(public_id: str, reason: str = "") -> dict[str, Any]:
|
||
crawler_task_id = None
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row:
|
||
raise TaskError("TASK_NOT_FOUND", "任务不存在", 404)
|
||
if row["status"] == "cancelled":
|
||
return _task_dict(conn, row)
|
||
if row["status"] not in {"waiting", "running"}:
|
||
raise TaskError("TASK_STATUS_CONFLICT", "当前任务状态不允许取消", 409, {"current_status": row["status"]})
|
||
crawler_task_id = row["crawler_task_id"]
|
||
if not crawler_task_id and row["result_json"]:
|
||
saved_result = json_loads(row["result_json"])
|
||
if isinstance(saved_result, dict):
|
||
crawler_task_id = saved_result.get("task_id")
|
||
if crawler_task_id:
|
||
try:
|
||
cancel_collection(str(crawler_task_id))
|
||
except CrawlerError as exc:
|
||
raise TaskError("COLLECTOR_UNAVAILABLE", f"爬虫任务停止失败：{exc}", 503) from exc
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row:
|
||
raise TaskError("TASK_NOT_FOUND", "任务不存在", 404)
|
||
if row["status"] not in {"waiting", "running", "cancelled"}:
|
||
raise TaskError("TASK_STATUS_CONFLICT", "任务状态已变化，无法停止", 409)
|
||
now = now_iso()
|
||
conn.execute(
|
||
"""UPDATE research_tasks SET status='cancelled',execution_stage='cancelled',
|
||
cancel_requested_at=?,completed_at=?,updated_at=?,version=version+1 WHERE id=?""",
|
||
(now, now, now, row["id"]),
|
||
)
|
||
_event(conn, row["id"], "cancelled", "cancelled", row["progress_percent"], "user",
|
||
payload={"reason": reason, "crawler_task_id": crawler_task_id})
|
||
if crawler_task_id:
|
||
threading.Thread(
|
||
target=_sync_cancelled_result,
|
||
args=(public_id, str(crawler_task_id)),
|
||
daemon=True,
|
||
name=f"holy-crab-cancel-{public_id}",
|
||
).start()
|
||
return get_task(public_id)
|
||
|
||
|
||
def _sync_cancelled_result(public_id: str, crawler_task_id: str) -> None:
|
||
"""等待爬虫保存停止快照，并同步到业务任务结果和文件。"""
|
||
|
||
deadline = time.time() + max(60, Config.CRAWLER_TIMEOUT_SECONDS)
|
||
latest = None
|
||
while time.time() < deadline:
|
||
try:
|
||
latest = query_collection(crawler_task_id)
|
||
except CrawlerError:
|
||
time.sleep(Config.CRAWLER_POLL_INTERVAL_SECONDS)
|
||
continue
|
||
if (
|
||
latest.get("status") == "cancelled"
|
||
and latest.get("current_stage") == "cancelled"
|
||
and latest.get("data") is not None
|
||
):
|
||
break
|
||
time.sleep(Config.CRAWLER_POLL_INTERVAL_SECONDS)
|
||
if latest and latest.get("data") is not None:
|
||
_write_result_files(public_id, latest)
|
||
_update(
|
||
public_id, status="cancelled", stage="cancelled",
|
||
progress=float(latest.get("progress") or 0),
|
||
crawler_task_id=crawler_task_id, result=latest, completed=True,
|
||
error_code="TASK_CANCELLED",
|
||
error_message="任务已停止，返回停止前已采集数据",
|
||
)
|
||
|
||
|
||
def resume_task(public_id: str) -> dict[str, Any]:
|
||
"""继续原业务任务：有爬虫任务 ID 时断点恢复，否则从初始化阶段重试。"""
|
||
with connection() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)
|
||
).fetchone()
|
||
if not row:
|
||
raise TaskError("TASK_NOT_FOUND", "任务不存在", 404)
|
||
if row["status"] not in {"failed", "cancelled", "partial"}:
|
||
raise TaskError(
|
||
"TASK_STATUS_CONFLICT",
|
||
"只有失败、部分完成或已取消任务可以继续采集",
|
||
409,
|
||
)
|
||
crawler_task_id = row["crawler_task_id"]
|
||
if not crawler_task_id and row["result_json"]:
|
||
saved_result = json_loads(row["result_json"])
|
||
if isinstance(saved_result, dict):
|
||
crawler_task_id = saved_result.get("task_id")
|
||
if not crawler_task_id:
|
||
# 首次请求爬虫服务就失败时不会产生爬虫任务 ID，此时仍沿用
|
||
# 原业务任务 ID，从初始化阶段重试，避免“继续采集”创建新任务。
|
||
now = now_iso()
|
||
conn.execute(
|
||
"""UPDATE research_tasks
|
||
SET status='waiting',execution_stage='queued',progress_percent=0,
|
||
cancel_requested_at=NULL,completed_at=NULL,error_code=NULL,
|
||
error_message=NULL,result_json=NULL,updated_at=?,version=version+1
|
||
WHERE id=?""",
|
||
(now, row["id"]),
|
||
)
|
||
_event(
|
||
conn, row["id"], "retried", "queued", 0, "user",
|
||
payload={"resume_mode": "restart_initialization"},
|
||
)
|
||
|
||
if not crawler_task_id:
|
||
threading.Thread(
|
||
target=_execute_task,
|
||
args=(public_id,),
|
||
daemon=True,
|
||
name=f"holy-crab-retry-{public_id}",
|
||
).start()
|
||
return get_task(public_id)
|
||
|
||
try:
|
||
resumed = resume_collection(str(crawler_task_id))
|
||
except CrawlerError as exc:
|
||
raise TaskError(
|
||
"COLLECTOR_UNAVAILABLE", f"爬虫任务恢复失败：{exc}", 503
|
||
) from exc
|
||
|
||
with connection() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)
|
||
).fetchone()
|
||
if not row:
|
||
raise TaskError("TASK_NOT_FOUND", "任务不存在", 404)
|
||
now = now_iso()
|
||
progress = max(0, min(99, int(float(resumed.get("progress") or row["progress_percent"] or 0))))
|
||
conn.execute(
|
||
"""UPDATE research_tasks
|
||
SET status='running',execution_stage='resuming',progress_percent=?,
|
||
cancel_requested_at=NULL,completed_at=NULL,error_code=NULL,
|
||
error_message=NULL,result_json=?,updated_at=?,version=version+1
|
||
WHERE id=?""",
|
||
(progress, json_dumps(resumed), now, row["id"]),
|
||
)
|
||
_event(
|
||
conn, row["id"], "resumed", "resuming", progress, "user",
|
||
payload={"crawler_task_id": crawler_task_id},
|
||
)
|
||
|
||
threading.Thread(
|
||
target=_monitor_resumed_task,
|
||
args=(public_id, str(crawler_task_id)),
|
||
daemon=True,
|
||
name=f"holy-crab-resume-{public_id}",
|
||
).start()
|
||
return get_task(public_id)
|
||
|
||
|
||
def _monitor_resumed_task(public_id: str, crawler_task_id: str) -> None:
|
||
"""轮询恢复后的原爬虫任务，并覆盖原业务任务的结果文件。"""
|
||
try:
|
||
while True:
|
||
time.sleep(Config.CRAWLER_POLL_INTERVAL_SECONDS)
|
||
with connection() as conn:
|
||
local = conn.execute(
|
||
"SELECT status FROM research_tasks WHERE public_task_id=?",
|
||
(public_id,),
|
||
).fetchone()
|
||
if not local or local["status"] == "cancelled":
|
||
return
|
||
queried = query_collection(crawler_task_id)
|
||
status = str(queried.get("status") or "running")
|
||
if status in {"completed", "partial"}:
|
||
_complete_task(
|
||
public_id, queried, status=status,
|
||
crawler_task_id=crawler_task_id,
|
||
)
|
||
return
|
||
if status == "failed":
|
||
raise CrawlerError(
|
||
str(queried.get("error_message") or "恢复采集失败")
|
||
)
|
||
_update(
|
||
public_id,
|
||
status="running",
|
||
stage=str(queried.get("current_stage") or "resuming"),
|
||
progress=float(queried.get("progress") or 0),
|
||
crawler_task_id=crawler_task_id,
|
||
result=queried,
|
||
)
|
||
except Exception as exc:
|
||
_update(
|
||
public_id,
|
||
status="failed",
|
||
stage="resume_failed",
|
||
progress=0,
|
||
error_code="COLLECTOR_RESUME_FAILED",
|
||
error_message=str(exc),
|
||
completed=True,
|
||
)
|
||
|
||
|
||
def recollect_task(public_id: str) -> dict[str, Any]:
|
||
"""复制原配置并创建全新的业务任务和爬虫任务，从头采集。"""
|
||
with connection() as conn:
|
||
row = conn.execute("SELECT * FROM research_tasks WHERE public_task_id=?", (public_id,)).fetchone()
|
||
if not row:
|
||
raise TaskError("TASK_NOT_FOUND", "任务不存在", 404)
|
||
if row["status"] in {"waiting", "running"}:
|
||
raise TaskError("TASK_STATUS_CONFLICT", "运行中的任务不能重新采集", 409)
|
||
keywords = _keywords(conn, row["id"])
|
||
data = {
|
||
"name": row["name"], "mode": row["mode"], "platform": row["platform"],
|
||
"keywords": keywords, "created_by_id": row["created_by_id"],
|
||
"created_by_name": row["created_by_name"],
|
||
}
|
||
source_pk = row["id"]
|
||
return create_task(data, retry_of_task_id=source_pk)
|
||
|
||
|
||
# 保留旧函数名供其他内部调用兼容；语义已经调整为断点恢复。
|
||
retry_task = resume_task
|
||
|
||
|
||
def get_file(public_id: str, file_id: str) -> tuple[dict[str, Any], Path]:
|
||
with connection() as conn:
|
||
row = conn.execute(
|
||
"""SELECT f.* FROM task_files f JOIN research_tasks t ON t.id=f.task_id
|
||
WHERE t.public_task_id=? AND f.public_file_id=? AND f.available=1""",
|
||
(public_id, file_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise TaskError("FILE_NOT_FOUND", "文件不存在", 404)
|
||
return _file_dict(row), Path(row["storage_key"])
|
||
|
||
|
||
def resume_unfinished_tasks() -> None:
|
||
"""服务重启后恢复等待任务；已有爬虫任务则继续查询，否则重新提交。"""
|
||
|
||
with connection() as conn:
|
||
rows = conn.execute(
|
||
"SELECT public_task_id FROM research_tasks WHERE status IN ('waiting','running')"
|
||
).fetchall()
|
||
for row in rows:
|
||
threading.Thread(
|
||
target=_execute_task, args=(row["public_task_id"],), daemon=True,
|
||
name=f"holy-crab-resume-{row['public_task_id']}",
|
||
).start()
|