#!/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())