CIBank/backend/daily_pipeline.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import fcntl
import json
import os
import sqlite3
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path

from dotenv import load_dotenv

BASE_DIR = Path(__file__).resolve().parent
load_dotenv(BASE_DIR / ".env")
DATA_DIR = BASE_DIR / "data"
DB_PATH = Path(os.getenv("CIBANK_DB", DATA_DIR / "cibank.db"))
LOCK_PATH = DATA_DIR / "pipeline.lock"
LOG_DIR = DATA_DIR / "logs"
BACKUP_DIR = DATA_DIR / "backups"
SCHEMA_PATH = BASE_DIR / "schema.sql"


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


def open_db() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA_PATH.read_text(encoding="utf-8"))
    return conn


def acquire_lock():
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    lock_file = LOCK_PATH.open("a+")
    try:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError as exc:
        lock_file.close()
        raise RuntimeError("已有流水线任务正在运行") from exc
    lock_file.seek(0)
    lock_file.truncate()
    lock_file.write(f"{os.getpid()}\n{now_iso()}\n")
    lock_file.flush()
    return lock_file


def backup_database() -> str:
    BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    backup_path = BACKUP_DIR / f"cibank-{datetime.now().strftime('%Y%m%d-%H%M%S')}.db"
    source = sqlite3.connect(DB_PATH)
    target = sqlite3.connect(backup_path)
    try:
        source.backup(target)
    finally:
        target.close()
        source.close()
    backups = sorted(BACKUP_DIR.glob("cibank-*.db"), key=lambda path: path.stat().st_mtime, reverse=True)
    for expired in backups[14:]:
        expired.unlink()
    return str(backup_path)


def create_run(run_type: str) -> int:
    timestamp = now_iso()
    with open_db() as conn:
        cursor = conn.execute(
            """
            INSERT INTO pipeline_runs
              (run_type, started_at, status, current_step, steps_json, created_at)
            VALUES (?, ?, 'running', 'preflight', '{}', ?)
            """,
            (run_type, timestamp, timestamp),
        )
        conn.commit()
        return cursor.lastrowid


def update_run(run_id: int, **values) -> None:
    allowed = {"finished_at", "status", "current_step", "steps_json", "error_message"}
    fields = [(key, value) for key, value in values.items() if key in allowed]
    if not fields:
        return
    assignments = ", ".join(f"{key}=?" for key, _ in fields)
    with open_db() as conn:
        conn.execute(
            f"UPDATE pipeline_runs SET {assignments} WHERE id=?",
            [value for _, value in fields] + [run_id],
        )
        conn.commit()


def run_step(run_id: int, name: str, arguments: list[str], log_file, steps: dict) -> bool:
    update_run(run_id, current_step=name)
    started = time.monotonic()
    log_file.write(f"\n[{now_iso()}] START {name}: {' '.join(arguments)}\n")
    log_file.flush()
    env = os.environ.copy()
    env["PATH"] = "/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin:" + env.get("PATH", "")
    result = subprocess.run(
        [sys.executable, *arguments],
        cwd=BASE_DIR,
        env=env,
        stdout=log_file,
        stderr=subprocess.STDOUT,
        text=True,
        check=False,
    )
    duration = round(time.monotonic() - started, 2)
    steps[name] = {"status": "success" if result.returncode == 0 else "failed", "returncode": result.returncode, "duration": duration}
    update_run(run_id, steps_json=json.dumps(steps, ensure_ascii=False))
    log_file.write(f"[{now_iso()}] END {name}: code={result.returncode}, duration={duration}s\n")
    log_file.flush()
    return result.returncode == 0


def validate_environment() -> None:
    if not DB_PATH.exists():
        raise RuntimeError(f"数据库不存在：{DB_PATH}")
    api_key = os.getenv("DASHSCOPE_API_KEY", "").strip()
    if not api_key or api_key == "replace_with_your_key":
        raise RuntimeError("DASHSCOPE_API_KEY 未配置")
    with sqlite3.connect(DB_PATH) as conn:
        result = conn.execute("PRAGMA integrity_check").fetchone()[0]
    if result != "ok":
        raise RuntimeError(f"数据库完整性检查失败：{result}")


def show_status() -> int:
    with open_db() as conn:
        row = conn.execute("SELECT * FROM pipeline_runs ORDER BY id DESC LIMIT 1").fetchone()
    print(json.dumps(dict(row) if row else {"status": "never_run"}, ensure_ascii=False, indent=2))
    return 0


def run_pipeline(run_type: str, skip_crawl: bool) -> int:
    lock_file = acquire_lock()
    run_id = None
    LOG_DIR.mkdir(parents=True, exist_ok=True)
    log_path = LOG_DIR / f"pipeline-{datetime.now().strftime('%Y%m%d-%H%M%S')}.log"
    try:
        validate_environment()
        run_id = create_run(run_type)
        steps = {"backup": {"status": "success", "path": backup_database()}}
        update_run(run_id, steps_json=json.dumps(steps, ensure_ascii=False))
        commands = []
        if not skip_crawl:
            commands.extend([
                ("crawl_mp", ["crawler.py", "--source-type", "公众号"]),
                ("crawl_video", ["crawler.py", "--source-type", "视频号"]),
            ])
        analyze_args = ["structurer.py", "analyze", "--pause", "0.6", "--retries", "3"]
        embed_args = ["structurer.py", "embed"]
        if run_type == "full":
            analyze_args.append("--force")
            embed_args.append("--force")
        commands.extend([
            ("analyze", analyze_args),
            ("embed", embed_args),
            ("signals", ["structurer.py", "signals", "--days", "30"]),
        ])
        all_success = True
        with log_path.open("a", encoding="utf-8") as log_file:
            log_file.write(f"[{now_iso()}] pipeline run_id={run_id}, type={run_type}\n")
            for name, arguments in commands:
                if not run_step(run_id, name, arguments, log_file, steps):
                    all_success = False
        status = "success" if all_success else "partial"
        update_run(run_id, finished_at=now_iso(), status=status, current_step="completed", steps_json=json.dumps(steps, ensure_ascii=False))
        print(json.dumps({"run_id": run_id, "status": status, "log": str(log_path), "steps": steps}, ensure_ascii=False, indent=2))
        return 0 if all_success else 1
    except Exception as exc:
        if run_id is not None:
            update_run(run_id, finished_at=now_iso(), status="failed", current_step="failed", error_message=str(exc))
        print(json.dumps({"status": "failed", "error": str(exc), "log": str(log_path)}, ensure_ascii=False), file=sys.stderr)
        return 1
    finally:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
        lock_file.close()


def main() -> int:
    parser = argparse.ArgumentParser(description="CIBank 自动增量处理流水线")
    parser.add_argument("command", choices=["daily", "full", "status"])
    parser.add_argument("--skip-crawl", action="store_true")
    args = parser.parse_args()
    if args.command == "status":
        return show_status()
    return run_pipeline(args.command, args.skip_crawl)


if __name__ == "__main__":
    raise SystemExit(main())