HyperNode-AI/app/schema_upgrade.py

"""给已经上线的库补上认证字段。create_all 不会修改旧表。"""

from sqlalchemy import inspect, text

from app.database import engine


def upgrade_schema() -> None:
    inspector = inspect(engine)
    tables = set(inspector.get_table_names())
    statements: list[str] = []
    need_username_index = False
    node_columns: set[str] = set()
    if "users" in tables:
        columns = {column["name"] for column in inspector.get_columns("users")}
        if "username" not in columns:
            statements.append("ALTER TABLE users ADD COLUMN username VARCHAR(32)")
        if "password_hash" not in columns:
            statements.append("ALTER TABLE users ADD COLUMN password_hash VARCHAR(255)")
        if "role" not in columns:
            statements.append("ALTER TABLE users ADD COLUMN role VARCHAR(16) DEFAULT 'user' NOT NULL")
        need_username_index = not _has_username_index(inspector)
    if "compute_tasks" in tables:
        task_columns = {column["name"] for column in inspector.get_columns("compute_tasks")}
        if "result" not in task_columns:
            statements.append("ALTER TABLE compute_tasks ADD COLUMN result JSON")
    if "compute_nodes" in tables:
        node_columns = {column["name"] for column in inspector.get_columns("compute_nodes")}
        if "node_token" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN node_token VARCHAR(128)")
            statements.append(
                "CREATE UNIQUE INDEX IF NOT EXISTS uq_compute_nodes_node_token ON compute_nodes (node_token)"
            )
        if "gpu_model" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN gpu_model VARCHAR(128)")
        if "vram_gb" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN vram_gb FLOAT")
        if "cpu_count" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN cpu_count INTEGER")
        if "memory_gb" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN memory_gb FLOAT")
        if "hardware_weight" not in node_columns:
            statements.append(
                "ALTER TABLE compute_nodes ADD COLUMN hardware_weight FLOAT DEFAULT 0.05 NOT NULL"
            )
        if "gpu_generation" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN gpu_generation VARCHAR(8) DEFAULT '10' NOT NULL")
        if "recovery_days" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN recovery_days INTEGER DEFAULT 540 NOT NULL")
        if "generation_locked" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN generation_locked BOOLEAN DEFAULT 0 NOT NULL")
        if "recoverable_stake_amount" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN recoverable_stake_amount NUMERIC(24, 8)")
        if "node_index" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN node_index INTEGER")
        if "is_whitelist" not in node_columns:
            statements.append(
                "ALTER TABLE compute_nodes ADD COLUMN is_whitelist BOOLEAN DEFAULT 0 NOT NULL"
            )
        if "is_banned" not in node_columns:
            statements.append(
                "ALTER TABLE compute_nodes ADD COLUMN is_banned BOOLEAN DEFAULT 0 NOT NULL"
            )
        if "client_ip" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN client_ip VARCHAR(64)")
            statements.append(
                "CREATE INDEX IF NOT EXISTS ix_compute_nodes_client_ip ON compute_nodes (client_ip)"
            )
        if "hardware_fingerprint" not in node_columns:
            statements.append("ALTER TABLE compute_nodes ADD COLUMN hardware_fingerprint VARCHAR(64)")
            statements.append(
                "CREATE UNIQUE INDEX IF NOT EXISTS uq_compute_nodes_hardware_fingerprint "
                "ON compute_nodes (hardware_fingerprint)"
            )
    index_ready = "compute_nodes" in tables and ("node_index" in node_columns or any("node_index" in item for item in statements))
    if not statements and not need_username_index and not index_ready:
        return
    with engine.begin() as conn:
        for statement in statements:
            conn.execute(text(statement))
        if need_username_index:
            conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS uq_users_username ON users (username)"))
        if index_ready:
            _backfill_node_index(conn, "recovery_days" in node_columns or any("recovery_days" in item for item in statements))


def _backfill_node_index(conn, has_recovery: bool) -> None:
    missing = conn.execute(text("SELECT id FROM compute_nodes WHERE node_index IS NULL ORDER BY id")).all()
    start = conn.execute(text("SELECT COALESCE(MAX(node_index), 0) FROM compute_nodes")).scalar() or 0
    for offset, (row_id,) in enumerate(missing, start=1):
        conn.execute(
            text("UPDATE compute_nodes SET node_index = :index WHERE id = :id"),
            {"index": int(start) + offset, "id": row_id},
        )
    conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS uq_compute_nodes_node_index ON compute_nodes (node_index)"))
    if has_recovery:
        conn.execute(text("UPDATE compute_nodes SET recovery_days = 30 WHERE node_index BETWEEN 1 AND 20"))


def _has_username_index(inspector) -> bool:
    if "users" not in inspector.get_table_names():
        return False
    for index in inspector.get_indexes("users"):
        if index.get("column_names") == ["username"] and index.get("unique"):
            return True
    return False