"""给已经上线的库补上认证字段。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 statements or need_username_index or index_ready: 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)) _relax_task_node_id() 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 _relax_task_node_id(bind=None) -> None: """旧库 compute_tasks.node_id 为 NOT NULL，排队任务需要改成可空。SQLite 不能直接改空约束，整表重建。""" target = bind or engine inspector = inspect(target) if "compute_tasks" not in inspector.get_table_names(): return columns = {column["name"]: column for column in inspector.get_columns("compute_tasks")} node_column = columns.get("node_id") if node_column is None or node_column.get("nullable", True): return names = [ name for name in ( "id", "node_id", "task_type", "payload", "status", "push_status", "push_error", "dispatched_at", "deadline_at", "responded_at", "success", "result", ) if name in columns ] copied = ", ".join(names) with target.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: conn.execute(text("PRAGMA foreign_keys=OFF")) conn.execute(text("DROP TABLE IF EXISTS compute_tasks_queue")) conn.execute( text( """ CREATE TABLE compute_tasks_queue ( id VARCHAR(36) NOT NULL PRIMARY KEY, node_id INTEGER, task_type VARCHAR(32) NOT NULL, payload JSON NOT NULL, status VARCHAR(16) NOT NULL, push_status VARCHAR(16) NOT NULL, push_error VARCHAR(255), dispatched_at DATETIME NOT NULL, deadline_at DATETIME NOT NULL, responded_at DATETIME, success BOOLEAN, result JSON, FOREIGN KEY(node_id) REFERENCES compute_nodes (id) ) """ ) ) conn.execute(text(f"INSERT INTO compute_tasks_queue ({copied}) SELECT {copied} FROM compute_tasks")) conn.execute(text("DROP TABLE compute_tasks")) conn.execute(text("ALTER TABLE compute_tasks_queue RENAME TO compute_tasks")) conn.execute(text("CREATE INDEX IF NOT EXISTS ix_compute_tasks_node_id ON compute_tasks (node_id)")) conn.execute(text("CREATE INDEX IF NOT EXISTS ix_compute_tasks_status ON compute_tasks (status)")) conn.execute(text("CREATE INDEX IF NOT EXISTS ix_task_node_status ON compute_tasks (node_id, status)")) conn.execute(text("CREATE INDEX IF NOT EXISTS ix_task_dispatched_at ON compute_tasks (dispatched_at)")) conn.execute(text("PRAGMA foreign_keys=ON")) 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