chuanqi/cq2_sdk/scripts/demo_task_log_store.py

"""End-to-end demo: TaskLogStore against real MySQL.

Simulates the 13-task driver writing per-task results to the task_logs
table, then demonstrates the resume / query patterns:

  1. setup: create user + server + zone + run via the other Stores
  2. append 5 task logs (t01..t05) with mixed success/failure
  3. list_for_run: read all logs back in order
  4. list_for_run_task: get the latest log for a specific task
  5. count_for_run: total / success / failure counts
  6. completed_task_ids: the set the resume logic would use

Run:
  uv run python scripts/demo_task_log_store.py
  uv run python scripts/demo_task_log_store.py --cleanup
"""
from __future__ import annotations
import argparse
import sys
from datetime import datetime
from pathlib import Path

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE.parent / "src"))

from cq2_sdk.env_loader import load_env  # noqa: E402
_PROJECT_ROOT = _HERE.parent
load_env(_PROJECT_ROOT / ".env")

from cq2_sdk.db_models import Zone as DbZone  # noqa: E402
from cq2_sdk.store import (  # noqa: E402
    RunStore,
    ServerStore,
    TaskLogStore,
    UserStore,
    ZoneStore,
    connection,
)

TEST_USERNAME = "demo_task_log_user"
TEST_SERVER = "DEMO_TASKLOG_SERVER_玖扬复古176"
TEST_ZONE_ID = "demo_tasklog_zone_001"


def setup_run(servers, zones, users, runs) -> tuple[int, int, int]:
    """Create user + server + zone + run. Returns (user_id, server_id, run_id)."""
    srv = servers.upsert(TEST_SERVER, 99999999, game_id=206923)
    zones.bulk_upsert(srv.id, [
        DbZone(zone_id=TEST_ZONE_ID, zone_name="TaskLog Demo",
               host="demo.host", port=12345,
               open_time=datetime.now().replace(microsecond=0),
               status="open"),
    ])
    user = users.upsert_by_username(TEST_USERNAME, "demo_token_for_tasklog")
    run = runs.create(user_id=user.id, server_id=srv.id, zone_id=1)
    return user.id, srv.id, run.id


def cleanup_run(cur, run_id: int, user_id: int) -> None:
    cur.execute("DELETE FROM task_logs WHERE run_id = %s", (run_id,))
    cur.execute("DELETE FROM runs WHERE id = %s", (run_id,))
    cur.execute("DELETE FROM users WHERE id = %s", (user_id,))


def main() -> int:
    parser = argparse.ArgumentParser(description="Demo TaskLogStore")
    parser.add_argument("--cleanup", action="store_true",
                        help="Delete the test user/run/logs after demo")
    args = parser.parse_args()

    print("=" * 70)
    print("TaskLogStore end-to-end demo")
    print("=" * 70)

    with connection() as conn:
        servers = ServerStore(conn)
        zones = ZoneStore(conn)
        users = UserStore(conn)
        runs = RunStore(conn)
        logs = TaskLogStore(conn)

        # ---------- setup ----------
        print("\n[setup] create user + server + zone + run...")
        user_id, server_id, run_id = setup_run(servers, zones, users, runs)
        print(f"        user_id={user_id}  server_id={server_id}  run_id={run_id}")

        # ---------- 1. append 5 task logs ----------
        print("\n[1] TaskLogStore.append() — simulate 5 task results")
        scenarios = [
            ("t01", True,  2.1,  None),
            ("t02", True,  5.4,  None),
            ("t03", False, 1.0,  "geetest validate token expired"),
            ("t04", True,  0.8,  None),
            ("t05", True,  3.2,  None),
        ]
        appended_ids: list[int] = []
        for tid, ok, elapsed, err in scenarios:
            log = logs.append(run_id=run_id, task_id=tid, success=ok,
                              elapsed_seconds=elapsed, error=err)
            appended_ids.append(log.id)
            status = "OK " if ok else "ERR"
            err_part = f" err={err!r}" if err else ""
            print(f"        append({tid}, {status}) -> id={log.id}  "
                  f"elapsed={elapsed}s{err_part}")

        # ---------- 2. list_for_run ----------
        print("\n[2] TaskLogStore.list_for_run(run_id) — read back in order")
        all_logs = logs.list_for_run(run_id)
        for l in all_logs:
            mark = "[OK]" if l.success else "[ERR]"
            err_part = f" err={l.error!r}" if l.error else ""
            print(f"        {mark} {l.task_id}  started={l.started_at}  "
                  f"elapsed={l.elapsed_seconds}s{err_part}")

        # ---------- 3. list_for_run_task ----------
        print("\n[3] TaskLogStore.list_for_run_task(run_id, 't03') — "
              "find the failure")
        fail_log = logs.list_for_run_task(run_id, "t03")
        assert fail_log is not None
        assert fail_log.success is False
        assert fail_log.error == "geetest validate token expired"
        print(f"        t03.success = {fail_log.success}  error = {fail_log.error!r}")

        # ---------- 4. counts ----------
        print("\n[4] TaskLogStore.count_for_run(run_id) — summary")
        total = logs.count_for_run(run_id)
        ok = logs.count_for_run(run_id, success=True)
        fail = logs.count_for_run(run_id, success=False)
        print(f"        total  = {total}")
        print(f"        success= {ok}")
        print(f"        fail   = {fail}")
        assert total == 5 and ok == 4 and fail == 1

        # ---------- 5. completed_task_ids ----------
        print("\n[5] TaskLogStore.completed_task_ids(run_id) — resume set")
        done = logs.completed_task_ids(run_id)
        print(f"        completed = {sorted(done)}")
        # E2 resume logic would skip these and only run t06..t13
        assert done == {"t01", "t02", "t04", "t05"}

        # ---------- summary ----------
        print("\n[summary] persisted in MySQL 996cq:")
        print(f"        task_logs: 5 rows for run_id={run_id} "
              f"(4 success, 1 failure)")
        print(f"        ids: {appended_ids}")

        if args.cleanup:
            print("\n[cleanup] deleting test user/run/logs...")
            with conn.cursor() as cur:
                cleanup_run(cur, run_id, user_id)
            print("        done")
        else:
            print("\n  (use --cleanup to delete the test rows)")

    print("\n[OK] demo complete.")
    return 0


if __name__ == "__main__":
    sys.exit(main())