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);
            """
        )


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)