322 lines
9.3 KiB
Python
322 lines
9.3 KiB
Python
"""跨数据库兼容的轻量 schema 迁移。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from sqlalchemy import inspect, text
|
|
|
|
|
|
def _table_columns(conn, table: str) -> set[str]:
|
|
try:
|
|
insp = inspect(conn)
|
|
return {col["name"] for col in insp.get_columns(table)}
|
|
except Exception:
|
|
return set()
|
|
|
|
|
|
def _dialect(conn) -> str:
|
|
return conn.dialect.name
|
|
|
|
|
|
def _bool_default(conn, value: bool = True) -> str:
|
|
if _dialect(conn) == "postgresql":
|
|
return "TRUE" if value else "FALSE"
|
|
return "1" if value else "0"
|
|
|
|
|
|
def add_column_if_missing(conn, table: str, column: str, ddl_by_dialect: dict[str, str]) -> None:
|
|
cols = _table_columns(conn, table)
|
|
if not cols or column in cols:
|
|
return
|
|
dialect = _dialect(conn)
|
|
ddl = ddl_by_dialect.get(dialect) or ddl_by_dialect.get("default")
|
|
if ddl:
|
|
conn.execute(text(ddl))
|
|
|
|
|
|
def add_index_if_missing(
|
|
conn,
|
|
table: str,
|
|
index_name: str,
|
|
columns: tuple[str, ...],
|
|
) -> None:
|
|
"""Create one portable index without relying on dialect-specific IF NOT EXISTS."""
|
|
try:
|
|
insp = inspect(conn)
|
|
if not insp.has_table(table):
|
|
return
|
|
existing = {item.get("name") for item in insp.get_indexes(table)}
|
|
except Exception:
|
|
return
|
|
if index_name in existing:
|
|
return
|
|
safe_columns = ", ".join(columns)
|
|
conn.execute(text(f"CREATE INDEX {index_name} ON {table} ({safe_columns})"))
|
|
|
|
|
|
def migrate_accounts_table(conn) -> None:
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"cookie_data",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN cookie_data TEXT"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"cookie_updated_at",
|
|
{
|
|
"default": "ALTER TABLE accounts ADD COLUMN cookie_updated_at DATETIME",
|
|
"postgresql": "ALTER TABLE accounts ADD COLUMN cookie_updated_at TIMESTAMP",
|
|
},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"im_session_data",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN im_session_data TEXT"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"reply_delay_seconds",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN reply_delay_seconds INTEGER DEFAULT 0"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"reply_cooldown_seconds",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN reply_cooldown_seconds INTEGER"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"follow_welcome_enabled",
|
|
{
|
|
"default": "ALTER TABLE accounts ADD COLUMN follow_welcome_enabled BOOLEAN DEFAULT 0",
|
|
"postgresql": "ALTER TABLE accounts ADD COLUMN follow_welcome_enabled BOOLEAN DEFAULT FALSE",
|
|
},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"follow_welcome_content",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN follow_welcome_content TEXT"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"owner_id",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN owner_id INTEGER"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"user_agent",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN user_agent TEXT"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"avatar_url",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN avatar_url TEXT"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"douyin_uid",
|
|
{"default": "ALTER TABLE accounts ADD COLUMN douyin_uid VARCHAR(64)"},
|
|
)
|
|
|
|
|
|
def migrate_account_videos_table(conn) -> None:
|
|
add_column_if_missing(
|
|
conn,
|
|
"account_videos",
|
|
"media_type",
|
|
{"default": "ALTER TABLE account_videos ADD COLUMN media_type VARCHAR(20)"},
|
|
)
|
|
|
|
|
|
def migrate_message_logs_table(conn) -> None:
|
|
add_column_if_missing(
|
|
conn,
|
|
"message_logs",
|
|
"sender_avatar",
|
|
{"default": "ALTER TABLE message_logs ADD COLUMN sender_avatar TEXT"},
|
|
)
|
|
add_index_if_missing(
|
|
conn,
|
|
"message_logs",
|
|
"ix_message_logs_account_created_at",
|
|
("account_id", "created_at"),
|
|
)
|
|
add_index_if_missing(
|
|
conn,
|
|
"message_logs",
|
|
"ix_message_logs_created_at",
|
|
("created_at",),
|
|
)
|
|
add_index_if_missing(
|
|
conn,
|
|
"message_logs",
|
|
"ix_message_logs_status_account_id",
|
|
("status", "account_id"),
|
|
)
|
|
add_index_if_missing(
|
|
conn,
|
|
"received_message_logs",
|
|
"ix_received_message_logs_account_created_at",
|
|
("account_id", "created_at"),
|
|
)
|
|
add_index_if_missing(
|
|
conn,
|
|
"system_logs",
|
|
"ix_system_logs_account_created_at",
|
|
("account_id", "created_at"),
|
|
)
|
|
|
|
|
|
def migrate_rules_table(conn) -> None:
|
|
add_column_if_missing(
|
|
conn,
|
|
"rules",
|
|
"owner_id",
|
|
{"default": "ALTER TABLE rules ADD COLUMN owner_id INTEGER"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"rules",
|
|
"sort_order",
|
|
{"default": "ALTER TABLE rules ADD COLUMN sort_order INTEGER DEFAULT 0"},
|
|
)
|
|
cols = _table_columns(conn, "rules")
|
|
if cols and "sort_order" in cols:
|
|
conn.execute(
|
|
text("UPDATE rules SET sort_order = id WHERE sort_order IS NULL OR sort_order = 0")
|
|
)
|
|
|
|
|
|
def migrate_roles_table(conn) -> None:
|
|
"""Ensure roles table exists (create_all usually handles this; keep as safety net)."""
|
|
try:
|
|
insp = inspect(conn)
|
|
if insp.has_table("roles"):
|
|
return
|
|
except Exception:
|
|
return
|
|
dialect = _dialect(conn)
|
|
if dialect == "postgresql":
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
CREATE TABLE roles (
|
|
id SERIAL PRIMARY KEY,
|
|
code VARCHAR(50) NOT NULL UNIQUE,
|
|
label VARCHAR(100) NOT NULL,
|
|
description VARCHAR(255),
|
|
is_system BOOLEAN DEFAULT FALSE,
|
|
is_admin BOOLEAN DEFAULT FALSE,
|
|
permissions TEXT NOT NULL DEFAULT '[]',
|
|
created_at TIMESTAMP,
|
|
updated_at TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
else:
|
|
conn.execute(
|
|
text(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS roles (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
code VARCHAR(50) NOT NULL UNIQUE,
|
|
label VARCHAR(100) NOT NULL,
|
|
description VARCHAR(255),
|
|
is_system BOOLEAN DEFAULT 0,
|
|
is_admin BOOLEAN DEFAULT 0,
|
|
permissions TEXT NOT NULL DEFAULT '[]',
|
|
created_at DATETIME,
|
|
updated_at DATETIME
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
add_index_if_missing(conn, "roles", "ix_roles_code", ("code",))
|
|
|
|
|
|
def migrate_users_table(conn) -> None:
|
|
add_column_if_missing(
|
|
conn,
|
|
"users",
|
|
"email",
|
|
{"default": "ALTER TABLE users ADD COLUMN email VARCHAR(255)"},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"users",
|
|
"email_verified",
|
|
{
|
|
"default": "ALTER TABLE users ADD COLUMN email_verified BOOLEAN DEFAULT 0",
|
|
"postgresql": "ALTER TABLE users ADD COLUMN email_verified BOOLEAN DEFAULT FALSE",
|
|
},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"users",
|
|
"email_verified_at",
|
|
{
|
|
"default": "ALTER TABLE users ADD COLUMN email_verified_at DATETIME",
|
|
"postgresql": "ALTER TABLE users ADD COLUMN email_verified_at TIMESTAMP",
|
|
},
|
|
)
|
|
add_column_if_missing(
|
|
conn,
|
|
"users",
|
|
"max_accounts",
|
|
{"default": "ALTER TABLE users ADD COLUMN max_accounts INTEGER DEFAULT 3"},
|
|
)
|
|
cols = _table_columns(conn, "users")
|
|
if not cols:
|
|
return
|
|
verified = _bool_default(conn, True)
|
|
conn.execute(text(f"UPDATE users SET email_verified = {verified} WHERE email_verified IS NULL"))
|
|
conn.execute(text("UPDATE users SET max_accounts = -1 WHERE role = 'admin'"))
|
|
conn.execute(text("UPDATE users SET max_accounts = 3 WHERE max_accounts IS NULL"))
|
|
|
|
|
|
def migrate_payment_orders_table(conn) -> None:
|
|
cols = _table_columns(conn, "payment_orders")
|
|
if not cols:
|
|
return
|
|
add_column_if_missing(
|
|
conn,
|
|
"payment_orders",
|
|
"slots_applied",
|
|
{
|
|
"default": "ALTER TABLE payment_orders ADD COLUMN slots_applied BOOLEAN DEFAULT 0",
|
|
"postgresql": "ALTER TABLE payment_orders ADD COLUMN slots_applied BOOLEAN DEFAULT FALSE",
|
|
},
|
|
)
|
|
applied = _bool_default(conn, True)
|
|
conn.execute(
|
|
text(
|
|
"UPDATE payment_orders SET slots_applied = "
|
|
f"{applied} WHERE status = 'paid' AND (slots_applied IS NULL OR slots_applied = 0)"
|
|
)
|
|
)
|
|
|
|
|
|
def migrate_accounts_quota_disabled(conn) -> None:
|
|
cols = _table_columns(conn, "accounts")
|
|
if not cols:
|
|
return
|
|
add_column_if_missing(
|
|
conn,
|
|
"accounts",
|
|
"quota_disabled",
|
|
{
|
|
"default": "ALTER TABLE accounts ADD COLUMN quota_disabled BOOLEAN DEFAULT 0",
|
|
"postgresql": "ALTER TABLE accounts ADD COLUMN quota_disabled BOOLEAN DEFAULT FALSE",
|
|
},
|
|
)
|