AiDOPWarehouse/doc/plan/sql/uat-data/run_uat_full_module_refresh.py

#!/usr/bin/env python3
"""Generate S2 schedules and rebuild S1-S7 for all UAT tenants."""

from __future__ import annotations

import argparse
import json
import time
from typing import Any

import run_wp_wb8_api_loop as wb8
from apply_sql_file import connect


PASSWORD = "1234567890dop"
TENANTS = (wb8.A, wb8.B, wb8.DEMO)
MODULES = ("S1", "S2", "S3", "S4", "S5", "S6", "S7")
TERMINAL = {"SUCCESS", "FAILED", "CANCELLED"}


def value(body: Any, *names: str, default: Any = None) -> Any:
    data = wb8.unwrap(body)
    return wb8.pick(data, *names, default=default)


def call(method: str, path: str, token: str, timeout: int = 300) -> Any:
    status, body = wb8.request(method, path, token=token, timeout=timeout)
    if status < 200 or status >= 300:
        raise RuntimeError(f"{method} {path}: HTTP {status} {body}")
    return body


def wait_schedule(token: str, run_id: int, timeout_seconds: int) -> dict[str, Any]:
    deadline = time.monotonic() + timeout_seconds
    while time.monotonic() < deadline:
        body = call("GET", f"/api/Production/scheduling/generate-status?runId={run_id}", token)
        status = str(value(body, "status", "Status", default="")).upper()
        if status in TERMINAL:
            return {"runId": run_id, "status": status, "message": value(body, "message", "Message")}
        time.sleep(3)
    return {"runId": run_id, "status": "TIMEOUT"}


def running_schedule_id(tenant_id: str) -> int:
    conn = connect()
    try:
        with conn.cursor() as cursor:
            cursor.execute(
                """
                SELECT id
                FROM aidop_action_run_log
                WHERE tenant_id=%s AND action_code='S2_SCHEDULE_GENERATE' AND status='RUNNING'
                ORDER BY id DESC
                LIMIT 1
                """,
                (int(tenant_id),),
            )
            row = cursor.fetchone()
            return int(row["id"]) if row else 0
    finally:
        conn.close()


def wait_rebuild(token: str, module: str, factory_id: str, job_id: int, timeout_seconds: int) -> dict[str, Any]:
    deadline = time.monotonic() + timeout_seconds
    path = f"/api/AidopKanban/{module}/rebuild-jobs/{job_id}?factoryId={factory_id}"
    while time.monotonic() < deadline:
        body = call("GET", path, token)
        status = str(value(body, "status", "Status", default="")).upper()
        if status in TERMINAL:
            return {
                "jobId": job_id,
                "status": status,
                "error": value(body, "errorMessage", "ErrorMessage"),
            }
        time.sleep(5)
    return {"jobId": job_id, "status": "TIMEOUT"}


