from __future__ import annotations import json import os from dataclasses import dataclass from pathlib import Path from typing import Any from alembic.config import Config from alembic.migration import MigrationContext from alembic.script import ScriptDirectory from sqlalchemy import inspect, text from supply_agent.config import Settings from supply_infra.aigc.plan_map import list_unique_plan_pairs from supply_infra.config import InfraSettings from supply_infra.db.session import dispose_engine, get_engine from supply_infra.pipeline.dag import PIPELINE_STEPS, validate_dag from supply_infra.pipeline.registry import STEP_REGISTRY REQUIRED_PIPELINE_TABLES = frozenset( { "pipeline_run", "pipeline_step_run", "pipeline_lock", "pipeline_outbox", "global_tree_category", "multi_demand_pool_di", "demand_belong_category", "demand_belong_pool_rel", "demand_popularity_stats", "category_tree_weight", "multi_demand_video_detail", "multi_demand_video_point", "demand_grade", "demand_grade_category_rel", "demand_grade_plan", "demand_grade_plan_group", "demand_grade_plan_group_item", "demand_video_expansion", "demand_video_expansion_run", "video_discovery_run", "video_discovery_search", "video_discovery_candidate", "evidence_package", "raw_demand_expression", "standard_demand_term", "standard_demand_alias", "platform_demand", "platform_demand_version", "platform_demand_category_rel", "strategy_version", "demand_evaluation_snapshot", "demand_score_contribution", "daily_demand_package", "daily_demand_task", "content_demand_coverage", "content_performance_fact", "demand_attribution_snapshot", "strategy_change_proposal", "demand_feedback", } ) @dataclass(frozen=True) class PreflightResult: passed: bool errors: list[str] warnings: list[str] checks: dict[str, Any] def to_dict(self) -> dict[str, Any]: return { "passed": self.passed, "errors": self.errors, "warnings": self.warnings, "checks": self.checks, } def _env_true(name: str, *, default: bool) -> bool: raw = os.getenv(name) if raw is None: return default return raw.strip().lower() in {"1", "true", "yes", "on"} def evaluate_preflight( *, infra: InfraSettings, openrouter_configured: bool, table_names: set[str], current_revision: str | None, expected_revision: str, require_scheduler: bool, require_secure_cookie: bool, ) -> PreflightResult: errors: list[str] = [] warnings: list[str] = [] missing_tables = sorted(REQUIRED_PIPELINE_TABLES - table_names) dag_keys = {step.key for step in PIPELINE_STEPS} registry_keys = set(STEP_REGISTRY) if require_scheduler and not infra.scheduler_enabled: errors.append( "SCHEDULER_ENABLED 必须为 true，否则不会自动提交每日任务" ) if not infra.mysql_password: errors.append("MYSQL_PASSWORD 未配置") if not (infra.odps_access_id and infra.odps_access_key and infra.odps_project): errors.append("ODPS_ACCESS_ID/ODPS_ACCESS_KEY/ODPS_PROJECT 未完整配置") if not openrouter_configured: errors.append("OPENROUTER_API_KEY 未配置，Agent 步骤无法运行") if not infra.category_match_api_url.strip(): errors.append("CATEGORY_MATCH_API_URL 未配置，需求词无法挂树") if not infra.aigc_api_token: errors.append("AIGC_API_TOKEN 未配置，最后发布步骤无法运行") if infra.log_oss_upload_enabled and not ( infra.aliyun_oss_access_key_id and infra.aliyun_oss_access_key_secret and infra.aliyun_oss_bucket ): errors.append("LOG_OSS_UPLOAD_ENABLED=true，但 OSS 凭证未完整配置") if require_secure_cookie and not infra.auth_cookie_secure: errors.append("HTTPS 上线要求 AUTH_COOKIE_SECURE=true") elif not infra.auth_cookie_secure: warnings.append("AUTH_COOKIE_SECURE=false；仅适合非 HTTPS 本地环境") if current_revision != expected_revision: errors.append( "数据库迁移版本不一致：" f"current={current_revision or 'none'} expected={expected_revision}" ) if missing_tables: errors.append(f"缺少流水线依赖表：{', '.join(missing_tables)}") if dag_keys != registry_keys: errors.append( "DAG 与处理器注册不一致：" f"missing={sorted(dag_keys - registry_keys)} " f"orphan={sorted(registry_keys - dag_keys)}" ) if not list_unique_plan_pairs(): errors.append("AIGC 生成/发布计划映射为空") checks = { "database_revision": current_revision, "expected_revision": expected_revision, "database_table_count": len(table_names), "missing_table_count": len(missing_tables), "pipeline_step_count": len(PIPELINE_STEPS), "scheduler_enabled": infra.scheduler_enabled, "scheduler_cron": ( f"{infra.scheduler_cron_hour:02d}:" f"{infra.scheduler_cron_minute:02d} " f"{infra.scheduler_timezone}" ), "worker_processes": infra.pipeline_worker_processes, "max_active_steps": infra.pipeline_max_active_steps, "mysql_required_connections": infra.pipeline_required_connections, "mysql_connection_budget": infra.mysql_connection_budget, "aigc_plan_pair_count": len(list_unique_plan_pairs()), "category_match_api_configured": bool( infra.category_match_api_url.strip() ), } return PreflightResult( passed=not errors, errors=errors, warnings=warnings, checks=checks, ) def run_preflight() -> PreflightResult: infra = InfraSettings() agent = Settings() validate_dag() config_path = Path(__file__).resolve().parents[2] / "alembic.ini" alembic_config = Config(str(config_path)) expected_revision = ScriptDirectory.from_config( alembic_config ).get_current_head() engine = get_engine() try: with engine.connect() as connection: connection.execute(text("SELECT 1")) table_names = set(inspect(connection).get_table_names()) current_revision = MigrationContext.configure( connection ).get_current_revision() except Exception as exc: return PreflightResult( passed=False, errors=[ "MySQL 连接或元数据检查失败：" f"{type(exc).__name__}（详细信息见数据库日志）" ], warnings=[], checks={"database_connected": False}, ) finally: dispose_engine() return evaluate_preflight( infra=infra, openrouter_configured=bool(agent.openrouter_api_key.strip()), table_names=table_names, current_revision=current_revision, expected_revision=expected_revision, require_scheduler=_env_true( "PIPELINE_REQUIRE_SCHEDULER", default=True, ), require_secure_cookie=_env_true( "PIPELINE_REQUIRE_SECURE_COOKIE", default=False, ), ) def main() -> None: result = run_preflight() print(json.dumps(result.to_dict(), ensure_ascii=False, indent=2)) raise SystemExit(0 if result.passed else 1) if __name__ == "__main__": main()