chuanqi/cq2_sdk/scripts/demo_user_run_store.py

"""End-to-end demo: UserStore + RunStore against real MySQL.

Setup (raw SQL, since ServerStore/ZoneStore are fold 2):
  1. Open connection
  2. Insert a test server (idempotent on server_name)
  3. Insert a test zone (idempotent on server_id+zone_id+open_time)

Then exercise the new Stores:
  4. UserStore.upsert_by_username -> save a fake token
  5. UserStore.get_by_username    -> read it back, verify id
  6. RunStore.create              -> create a run for that user
  7. RunStore.update_status       -> flip to success after "completion"
  8. RunStore.list_recent         -> read it back, verify
  9. Print a summary table

Cleanup (--cleanup flag):
  10. Delete the test user and run
"""
from __future__ import annotations
import argparse
import os
import sys
from datetime import datetime, timedelta
from pathlib import Path

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

# load .env from project root (not CWD) so script works regardless of where it's invoked
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 Run, User  # noqa: E402
from cq2_sdk.store import RunStore, ServerStore, UserStore, ZoneStore, connection  # noqa: E402
from cq2_sdk.db_models import Zone as DbZone  # noqa: E402

TEST_USERNAME = "demo_user_run_store"
TEST_SERVER = "DEMO_SERVER_玖扬复古176"
TEST_ZONE_ID = "demo_zone_001"


def setup_server_and_zone(servers: ServerStore, zones: ZoneStore) -> tuple[int, int]:
    """Insert test server + zone via the new stores (no raw SQL)."""
    srv = servers.upsert(TEST_SERVER, 99999999, game_id=206923)
    open_time = datetime.now().replace(microsecond=0)
    zones.bulk_upsert(srv.id, [
        DbZone(zone_id=TEST_ZONE_ID, zone_name="Demo Zone",
               host="demo.host", port=12345, open_time=open_time, status="open"),
    ])
    zone = zones.get_latest_openable(srv.id)
    if zone is None:
        raise RuntimeError("setup failed: zone not openable after bulk_upsert")
    return srv.id, zone.id


def cleanup(cur, user_id: int, run_id: int | None) -> None:
    if run_id is not None:
        cur.execute("DELETE FROM runs WHERE id = %s", (run_id,))
    if user_id:
        cur.execute("DELETE FROM users WHERE id = %s", (user_id,))


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

    print("=" * 70)
    print("UserStore + RunStore end-to-end demo")
    print("=" * 70)

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

        # 0. setup server + zone via the new Stores (no raw SQL)
        print("\n[0] setup test server + zone via ServerStore/ZoneStore...")
        server_id, zone_id = setup_server_and_zone(servers, zones)
        print(f"    server_id = {server_id}, zone_id = {zone_id}")

        # 1. user upsert
        print("\n[1] UserStore.upsert_by_username...")
        expires = datetime.now() + timedelta(hours=24)
        user = users.upsert_by_username(TEST_USERNAME, "demo_jwt_token_xyz", expires)
        print(f"    user.id          = {user.id}")
        print(f"    user.username    = {user.username}")
        print(f"    user.token       = {user.token!r}")
        print(f"    user.last_login  = {user.last_login_at}")

        # 2. user read back
        print("\n[2] UserStore.get_by_username...")
        refetched = users.get_by_username(TEST_USERNAME)
        assert refetched is not None, "upsert succeeded but can't read back"
        assert refetched.id == user.id, "id changed after upsert"
        assert refetched.token == "demo_jwt_token_xyz"
        print(f"    refetched.id   = {refetched.id} (matches)")
        print(f"    refetched.token = {refetched.token!r} (matches)")

        # 3. run create
        print("\n[3] RunStore.create...")
        run = runs.create(user_id=user.id, server_id=server_id, zone_id=zone_id)
        print(f"    run.id        = {run.id}")
        print(f"    run.user_id   = {run.user_id}")
        print(f"    run.status    = {run.status}")
        print(f"    run.started_at = {run.started_at}")

        # 4. run update (running -> success)
        print("\n[4] RunStore.update_status(running -> success)...")
        updated = runs.update_status(run.id, "success", current_task="t13")
        print(f"    run.status       = {updated.status}")
        print(f"    run.current_task = {updated.current_task}")
        print(f"    run.completed_at = {updated.completed_at}")

        # 5. list recent
        print("\n[5] RunStore.list_recent(user_id)...")
        recent = runs.list_recent(user.id, limit=5)
        print(f"    found {len(recent)} run(s) for user_id={user.id}")
        for r in recent:
            print(f"      - run.id={r.id} status={r.status} task={r.current_task} "
                  f"started={r.started_at} completed={r.completed_at}")

        # 6. summary
        print("\n[summary] persisted state in MySQL 996cq:")
        print(f"    users.id  = {user.id}  username='{user.username}'")
        print(f"    runs.id   = {run.id}   user_id={run.user_id} status='{updated.status}'")

        if args.cleanup:
            print("\n[6] --cleanup: deleting test rows...")
            with conn.cursor() as cur:
                cleanup(cur, user.id, run.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())