hyapp-server/scripts/ops/grant_all_users_coin.py

#!/usr/bin/env python3
"""Grant COIN to every current production user.

This script intentionally creates a fresh command_id for every run. Running it
twice credits users twice. If a run fails halfway, re-running it credits the
already processed users again, matching the requested non-idempotent behavior.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import shlex
import subprocess
import time
import uuid
from pathlib import Path
from typing import Any

import pymysql
import yaml


ASSET_COIN = "COIN"
BIZ_TYPE_MANUAL_CREDIT = "manual_credit"
OUTBOX_STATUS_PENDING = "pending"
TRANSACTION_STATUS_SUCCEEDED = "succeeded"
MAX_INT64 = 2**63 - 1


def normalize_app_code(value: str) -> str:
    value = (value or "").strip().lower()
    return value or "lalu"


def stable_hash(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def transaction_id(app_code: str, command_id: str) -> str:
    return "wtx_" + stable_hash(f"{normalize_app_code(app_code)}|{command_id}")


def event_id(tx_id: str, event_type: str, user_id: int, asset_type: str) -> str:
    return "wev_" + stable_hash(f"{tx_id}|{event_type}|{user_id}|{asset_type}")


def admin_credit_request_hash(
    app_code: str,
    target_user_id: int,
    asset_type: str,
    amount: int,
    operator_user_id: int,
    reason: str,
    evidence_ref: str,
) -> str:
    return stable_hash(
        f"admin_credit|{normalize_app_code(app_code)}|{target_user_id}|{asset_type}|{amount}|{operator_user_id}|{reason}|{evidence_ref}"
    )


def load_yaml(path: str) -> dict[str, Any]:
    payload = yaml.safe_load(read_config_text(path)) or {}
    if not isinstance(payload, dict):
        raise RuntimeError(f"config must be a YAML object: {path}")
    return payload


def read_config_text(location: str) -> str:
    location = location.strip()
    if looks_like_ssh_location(location):
        host, remote_path = location.split(":", 1)
        return subprocess.check_output(
            ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5", host, "cat " + shlex.quote(remote_path)],
            text=True,
        )
    return Path(location).expanduser().read_text(encoding="utf-8")


def looks_like_ssh_location(location: str) -> bool:
    if location.startswith("/") or location.startswith("./") or location.startswith("../") or "://" in location:
        return False
    head, sep, tail = location.partition(":")
    return bool(sep and tail and ("@" in head or "." in head or head.replace("-", "").isalnum()))


def parse_go_mysql_dsn(dsn: str) -> dict[str, Any]:
    # Supports the project form: USER:PASSWORD@tcp(HOST:PORT)/DATABASE?params.
    prefix, rest = dsn.split("@tcp(", 1)
    user, password = prefix.split(":", 1) if ":" in prefix else (prefix, "")
    addr, rest = rest.split(")/", 1)
    host, port = addr.rsplit(":", 1) if ":" in addr else (addr, "3306")
    database = rest.split("?", 1)[0]
    return {
        "host": host,
        "port": int(port),
        "user": user,
        "password": password,
        "database": database,
        "charset": "utf8mb4",
        "autocommit": False,
        "connect_timeout": 8,
        "read_timeout": 60,
        "write_timeout": 60,
    }


def connect_from_config(path: str):
    cfg = load_yaml(path)
    dsn = str(cfg.get("mysql_dsn") or "").strip()
    if not dsn:
        raise RuntimeError(f"mysql_dsn is required in {path}")
    return pymysql.connect(**parse_go_mysql_dsn(dsn))


def current_ms() -> int:
    return int(time.time() * 1000)


def fetch_users(conn, app_code: str, statuses: list[str]) -> list[int]:
    where = ["app_code = %s"]
    args: list[Any] = [app_code]
    if statuses:
        where.append("status IN (" + ",".join(["%s"] * len(statuses)) + ")")
        args.extend(statuses)
    with conn.cursor() as cursor:
        cursor.execute(f"SELECT user_id FROM users WHERE {' AND '.join(where)} ORDER BY user_id ASC", args)
        return [int(row[0]) for row in cursor.fetchall()]


def fetch_user_status_counts(conn, app_code: str) -> dict[str, int]:
    with conn.cursor() as cursor:
        cursor.execute("SELECT status, COUNT(*) FROM users WHERE app_code = %s GROUP BY status ORDER BY status", (app_code,))
        return {str(status): int(count) for status, count in cursor.fetchall()}


def grant_one(conn, *, app_code: str, user_id: int, amount: int, operator_user_id: int, reason: str, evidence_ref: str, run_id: str) -> tuple[str, int]:
    command_id = f"bulk_coin_grant:{run_id}:{user_id}"
    tx_id = transaction_id(app_code, command_id)
    req_hash = admin_credit_request_hash(app_code, user_id, ASSET_COIN, amount, operator_user_id, reason, evidence_ref)
    now_ms = current_ms()
    metadata = {
        "app_code": app_code,
        "target_user_id": user_id,
        "asset_type": ASSET_COIN,
        "amount": amount,
        "operator_user_id": operator_user_id,
        "reason": reason,
        "evidence_ref": evidence_ref,
    }

    cursor = conn.cursor()
    try:
        cursor.execute("START TRANSACTION")
        # Positive manual credit may create a missing COIN account; version then advances exactly like wallet-service.
        cursor.execute(
            """
            INSERT IGNORE INTO wallet_accounts (
              app_code, user_id, asset_type, available_amount, frozen_amount, version, created_at_ms, updated_at_ms
            ) VALUES (%s, %s, %s, 0, 0, 1, %s, %s)
            """,
            (app_code, user_id, ASSET_COIN, now_ms, now_ms),
        )
        cursor.execute(
            """
            SELECT available_amount, frozen_amount, version
            FROM wallet_accounts
            WHERE app_code = %s AND user_id = %s AND asset_type = %s
            FOR UPDATE
            """,
            (app_code, user_id, ASSET_COIN),
        )
        row = cursor.fetchone()
        if row is None:
            raise RuntimeError(f"wallet account creation failed for user_id={user_id}")
        available, frozen, version = int(row[0]), int(row[1]), int(row[2])
        available_after = available + amount
        if available_after < 0 or available_after > MAX_INT64 - frozen:
            raise RuntimeError(f"wallet balance overflow for user_id={user_id}")

        cursor.execute(
            """
            INSERT INTO wallet_transactions (
              app_code, transaction_id, command_id, biz_type, status, request_hash, external_ref,
              metadata_json, created_at_ms, updated_at_ms
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            (
                app_code,
                tx_id,
                command_id,
                BIZ_TYPE_MANUAL_CREDIT,
                TRANSACTION_STATUS_SUCCEEDED,
                req_hash,
                evidence_ref,
                json.dumps(metadata, separators=(",", ":"), ensure_ascii=False),
                now_ms,
                now_ms,
            ),
        )
        cursor.execute(
            """
            UPDATE wallet_accounts
            SET available_amount = %s, frozen_amount = %s, version = version + 1, updated_at_ms = %s
            WHERE app_code = %s AND user_id = %s AND asset_type = %s AND version = %s
            """,
            (available_after, frozen, now_ms, app_code, user_id, ASSET_COIN, version),
        )
        if cursor.rowcount != 1:
            raise RuntimeError(f"wallet account version conflict for user_id={user_id}")
        cursor.execute(
            """
            INSERT INTO wallet_entries (
              app_code, transaction_id, user_id, asset_type, available_delta, frozen_delta,
              available_after, frozen_after, counterparty_user_id, room_id, created_at_ms
            ) VALUES (%s, %s, %s, %s, %s, 0, %s, %s, 0, '', %s)
            """,
            (app_code, tx_id, user_id, ASSET_COIN, amount, available_after, frozen, now_ms),
        )

        event_payload = {
            "transaction_id": tx_id,
            "command_id": command_id,
            "user_id": user_id,
            "asset_type": ASSET_COIN,
            "available_delta": amount,
            "frozen_delta": 0,
            "available_after": available_after,
            "frozen_after": frozen,
            "balance_version": version + 1,
            "metadata": metadata,
            "created_at_ms": now_ms,
        }
        cursor.execute(
            """
            INSERT INTO wallet_outbox (
              app_code, event_id, event_type, transaction_id, command_id, user_id, asset_type,
              available_delta, frozen_delta, payload, status, retry_count, created_at_ms, updated_at_ms
            ) VALUES (%s, %s, 'WalletBalanceChanged', %s, %s, %s, %s, %s, 0, %s, %s, 0, %s, %s)
            """,
            (
                app_code,
                event_id(tx_id, "WalletBalanceChanged", user_id, ASSET_COIN),
                tx_id,
                command_id,
                user_id,
                ASSET_COIN,
                amount,
                json.dumps(event_payload, separators=(",", ":"), ensure_ascii=False),
                OUTBOX_STATUS_PENDING,
                now_ms,
                now_ms,
            ),
        )
        conn.commit()
        return tx_id, available_after
    except Exception:
        conn.rollback()
        raise
    finally:
        cursor.close()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Non-idempotently grant COIN to every current user.")
    parser.add_argument(
        "--user-config",
        default="services/user-service/configs/config.yaml",
        help="YAML containing user-service mysql_dsn. Supports local path or user@host:/path/config.yaml.",
    )
    parser.add_argument(
        "--wallet-config",
        default="services/wallet-service/configs/config.yaml",
        help="YAML containing wallet-service mysql_dsn. Supports local path or user@host:/path/config.yaml.",
    )
    parser.add_argument("--app-code", default="lalu", help="Tenant app_code to grant.")
    parser.add_argument("--amount", type=int, default=1_000_000, help="COIN amount credited to each user per execution.")
    parser.add_argument("--operator-user-id", type=int, default=0, help="Required with --yes; stored in manual credit metadata.")
    parser.add_argument("--reason", default="bulk coin grant to all users", help="Manual credit reason stored with every transaction.")
    parser.add_argument("--evidence-ref-prefix", default="ops-bulk-coin-grant", help="Evidence reference prefix; run_id and user_id are appended.")
    parser.add_argument("--user-status", action="append", default=[], help="Optional repeatable users.status filter. Omit to include all users.")
    parser.add_argument("--progress-every", type=int, default=100, help="Print progress every N credited users.")
    parser.add_argument("--yes", action="store_true", help="Actually write wallet credits. Omit for dry-run.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    app_code = normalize_app_code(args.app_code)
    if args.amount <= 0 or args.amount > MAX_INT64:
        raise RuntimeError("--amount must be a positive int64")
    if args.yes and args.operator_user_id <= 0:
        raise RuntimeError("--operator-user-id is required with --yes")
    if args.yes and not args.reason.strip():
        raise RuntimeError("--reason is required with --yes")

    user_conn = connect_from_config(args.user_config)
    wallet_conn = connect_from_config(args.wallet_config)
    try:
        status_counts = fetch_user_status_counts(user_conn, app_code)
        users = fetch_users(user_conn, app_code, [status.strip() for status in args.user_status if status.strip()])
        summary = {
            "mode": "execute" if args.yes else "dry_run",
            "app_code": app_code,
            "amount_per_user": args.amount,
            "target_user_count": len(users),
            "total_coin_delta": len(users) * args.amount,
            "user_status_counts": status_counts,
            "user_status_filter": args.user_status,
        }
        print(json.dumps(summary, ensure_ascii=False, indent=2))
        if not args.yes:
            print("dry-run only; rerun with --yes --operator-user-id <id> to credit users")
            return 0

        run_id = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) + "_" + uuid.uuid4().hex[:12]
        credited = 0
        for user_id in users:
            evidence_ref = f"{args.evidence_ref_prefix}:{run_id}:{user_id}"
            grant_one(
                wallet_conn,
                app_code=app_code,
                user_id=user_id,
                amount=args.amount,
                operator_user_id=args.operator_user_id,
                reason=args.reason.strip(),
                evidence_ref=evidence_ref,
                run_id=run_id,
            )
            credited += 1
            if args.progress_every > 0 and (credited % args.progress_every == 0 or credited == len(users)):
                print(json.dumps({"run_id": run_id, "credited": credited, "total": len(users)}, ensure_ascii=False))
        print(json.dumps({"ok": True, "run_id": run_id, "credited": credited, "amount_per_user": args.amount}, ensure_ascii=False, indent=2))
        return 0
    finally:
        user_conn.close()
        wallet_conn.close()


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