200 lines
7.0 KiB
Python
200 lines
7.0 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()
|
||
_IMPORT_EXECUTION_LOCK = threading.Lock()
|
||
_SOURCE_TYPE_LABELS = {
|
||
6: "手工录入",
|
||
7: "TXT 导入",
|
||
9: "其它",
|
||
}
|
||
|
||
|
||
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 _append_log_locked(target: dict, message: str) -> None:
|
||
target.setdefault("logs", [])
|
||
target["logs"].append(f"[{_now()}] {message}")
|
||
target["logs"] = target["logs"][-200:]
|
||
|
||
|
||
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 _update_task_with_log(task_id: str, log_message: 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)
|
||
_append_log_locked(target, log_message)
|
||
target["updated_at"] = _now()
|
||
_save_tasks(tasks)
|
||
return dict(target)
|
||
|
||
|
||
def _phase_label(phase: str) -> str:
|
||
mapping = {
|
||
"queued": "排队中",
|
||
"reading": "读取文件中",
|
||
"normalizing": "清洗中",
|
||
"importing": "入库中",
|
||
"completed": "已完成",
|
||
"failed": "失败",
|
||
}
|
||
return mapping.get(phase, phase)
|
||
|
||
|
||
def _source_label(source_type: int) -> str:
|
||
return _SOURCE_TYPE_LABELS.get(int(source_type or 7), "未知")
|
||
|
||
|
||
def _run_import_task(task_id: str, file_path: str, source_type: int = 7) -> None:
|
||
with _IMPORT_EXECUTION_LOCK:
|
||
_update_task_with_log(
|
||
task_id,
|
||
f"导入任务开始执行，来源类型：{_source_label(source_type)}",
|
||
status="running",
|
||
started_at=_now(),
|
||
message=f"导入任务开始执行，来源类型：{_source_label(source_type)}",
|
||
phase="reading",
|
||
phase_label=_phase_label("reading"),
|
||
)
|
||
try:
|
||
path = Path(file_path)
|
||
_update_task_with_log(
|
||
task_id,
|
||
f"开始读取文件：{path.name}",
|
||
phase="reading",
|
||
phase_label=_phase_label("reading"),
|
||
)
|
||
raw_lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
|
||
total_lines = len(raw_lines)
|
||
non_empty = sum(1 for line in raw_lines if line.strip())
|
||
_update_task_with_log(
|
||
task_id,
|
||
f"文件读取完成，共 {total_lines} 行，非空 {non_empty} 行",
|
||
phase="normalizing",
|
||
phase_label=_phase_label("normalizing"),
|
||
message=f"文件读取完成，准备清洗 {non_empty} 条域名",
|
||
)
|
||
|
||
result = import_domains_from_path(path, source_type=source_type)
|
||
stats = result.get("stats", {})
|
||
_update_task_with_log(
|
||
task_id,
|
||
(
|
||
f"导入完成：总数 {stats.get('total', 0)}，有效 {stats.get('valid', 0)}，"
|
||
f"新增 {stats.get('added', 0)}，已存在 {stats.get('exists', 0)}，无效 {stats.get('invalid', 0)}，"
|
||
f"来源类型 {result.get('source_label') or _source_label(source_type)}"
|
||
),
|
||
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)}"
|
||
),
|
||
phase="completed",
|
||
phase_label=_phase_label("completed"),
|
||
)
|
||
except Exception as exc:
|
||
_update_task_with_log(
|
||
task_id,
|
||
f"导入失败：{exc}",
|
||
status="failed",
|
||
completed_at=_now(),
|
||
message=f"导入失败：{exc}",
|
||
phase="failed",
|
||
phase_label=_phase_label("failed"),
|
||
)
|
||
|
||
|
||
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,
|
||
"source_label": _source_label(source_type),
|
||
"status": "queued",
|
||
"message": f"文件已接收，等待后台处理，来源类型：{_source_label(source_type)}",
|
||
"created_at": _now(),
|
||
"updated_at": _now(),
|
||
"started_at": "",
|
||
"completed_at": "",
|
||
"phase": "queued",
|
||
"phase_label": _phase_label("queued"),
|
||
"logs": [f"[{_now()}] 文件已接收，等待后台处理，来源类型：{_source_label(source_type)}"],
|
||
"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
|
||
target["phase"] = "queued"
|
||
target["phase_label"] = _phase_label("queued")
|
||
target["logs"] = [f"[{_now()}] 任务已重新加入队列，等待后台执行"]
|
||
_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
|