douyin/backend/models/db_migrate.py

"""跨数据库兼容的轻量 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 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"},
    )


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_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",
        },
    )