universal-crawler/store.py

# -*- coding: utf-8 -*-
"""任务 / 运行记录持久化层 (JSON 文件存储, 线程安全)"""
import json
import os
import threading
import uuid
from datetime import datetime

HERE = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(HERE, "data")
TASKS_FILE = os.path.join(DATA_DIR, "tasks.json")
RUNS_FILE = os.path.join(DATA_DIR, "runs.json")   # {task_id: [run, ...]}
AUTO_STATE_DIR = os.path.join(DATA_DIR, "auto_state")  # auto 任务待爬/已爬队列独立存储

_lock = threading.RLock()


def new_id(prefix):
    return f"{prefix}_{uuid.uuid4().hex[:12]}"


def now_str():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")


def _load(path, default):
    if not os.path.exists(path):
        return default
    try:
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return default


def _save(path, data):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    os.replace(tmp, path)


# ---------------- 任务 ----------------

def load_tasks():
    with _lock:
        return _load(TASKS_FILE, [])


def get_task(task_id):
    for t in load_tasks():
        if t["id"] == task_id:
            return t
    return None


def upsert_task(task):
    with _lock:
        tasks = _load(TASKS_FILE, [])
        for i, t in enumerate(tasks):
            if t["id"] == task["id"]:
                tasks[i] = task
                break
        else:
            tasks.append(task)
        _save(TASKS_FILE, tasks)


def delete_task(task_id):
    """彻底删除任务 (任务 + 运行记录)"""
    with _lock:
        tasks = _load(TASKS_FILE, [])
        tasks = [t for t in tasks if t["id"] != task_id]
        _save(TASKS_FILE, tasks)
        runs = _load(RUNS_FILE, {})
        runs.pop(task_id, None)
        _save(RUNS_FILE, runs)


# ---------------- 回收站 ----------------

def soft_delete_task(task_id, ts):
    """删除任务 -> 移入回收站 (软删除, 记录保留可恢复)"""
    with _lock:
        tasks = _load(TASKS_FILE, [])
        for t in tasks:
            if t["id"] == task_id:
                t["deleted_at"] = ts
                break
        _save(TASKS_FILE, tasks)


def restore_task(task_id):
    """从回收站恢复任务"""
    with _lock:
        tasks = _load(TASKS_FILE, [])
        for t in tasks:
            if t["id"] == task_id:
                t.pop("deleted_at", None)
                break
        _save(TASKS_FILE, tasks)


def list_trash():
    """回收站任务列表 (按删除时间倒序)"""
    with _lock:
        tasks = _load(TASKS_FILE, [])
        return sorted(
            [t for t in tasks if t.get("deleted_at")],
            key=lambda x: x.get("deleted_at", ""), reverse=True,
        )


def purge_task(task_id):
    """从回收站彻底删除单个任务 (任务 + 运行记录)"""
    with _lock:
        tasks = _load(TASKS_FILE, [])
        tasks = [t for t in tasks if t["id"] != task_id]
        _save(TASKS_FILE, tasks)
        runs = _load(RUNS_FILE, {})
        runs.pop(task_id, None)
        _save(RUNS_FILE, runs)


def purge_trash():
    """清空回收站, 返回被清空的任务 id 列表"""
    with _lock:
        tasks = _load(TASKS_FILE, [])
        kept, purged = [], []
        for t in tasks:
            if t.get("deleted_at"):
                purged.append(t["id"])
            else:
                kept.append(t)
        _save(TASKS_FILE, kept)
        runs = _load(RUNS_FILE, {})
        for tid in purged:
            runs.pop(tid, None)
        _save(RUNS_FILE, runs)
        return purged


# ---------------- auto 任务状态 (待爬队列/已爬集合) ----------------
# 独立于 tasks.json 存储: 避免上万条 URL 撑大任务文件导致全量读写变慢

def auto_state_path(task_id):
    return os.path.join(AUTO_STATE_DIR, f"{task_id}.json")


def load_auto_state(task_id, fallback_task=None):
    """读取 auto 任务的 pending/visited 状态; 无独立文件时从旧任务对象迁移"""
    with _lock:
        path = auto_state_path(task_id)
        if os.path.exists(path):
            st = _load(path, {}) or {}
            return {"pending": st.get("pending", []), "visited": st.get("visited", [])}
        if fallback_task:
            auto = fallback_task.get("auto") or {}
            st = {"pending": auto.get("pending", []), "visited": auto.get("visited", [])}
            if st["pending"] or st["visited"]:
                _save(path, st)   # 自动迁移旧数据
                return st
    return {"pending": [], "visited": []}


def save_auto_state(task_id, state):
    with _lock:
        _save(auto_state_path(task_id), {
            "pending": state.get("pending", []),
            "visited": state.get("visited", []),
        })


def clear_auto_state(task_id):
    """删除 auto 状态文件"""
    with _lock:
        try:
            os.remove(auto_state_path(task_id))
        except OSError:
            pass


# ---------------- 运行记录 ----------------

def load_runs_map():
    with _lock:
        return _load(RUNS_FILE, {})


def get_runs(task_id):
    with _lock:
        return _load(RUNS_FILE, {}).get(task_id, [])


def add_run(task_id, run):
    with _lock:
        runs = _load(RUNS_FILE, {})
        runs.setdefault(task_id, []).append(run)
        if len(runs[task_id]) > 50:          # 每个任务最多保留 50 次运行记录
            del runs[task_id][:-50]
        _save(RUNS_FILE, runs)


def save_run(task_id, run):
    with _lock:
        runs = _load(RUNS_FILE, {})
        lst = runs.setdefault(task_id, [])
        for i, r in enumerate(lst):
            if r["id"] == run["id"]:
                lst[i] = run
                break
        else:
            lst.append(run)
        if len(lst) > 50:
            del lst[:-50]
        _save(RUNS_FILE, runs)


def get_run(run_id):
    with _lock:
        runs = _load(RUNS_FILE, {})
        for lst in runs.values():
            for r in lst:
                if r["id"] == run_id:
                    return r
    return None