universal-crawler/db.py

# -*- coding: utf-8 -*-
"""
MySQL 记录层: 任务 / 运行 / 爬取结果 写入数据库
- 网页完整内容不入库 (只存磁盘文件), 库中只存标题、网址、状态、文件路径等元数据
- 爬取成功/失败均有 status 标记 (OK / FAIL), 运行记录有 run 状态标记
- 所有操作容错: 数据库不可用时不影响爬取主流程 (仅打印日志)
"""
import copy
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor

import pymysql

DB_CONFIG = dict(
    host="121.40.164.32",
    port=16006,
    user="uni_crawler",
    password="wleD6x2T",
    charset="utf8mb4",
)

DB_NAME = "uni_crawler"

DDL = [
    """
    CREATE TABLE IF NOT EXISTS crawl_tasks (
      id VARCHAR(32) PRIMARY KEY,
      name VARCHAR(200) NOT NULL,
      mode VARCHAR(20) NOT NULL,
      config TEXT,
      urls TEXT,
      auto_config TEXT,
      schedule_config TEXT,
      created_at DATETIME,
      updated_at DATETIME,
      deleted_at DATETIME NULL
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
    """
    CREATE TABLE IF NOT EXISTS crawl_runs (
      id VARCHAR(32) PRIMARY KEY,
      task_id VARCHAR(32) NOT NULL,
      mode VARCHAR(20) NOT NULL,
      status VARCHAR(20) NOT NULL DEFAULT 'running',
      total INT DEFAULT 0,
      done INT DEFAULT 0,
      ok_count INT DEFAULT 0,
      fail_count INT DEFAULT 0,
      image_count INT DEFAULT 0,
      started_at DATETIME NULL,
      finished_at DATETIME NULL,
      out_dir VARCHAR(500),
      KEY idx_task (task_id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
    """
    CREATE TABLE IF NOT EXISTS crawl_results (
      id BIGINT AUTO_INCREMENT PRIMARY KEY,
      run_id VARCHAR(32) NOT NULL,
      task_id VARCHAR(32) NOT NULL,
      mode VARCHAR(20),
      url VARCHAR(2000) NOT NULL,
      title VARCHAR(500),
      status VARCHAR(10) NOT NULL,
      error TEXT,
      crawl_time DATETIME,
      source_url VARCHAR(2000),
      depth INT,
      attempts INT DEFAULT 1,
      base_file VARCHAR(500),
      image_count INT DEFAULT 0,
      image_files TEXT,
      UNIQUE KEY uk_run_url (run_id, url(500))
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
    """,
]


def _migrate(conn):
    """存量表结构迁移: 合并 html_file/txt_file/meta_file 为 base_file"""
    with conn.cursor() as cur:
        cur.execute("SHOW COLUMNS FROM crawl_results LIKE 'html_file'")
        if cur.fetchone():
            cur.execute("ALTER TABLE crawl_results ADD COLUMN base_file VARCHAR(500) NULL AFTER meta_file")
            cur.execute("UPDATE crawl_results SET base_file = REPLACE(html_file, '.html', '') WHERE base_file IS NULL")
            cur.execute("ALTER TABLE crawl_results DROP COLUMN html_file, DROP COLUMN txt_file, DROP COLUMN meta_file")
            print("[db] 表结构迁移完成: html_file/txt_file/meta_file -> base_file")


def _conn():
    cfg = dict(DB_CONFIG)
    cfg["database"] = DB_NAME
    return pymysql.connect(**cfg, autocommit=True, connect_timeout=5)


def _safe(fn, *args, **kwargs):
    """执行数据库操作, 失败仅打印日志不抛出"""
    try:
        fn(*args, **kwargs)
    except Exception as e:
        print(f"[db] 操作失败: {e}", flush=True)


# ---------------- 异步同步 (避免远程 MySQL 阻塞爬虫/API 主流程) ----------------

_db_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="db-sync")


def upsert_task_async(task):
    """异步写入/更新任务 (不阻塞调用方)"""
    try:
        snap = copy.deepcopy(task)
        _db_executor.submit(upsert_task, snap)
    except Exception as e:
        print(f"[db] 异步任务同步失败: {e}", flush=True)


def sync_run_async(run):
    """异步同步运行记录+爬取结果; 完成后把 _db_count 回写到主线程 run 对象"""
    try:
        snap = copy.deepcopy(run)

        def _work():
            upsert_run(snap)
            results = snap.get("results", [])
            synced = snap.get("_db_count", 0)
            if len(results) > synced:
                insert_results(snap, results[synced:])
                snap["_db_count"] = len(results)
                run["_db_count"] = snap["_db_count"]

        _db_executor.submit(_work)
    except Exception as e:
        print(f"[db] 异步运行同步失败: {e}", flush=True)


def init_db():
    """建库建表 (幂等)"""
    try:
        conn = pymysql.connect(**DB_CONFIG, autocommit=True, connect_timeout=5)
        with conn.cursor() as cur:
            cur.execute(f"CREATE DATABASE IF NOT EXISTS `{DB_NAME}` DEFAULT CHARACTER SET utf8mb4")
        conn.close()
        conn = _conn()
        _migrate(conn)
        with conn.cursor() as cur:
            for ddl in DDL:
                cur.execute(ddl)
        conn.close()
        print("[db] 数据库初始化完成 (crawler: crawl_tasks / crawl_runs / crawl_results)")
    except Exception as e:
        print(f"[db] 数据库初始化失败: {e}", flush=True)


# ---------------- 序列化工具 ----------------

def _dt(v):
    """datetime 字符串 -> MySQL DATETIME (无效返回 None)"""
    if not v:
        return None
    s = str(v).replace("T", " ")
    if len(s) >= 19:
        return s[:19]
    return None


def _j(v):
    return json.dumps(v, ensure_ascii=False) if v else None


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

def upsert_task(task):
    """任务写入/更新 (含回收站状态 deleted_at)"""
    def _do():
        conn = _conn()
        with conn.cursor() as cur:
            cur.execute(
                """INSERT INTO crawl_tasks
                   (id, name, mode, config, urls, auto_config, schedule_config,
                    created_at, updated_at, deleted_at)
                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                   ON DUPLICATE KEY UPDATE
                   name=%s, mode=%s, config=%s, urls=%s, auto_config=%s,
                   schedule_config=%s, updated_at=%s, deleted_at=%s""",
                (
                    task["id"], task.get("name", ""), task.get("mode", ""),
                    _j(task.get("config")), _j(task.get("urls")),
                    _j(task.get("auto")), _j(task.get("schedule")),
                    _dt(task.get("created_at")), _dt(task.get("updated_at")),
                    _dt(task.get("deleted_at")),
                    task.get("name", ""), task.get("mode", ""),
                    _j(task.get("config")), _j(task.get("urls")),
                    _j(task.get("auto")), _j(task.get("schedule")),
                    _dt(task.get("updated_at")), _dt(task.get("deleted_at")),
                ),
            )
        conn.close()
    _safe(_do)


def purge_task_db(task_id):
    """彻底删除任务记录"""
    def _do():
        conn = _conn()
        with conn.cursor() as cur:
            cur.execute("DELETE FROM crawl_results WHERE task_id=%s", (task_id,))
            cur.execute("DELETE FROM crawl_runs WHERE task_id=%s", (task_id,))
            cur.execute("DELETE FROM crawl_tasks WHERE id=%s", (task_id,))
        conn.close()
    _safe(_do)


def purge_trash_db():
    """清空回收站 (删除所有已标记删除的任务记录)"""
    def _do():
        conn = _conn()
        with conn.cursor() as cur:
            cur.execute("SELECT id FROM crawl_tasks WHERE deleted_at IS NOT NULL")
            ids = [r[0] for r in cur.fetchall()]
            for tid in ids:
                cur.execute("DELETE FROM crawl_results WHERE task_id=%s", (tid,))
                cur.execute("DELETE FROM crawl_runs WHERE task_id=%s", (tid,))
                cur.execute("DELETE FROM crawl_tasks WHERE id=%s", (tid,))
        conn.close()
    _safe(_do)


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

def upsert_run(run):
    """运行记录写入/更新"""
    def _do():
        conn = _conn()
        st = run.get("stats") or {}
        with conn.cursor() as cur:
            cur.execute(
                """INSERT INTO crawl_runs
                   (id, task_id, mode, status, total, done,
                    ok_count, fail_count, image_count,
                    started_at, finished_at, out_dir)
                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                   ON DUPLICATE KEY UPDATE
                   status=%s, total=%s, done=%s, ok_count=%s, fail_count=%s,
                   image_count=%s, finished_at=%s, out_dir=%s""",
                (
                    run["id"], run.get("task_id", ""), run.get("mode", ""),
                    run.get("status", ""), run.get("progress", {}).get("total", 0),
                    run.get("progress", {}).get("done", 0),
                    st.get("ok", 0), st.get("fail", 0), st.get("images", 0),
                    _dt(run.get("started_at")), _dt(run.get("finished_at")),
                    run.get("out_dir", ""),
                    run.get("status", ""), run.get("progress", {}).get("total", 0),
                    run.get("progress", {}).get("done", 0),
                    st.get("ok", 0), st.get("fail", 0), st.get("images", 0),
                    _dt(run.get("finished_at")), run.get("out_dir", ""),
                ),
            )
        conn.close()
    _safe(_do)


# ---------------- 爬取结果 (每页一条, 成功失败均记录) ----------------

def _base_of(entry):
    """从结果条目提取基础文件名 (html/txt/meta 三个后缀共用同一前缀)"""
    for f in (entry.get("meta_file"), entry.get("html_file"), entry.get("txt_file")):
        f = f or ""
        if f.endswith(".meta.json"):
            return f[:-10]
        if f.endswith(".html"):
            return f[:-5]
        if f.endswith(".txt"):
            return f[:-4]
    return ""


def insert_results(run, results):
    """批量插入爬取结果 (增量); 文件只记基础名 base_file"""
    if not results:
        return

    def _do():
        conn = _conn()
        rows = []
        for r in results:
            rows.append((
                run["id"], run.get("task_id", ""), run.get("mode", ""),
                (r.get("url") or "")[:2000], (r.get("title") or "")[:500],
                r.get("status", "FAIL"), r.get("error"),
                _dt(r.get("crawl_time")), (r.get("source_url") or "")[:2000],
                r.get("depth"), r.get("attempts", 1),
                _base_of(r),
                len(r.get("images", []) or []),
                _j([im.get("file") for im in (r.get("images") or [])]),
            ))
        with conn.cursor() as cur:
            cur.executemany(
                """INSERT IGNORE INTO crawl_results
                   (run_id, task_id, mode, url, title, status, error,
                    crawl_time, source_url, depth, attempts,
                    base_file, image_count, image_files)
                   VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                rows,
            )
        conn.close()
    _safe(_do)


def sync_run(run, persist):
    """持久化回调: 同步运行记录 + 增量同步爬取结果
    persist: callable(task_id, run) 用于回写已同步进度标记
    """
    upsert_run(run)
    results = run.get("results", [])
    synced = run.get("_db_count", 0)
    if len(results) > synced:
        insert_results(run, results[synced:])
        run["_db_count"] = len(results)
        persist(run.get("task_id"), run)


# ---------------- 历史数据回填 (幂等) ----------------

def sync_all_history(task_ids=None):
    """把本地 JSON 中的历史任务/运行/结果全量回填数据库
    可重复执行 (INSERT IGNORE + 唯一键去重)
    task_ids: 指定只回填的任务ID列表, 默认全部
    返回统计 dict
    """
    import store as _store
    stats = {"tasks": 0, "runs": 0, "results": 0}
    for task in _store.load_tasks():
        if task_ids and task["id"] not in task_ids:
            continue
        upsert_task(task)
        stats["tasks"] += 1
        for run in _store.get_runs(task["id"]):
            upsert_run(run)
            stats["runs"] += 1
            results = run.get("results", [])
            if results:
                insert_results(run, results)
                stats["results"] += len(results)
                run["_db_count"] = len(results)
                _store.save_run(task["id"], run)   # 记录已同步标记, 避免运行中重复插入
    return stats