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