115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
from __future__ import annotations
|
||
|
||
import threading
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from uuid import uuid4
|
||
|
||
from app.core.files import import_root, load_import_records, save_import_records
|
||
from app.services.import_worker_service import import_domains_from_path
|
||
|
||
|
||
_IMPORT_TASK_LOCK = threading.Lock()
|
||
|
||
|
||
def _now() -> str:
|
||
return datetime.now().isoformat(sep=" ", timespec="seconds")
|
||
|
||
|
||
def list_import_tasks() -> list[dict]:
|
||
return load_import_records()
|
||
|
||
|
||
def _save_tasks(tasks: list[dict]) -> None:
|
||
save_import_records(tasks)
|
||
|
||
|
||
def _update_task(task_id: str, **patch: object) -> dict | None:
|
||
with _IMPORT_TASK_LOCK:
|
||
tasks = load_import_records()
|
||
target = next((item for item in tasks if item["task_id"] == task_id), None)
|
||
if not target:
|
||
return None
|
||
target.update(patch)
|
||
target["updated_at"] = _now()
|
||
_save_tasks(tasks)
|
||
return dict(target)
|
||
|
||
|
||
def _run_import_task(task_id: str, file_path: str, source_type: int = 7) -> None:
|
||
_update_task(task_id, status="running", started_at=_now(), message="导入任务开始执行")
|
||
try:
|
||
result = import_domains_from_path(Path(file_path), source_type=source_type)
|
||
stats = result.get("stats", {})
|
||
_update_task(
|
||
task_id,
|
||
status="completed",
|
||
completed_at=_now(),
|
||
result=result,
|
||
message=(
|
||
f"导入完成：总数 {stats.get('total', 0)}，有效 {stats.get('valid', 0)}，"
|
||
f"新增 {stats.get('added', 0)}，已存在 {stats.get('exists', 0)}，无效 {stats.get('invalid', 0)}"
|
||
),
|
||
)
|
||
except Exception as exc:
|
||
_update_task(
|
||
task_id,
|
||
status="failed",
|
||
completed_at=_now(),
|
||
message=f"导入失败：{exc}",
|
||
)
|
||
|
||
|
||
def create_import_task(content: bytes, filename: str, source_type: int = 7) -> dict:
|
||
task_id = uuid4().hex
|
||
safe_name = Path(filename).name or "domains.txt"
|
||
target = import_root() / f"{task_id}_{safe_name}"
|
||
target.write_bytes(content)
|
||
|
||
record = {
|
||
"task_id": task_id,
|
||
"filename": safe_name,
|
||
"stored_path": str(target),
|
||
"source_type": source_type,
|
||
"status": "queued",
|
||
"message": "文件已接收，等待处理",
|
||
"created_at": _now(),
|
||
"updated_at": _now(),
|
||
"started_at": "",
|
||
"completed_at": "",
|
||
"result": None,
|
||
}
|
||
|
||
with _IMPORT_TASK_LOCK:
|
||
tasks = load_import_records()
|
||
tasks.insert(0, record)
|
||
_save_tasks(tasks)
|
||
|
||
worker = threading.Thread(target=_run_import_task, args=(task_id, str(target), source_type), daemon=True)
|
||
worker.start()
|
||
return record
|
||
|
||
|
||
def retry_import_task(task_id: str) -> dict:
|
||
with _IMPORT_TASK_LOCK:
|
||
tasks = load_import_records()
|
||
target = next((item for item in tasks if item["task_id"] == task_id), None)
|
||
if not target:
|
||
raise ValueError("导入任务不存在")
|
||
if target.get("status") == "running":
|
||
raise ValueError("导入任务正在运行，不能重复执行")
|
||
target["status"] = "queued"
|
||
target["message"] = "任务已重新加入队列"
|
||
target["started_at"] = ""
|
||
target["completed_at"] = ""
|
||
target["updated_at"] = _now()
|
||
target["result"] = None
|
||
_save_tasks(tasks)
|
||
stored_path = target["stored_path"]
|
||
source_type = int(target.get("source_type", 7))
|
||
record = dict(target)
|
||
|
||
worker = threading.Thread(target=_run_import_task, args=(task_id, stored_path, source_type), daemon=True)
|
||
worker.start()
|
||
return record
|