SupplyAgent/supply_infra/pipeline/preflight.py

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()