"""跨数据库兼容的轻量 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)"}, ) add_index_if_missing( conn, "accounts", "ix_accounts_owner_id", ("owner_id",), ) 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_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", }, )