def run(
    timeout_seconds: int,
    modules: tuple[str, ...],
    skip_schedule: bool,
    factory_id_override: str | None = None,
    tenant_codes: set[str] | None = None,
) -> dict[str, Any]:
    result: dict[str, Any] = {"base": wb8.BASE, "tenants": {}}
    sessions: list[tuple[dict[str, Any], str]] = []
    selected_tenants = tuple(
        tenant for tenant in TENANTS
        if not tenant_codes or tenant["code"].upper() in tenant_codes
    )
    for tenant in selected_tenants:
        login = wb8.login(tenant["admin"]["account"], tenant["tenant_id"], PASSWORD)
        sessions.append((tenant, login["token"]))

    if not skip_schedule:
        # Production schedules are globally serialized by the backend gate.
        for tenant, token in sessions:
            code = tenant["code"]
            run_id = running_schedule_id(tenant["tenant_id"])
            if not run_id:
                body = call(
                    "POST",
                    f"/api/Production/scheduling/generate?domain={tenant['tenant_id']}&enableCapacityConstraint=false",
                    token,
                )
                run_id = int(value(body, "runId", "RunId", "id", "Id", default=0) or 0)
            else:
                body = {"result": {"runId": run_id}}
            result["tenants"].setdefault(code, {})["schedule"] = (
                wait_schedule(token, run_id, timeout_seconds) if run_id else {"status": "NO_RUN_ID", "body": wb8.unwrap(body)}
            )

    jobs: list[tuple[dict[str, Any], str, str, int, str]] = []
    for tenant, token in sessions:
        code = tenant["code"]
        factory_id = factory_id_override or tenant["factory_id"]
        result["tenants"].setdefault(code, {})["rebuilds"] = {}
        for module in modules:
            status_code, body = wb8.request(
                "POST",
                f"/api/AidopKanban/{module}/rebuild-jobs?factoryId={factory_id}",
                token=token,
            )
            if status_code not in (200, 201, 202, 409):
                raise RuntimeError(f"POST {module} rebuild: HTTP {status_code} {body}")
            job_id = int(value(body, "jobId", "JobId", "id", "Id", default=0) or 0)
            result["tenants"][code]["rebuilds"][module] = {"jobId": job_id, "status": "QUEUED"}
            if job_id:
                jobs.append((tenant, token, module, job_id, factory_id))

    deadline = time.monotonic() + timeout_seconds
    pending = list(jobs)
    while pending and time.monotonic() < deadline:
        next_pending: list[tuple[dict[str, Any], str, str, int, str]] = []
        for tenant, token, module, job_id, factory_id in pending:
            path = f"/api/AidopKanban/{module}/rebuild-jobs/{job_id}?factoryId={factory_id}"
            body = call("GET", path, token)
            status = str(value(body, "status", "Status", default="")).upper()
            if status in TERMINAL:
                result["tenants"][tenant["code"]]["rebuilds"][module] = {
                    "jobId": job_id,
                    "status": status,
                    "error": value(body, "errorMessage", "ErrorMessage"),
                }
            else:
                next_pending.append((tenant, token, module, job_id, factory_id))
        pending = next_pending
        if pending:
            time.sleep(5)
    for tenant, _, module, job_id, _ in pending:
        result["tenants"][tenant["code"]]["rebuilds"][module] = {"jobId": job_id, "status": "TIMEOUT"}
    return result


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", default="http://127.0.0.1:5005")
    parser.add_argument("--timeout-seconds", type=int, default=1800)
    parser.add_argument("--modules", default=",".join(MODULES))
    parser.add_argument("--skip-schedule", action="store_true")
    parser.add_argument("--factory-id", help="override each tenant's KPI factory id")
    parser.add_argument("--tenants", default="A,B,DEMO", help="comma-separated tenant codes")
    args = parser.parse_args()
    wb8.BASE = args.base.rstrip("/")
    modules = tuple(x.strip().upper() for x in args.modules.split(",") if x.strip())
    invalid = sorted(set(modules) - set(MODULES))
    if invalid:
        parser.error(f"unsupported modules: {','.join(invalid)}")
    tenant_codes = {x.strip().upper() for x in args.tenants.split(",") if x.strip()}
    valid_tenant_codes = {tenant["code"].upper() for tenant in TENANTS}
    invalid_tenants = sorted(tenant_codes - valid_tenant_codes)
    if invalid_tenants:
        parser.error(f"unsupported tenants: {','.join(invalid_tenants)}")
    result = run(args.timeout_seconds, modules, args.skip_schedule, args.factory_id, tenant_codes)
    print(json.dumps(result, ensure_ascii=False, indent=2, default=str))
    failed = [
        f"{tenant}:{module}:{data['status']}"
        for tenant, tenant_data in result["tenants"].items()
        for module, data in tenant_data.get("rebuilds", {}).items()
        if data.get("status") != "SUCCESS"
    ]
    return 1 if failed else 0


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