holy-python/database.py

import json
import sqlite3
from contextlib import contextmanager
from datetime import datetime
from typing import Any, Iterator

from config import Config


def now_iso() -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


@contextmanager
def connection() -> Iterator[sqlite3.Connection]:
    conn = sqlite3.connect(Config.DATABASE_PATH, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    try:
        yield conn
        conn.commit()
    finally:
        conn.close()


def init_database() -> None:
    """仅在服务启动时创建表，接口调用过程中不会重复执行建表。"""

    Config.DATABASE_PATH.parent.mkdir(parents=True, exist_ok=True)
    Config.FILE_STORAGE_PATH.mkdir(parents=True, exist_ok=True)
    with connection() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS research_tasks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                public_task_id TEXT NOT NULL UNIQUE,
                name TEXT NOT NULL,
                platform TEXT NOT NULL,
                mode TEXT NOT NULL,
                status TEXT NOT NULL,
                execution_stage TEXT NOT NULL DEFAULT 'queued',
                progress_percent INTEGER NOT NULL DEFAULT 0 CHECK(progress_percent BETWEEN 0 AND 100),
                created_by_id TEXT NOT NULL,
                created_by_name TEXT NOT NULL,
                retry_of_task_id INTEGER,
                crawler_task_id TEXT,
                error_code TEXT,
                error_message TEXT,
                result_json TEXT,
                cancel_requested_at TEXT,
                started_at TEXT,
                completed_at TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                version INTEGER NOT NULL DEFAULT 0,
                FOREIGN KEY(retry_of_task_id) REFERENCES research_tasks(id)
            );
            CREATE INDEX IF NOT EXISTS idx_tasks_status_created
                ON research_tasks(status, created_at DESC);
            CREATE INDEX IF NOT EXISTS idx_tasks_creator_created
                ON research_tasks(created_by_id, created_at DESC);

            CREATE TABLE IF NOT EXISTS task_keywords (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                keyword_order INTEGER NOT NULL,
                keyword TEXT NOT NULL,
                normalized_keyword TEXT NOT NULL,
                created_at TEXT NOT NULL,
                UNIQUE(task_id, keyword_order),
                UNIQUE(task_id, normalized_keyword),
                FOREIGN KEY(task_id) REFERENCES research_tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS task_events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                event_type TEXT NOT NULL,
                stage TEXT,
                progress_percent INTEGER,
                actor_type TEXT NOT NULL,
                actor_id TEXT,
                payload_json TEXT,
                occurred_at TEXT NOT NULL,
                FOREIGN KEY(task_id) REFERENCES research_tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS task_files (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                public_file_id TEXT NOT NULL UNIQUE,
                task_id INTEGER NOT NULL,
                name TEXT NOT NULL,
                stage TEXT NOT NULL,
                category TEXT NOT NULL,
                format TEXT NOT NULL,
                record_count INTEGER NOT NULL DEFAULT 0,
                size_bytes INTEGER NOT NULL DEFAULT 0,
                storage_key TEXT NOT NULL,
                sha256 TEXT NOT NULL,
                schema_version TEXT NOT NULL DEFAULT '1.0',
                generated_at TEXT NOT NULL,
                available INTEGER NOT NULL DEFAULT 1,
                UNIQUE(task_id, name),
                FOREIGN KEY(task_id) REFERENCES research_tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS auth_users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ding_user_id TEXT NOT NULL UNIQUE,
                union_id TEXT,
                open_id TEXT,
                name TEXT NOT NULL,
                avatar TEXT,
                mobile TEXT,
                email TEXT,
                raw_json TEXT,
                last_login_at TEXT NOT NULL,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_auth_users_name
                ON auth_users(name);

            CREATE TABLE IF NOT EXISTS auth_sessions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_hash TEXT NOT NULL UNIQUE,
                user_id INTEGER NOT NULL,
                expires_at TEXT NOT NULL,
                created_at TEXT NOT NULL,
                last_seen_at TEXT NOT NULL,
                revoked_at TEXT,
                FOREIGN KEY(user_id) REFERENCES auth_users(id) ON DELETE CASCADE
            );
            CREATE INDEX IF NOT EXISTS idx_auth_sessions_user
                ON auth_sessions(user_id, expires_at);

            CREATE TABLE IF NOT EXISTS acceptance_tasks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                public_task_id TEXT NOT NULL UNIQUE,
                name TEXT NOT NULL,
                acceptance_method TEXT NOT NULL,
                status TEXT NOT NULL DEFAULT 'waiting',
                progress_percent INTEGER NOT NULL DEFAULT 0,
                device_count INTEGER NOT NULL DEFAULT 1,
                created_by_id TEXT NOT NULL,
                created_by_name TEXT NOT NULL,
                error_message TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_acceptance_tasks_created
                ON acceptance_tasks(created_at DESC);

            CREATE TABLE IF NOT EXISTS acceptance_keywords (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                keyword_order INTEGER NOT NULL,
                keyword TEXT NOT NULL,
                created_at TEXT NOT NULL,
                UNIQUE(task_id, keyword_order),
                FOREIGN KEY(task_id) REFERENCES acceptance_tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS acceptance_feishu_sources (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                source_order INTEGER NOT NULL,
                source_url TEXT NOT NULL,
                final_url TEXT,
                source_type TEXT,
                title TEXT,
                content_text TEXT,
                content_json TEXT,
                fetch_status TEXT NOT NULL DEFAULT 'waiting',
                http_status INTEGER,
                error_message TEXT,
                fetched_at TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                UNIQUE(task_id, source_order),
                UNIQUE(task_id, source_url),
                FOREIGN KEY(task_id) REFERENCES acceptance_tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS acceptance_feishu_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                source_id INTEGER NOT NULL,
                row_number INTEGER NOT NULL,
                row_json TEXT NOT NULL,
                created_at TEXT NOT NULL,
                UNIQUE(source_id, row_number),
                FOREIGN KEY(source_id) REFERENCES acceptance_feishu_sources(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS acceptance_xhs_notes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                source_id INTEGER NOT NULL,
                source_row_number INTEGER,
                source_column_name TEXT,
                note_url TEXT NOT NULL,
                normalized_url TEXT NOT NULL,
                note_id TEXT,
                parse_status TEXT NOT NULL DEFAULT 'waiting',
                parsed_json TEXT,
                error_message TEXT,
                parsed_at TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                UNIQUE(task_id, normalized_url),
                FOREIGN KEY(task_id) REFERENCES acceptance_tasks(id) ON DELETE CASCADE,
                FOREIGN KEY(source_id) REFERENCES acceptance_feishu_sources(id) ON DELETE CASCADE
            );
            CREATE INDEX IF NOT EXISTS idx_acceptance_notes_task_status
                ON acceptance_xhs_notes(task_id, parse_status);

            CREATE TABLE IF NOT EXISTS feishu_notes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                note_url TEXT NOT NULL,
                normalized_url TEXT NOT NULL UNIQUE,
                note_id TEXT UNIQUE,
                title TEXT,
                body TEXT,
                author_nickname TEXT,
                author_avatar TEXT,
                likes INTEGER,
                collects INTEGER,
                comments INTEGER,
                images_json TEXT,
                tags_json TEXT,
                raw_json TEXT NOT NULL,
                parsed_at TEXT NOT NULL,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_feishu_notes_note_id
                ON feishu_notes(note_id);

            CREATE TABLE IF NOT EXISTS acceptance_keyword_runs (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                keyword TEXT NOT NULL,
                run_type TEXT NOT NULL,
                run_order INTEGER NOT NULL,
                status TEXT NOT NULL DEFAULT 'waiting',
                crawler_task_id TEXT,
                result_json TEXT,
                error_message TEXT,
                started_at TEXT,
                completed_at TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                UNIQUE(task_id, run_type, keyword),
                FOREIGN KEY(task_id) REFERENCES acceptance_tasks(id) ON DELETE CASCADE
            );
            CREATE INDEX IF NOT EXISTS idx_acceptance_keyword_runs_task
                ON acceptance_keyword_runs(task_id, status);

            CREATE TABLE IF NOT EXISTS acceptance_cookie_attempts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                keyword_run_id INTEGER NOT NULL,
                keyword TEXT NOT NULL,
                credential_id TEXT NOT NULL,
                attempt_number INTEGER NOT NULL,
                status TEXT NOT NULL,
                crawler_task_id TEXT,
                result_json TEXT,
                error_message TEXT,
                started_at TEXT NOT NULL,
                completed_at TEXT NOT NULL,
                created_at TEXT NOT NULL,
                UNIQUE(keyword_run_id, credential_id, attempt_number),
                FOREIGN KEY(task_id) REFERENCES acceptance_tasks(id) ON DELETE CASCADE,
                FOREIGN KEY(keyword_run_id) REFERENCES acceptance_keyword_runs(id) ON DELETE CASCADE
            );
            CREATE INDEX IF NOT EXISTS idx_cookie_attempts_run_status
                ON acceptance_cookie_attempts(keyword_run_id, status);

            CREATE TABLE IF NOT EXISTS acceptance_task_resources (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                resource_type TEXT NOT NULL,
                resource_id TEXT NOT NULL,
                resource_order INTEGER NOT NULL,
                created_at TEXT NOT NULL,
                UNIQUE(task_id, resource_type, resource_id),
                FOREIGN KEY(task_id) REFERENCES acceptance_tasks(id) ON DELETE CASCADE
            );
            CREATE INDEX IF NOT EXISTS idx_acceptance_task_resources
                ON acceptance_task_resources(task_id, resource_order);
            """
        )
        # 将旧版本保存在 acceptance_xhs_notes.parsed_json 中的成功结果迁移到
        # 全局去重的笔记业务表；任务关联仍由 acceptance_xhs_notes 保留。
        legacy_rows = conn.execute(
            """SELECT n.* FROM acceptance_xhs_notes n
               LEFT JOIN feishu_notes f ON f.normalized_url=n.normalized_url
               WHERE n.parse_status='completed' AND n.parsed_json IS NOT NULL AND f.id IS NULL"""
        ).fetchall()
        for row in legacy_rows:
            try:
                payload = json.loads(row["parsed_json"])
            except (TypeError, json.JSONDecodeError):
                continue
            upsert_feishu_note(conn, row, payload)
        # 为升级前已有的 Cookie 尝试补建任务资源记录，详情页可立即展示设备进度。
        legacy_resources = conn.execute(
            """SELECT task_id,credential_id,MIN(id) AS first_id
               FROM acceptance_cookie_attempts
               GROUP BY task_id,credential_id ORDER BY task_id,first_id"""
        ).fetchall()
        resource_orders: dict[int, int] = {}
        for row in legacy_resources:
            order = resource_orders.get(row["task_id"], 0)
            conn.execute(
                """INSERT OR IGNORE INTO acceptance_task_resources
                   (task_id,resource_type,resource_id,resource_order,created_at)
                   VALUES (?,'cookie',?,?,?)""",
                (row["task_id"], row["credential_id"], order, now_iso()),
            )
            resource_orders[row["task_id"]] = order + 1


def _integer_or_none(value: Any) -> int | None:
    if value is None or value == "":
        return None
    try:
        return int(str(value).replace(",", ""))
    except (TypeError, ValueError):
        return None


def upsert_feishu_note(conn: sqlite3.Connection, note_row: sqlite3.Row, payload: dict[str, Any]) -> None:
    """按规范化链接更新或新增笔记；任务关联保存在 acceptance_xhs_notes。"""

    now = now_iso()
    note_id = payload.get("noteId") or payload.get("note_id") or note_row["note_id"]
    existing = conn.execute(
        """SELECT id FROM feishu_notes
           WHERE normalized_url=? OR (? IS NOT NULL AND note_id=?)
           ORDER BY CASE WHEN note_id=? THEN 0 ELSE 1 END LIMIT 1""",
        (note_row["normalized_url"], note_id, note_id, note_id),
    ).fetchone()
    values = (
        note_row["note_url"], note_row["normalized_url"], note_id,
        payload.get("title"), payload.get("body") or payload.get("content"),
        payload.get("nickname") or payload.get("author_nickname"),
        payload.get("avatar") or payload.get("author_avatar"),
        _integer_or_none(payload.get("likes")), _integer_or_none(payload.get("collects")),
        _integer_or_none(payload.get("comments")),
        json.dumps(payload.get("images") or payload.get("images_list") or [], ensure_ascii=False),
        json.dumps(payload.get("tags") or [], ensure_ascii=False),
        json.dumps(payload, ensure_ascii=False), now, now,
    )
    if existing:
        conn.execute(
            """UPDATE feishu_notes SET
               note_url=?,normalized_url=?,note_id=?,title=?,body=?,author_nickname=?,author_avatar=?,
               likes=?,collects=?,comments=?,images_json=?,tags_json=?,raw_json=?,parsed_at=?,updated_at=?
               WHERE id=?""",
            (*values, existing["id"]),
        )
    else:
        conn.execute(
            """INSERT INTO feishu_notes
           (note_url,normalized_url,note_id,title,body,author_nickname,author_avatar,
            likes,collects,comments,images_json,tags_json,raw_json,parsed_at,created_at,updated_at)
           VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
            """,
            (*values, now),
        )


def json_dumps(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def json_loads(value: str | None) -> Any:
    if not value:
        return None
    return json.loads(value)