dekun
170 lines
8.3 KiB
Python
170 lines
8.3 KiB
Python
"""给已经上线的库补上认证字段。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
|