初始提交：餐智库 CIBank 餐饮行业知识库
This commit is contained in:
freedakgmail
@@ -0,0 +1,198 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user