worldquant-alpha-system/backend/app/cli.py

"""Administrative commands run inside the backend container."""

import argparse
import asyncio
import getpass
import math

from sqlalchemy import delete, update

from .config import Settings
from .db import create_database
from .models import Admin, LoginSession, MCPToken, now
from .security import password_hasher


async def reset_password():
    password = getpass.getpass("New admin password (12+ characters): ")
    if len(password) < 12 or password != getpass.getpass("Confirm password: "):
        raise SystemExit("Password too short or confirmation does not match")
    engine, sessions = create_database(Settings().database_url)
    async with sessions() as db:
        admin = await db.get(Admin, 1)
        admin.password_hash = password_hasher.hash(password)
        await db.execute(delete(LoginSession))
        await db.execute(update(MCPToken).where(MCPToken.revoked_at.is_(None)).values(revoked_at=now()))
        await db.commit()
    await engine.dispose()
    print("Admin password updated; all system sessions and MCP tokens revoked.")


async def token_command(args):
    import json

    from sqlalchemy import select

    from .mcp_api.auth import create_token
    from .models import MCPToken, now
    from .research.serialization import encode_snapshot

    engine, sessions = create_database(Settings().database_url)
    try:
        async with sessions.begin() as db:
            if args.command == "mcp-token-create":
                row, secret = await create_token(db, args.name, args.scope, args.days)
                result = {"id": row.id, "name": row.name, "scopes": row.scopes,
                          "expires_at": row.expires_at, "token": secret}
            elif args.command == "mcp-token-revoke":
                row = await db.get(MCPToken, args.token_id)
                if not row:
                    raise ValueError("令牌不存在")
                row.revoked_at = row.revoked_at or now()
                result = {"id": row.id, "revoked": True}
            else:
                rows = list(await db.scalars(select(MCPToken).order_by(MCPToken.created_at.desc())))
                result = [{k: getattr(row, k) for k in
                           ("id", "name", "scopes", "created_at", "expires_at", "revoked_at")} for row in rows]
        # Reveal only after the transaction has committed successfully.
        print(json.dumps(encode_snapshot(result), ensure_ascii=False, indent=2))
    finally:
        await engine.dispose()


async def catalog_sync_command(args):
    """Enqueue on the existing runner, then observe without owning the upstream session."""
    import json
    import time

    from fastapi import HTTPException

    from .business import Business
    from .catalog.contracts import CatalogJobInput, Scope
    from .catalog.service import Catalog
    from .models import Job

    engine, sessions = create_database(Settings().database_url)
    try:
        async with sessions.begin() as db:
            if args.resume_job:
                if args.region or args.universe or args.delay is not None:
                    raise ValueError("--resume-job 不能同时指定新范围")
                job = await db.get(Job, args.resume_job)
                if not job or job.kind != "catalog_full_sync":
                    raise ValueError("只能恢复已有全量目录任务")
                result = await Business(db).retry_job(job.id)
                job_id = result["id"]
            else:
                if not args.region or not args.universe or args.delay is None:
                    raise ValueError("需要 --region、--universe 和 --delay")
                scope = Scope(instrument_type=args.instrument_type, region=args.region,
                              universe=args.universe, delay=args.delay)
                job = await Catalog(db).create_job(CatalogJobInput(scope=scope), full=True)
                job_id = job.id
        deadline, previous = time.monotonic() + args.wait_timeout, None
        while True:
            async with sessions() as db:
                job = await db.get(Job, job_id)
                data = dict(job_id=job.id, status=job.status, processed=job.processed,
                            total=job.total, failed=job.failed, checkpoint=job.checkpoint, error=job.error)
            current = json.dumps(data, ensure_ascii=False, sort_keys=True)
            if current != previous:
                print(current, flush=True)
                previous = current
            if job.status == "completed":
                return 0
            if job.status in ("failed", "completed_with_errors", "cancelled"):
                return 2 if job.checkpoint.get("error_code") == "invalid_scope" else 1
            if job.status in ("waiting_auth", "waiting_connection"):
                return 3
            if time.monotonic() >= deadline:
                print(f"等待超时；后台任务 {job_id} 继续执行", flush=True)
                return 4
            await asyncio.sleep(min(5, max(0, deadline - time.monotonic())))
    except HTTPException as exc:
        print(str(exc.detail), flush=True)
        return 3 if exc.status_code == 409 else 1
    finally:
        await engine.dispose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("reset-password")
    create = commands.add_parser("mcp-token-create")
    create.add_argument("--name", required=True)
    create.add_argument("--scope", action="append", default=None)
    create.add_argument("--days", type=int, default=90)
    commands.add_parser("mcp-token-list")
    revoke = commands.add_parser("mcp-token-revoke")
    revoke.add_argument("token_id")
    sync = commands.add_parser("catalog-sync", help="全量同步一个范围的数据集及全部字段")
    sync.add_argument("--region")
    sync.add_argument("--universe")
    sync.add_argument("--delay", type=int, choices=range(0, 10))
    sync.add_argument("--instrument-type", default="EQUITY")
    sync.add_argument("--resume-job")
    sync.add_argument("--wait-timeout", type=float, default=21600)
    args = parser.parse_args()
    if args.command == "catalog-sync" and (not math.isfinite(args.wait_timeout) or args.wait_timeout <= 0):
        parser.error("--wait-timeout 必须大于 0")
    try:
        if args.command == "catalog-sync":
            raise SystemExit(asyncio.run(catalog_sync_command(args)))
        asyncio.run(reset_password() if args.command == "reset-password" else token_command(args))
    except ValueError as exc:
        parser.error(str(exc))