tg-spam-guard/db/__init__.py

"""Tortoise ORM 初始化"""
import logging
from tortoise import Tortoise, connections
import config

logger = logging.getLogger("spam_guard")

MODELS = [
    "db.models",
]

# aerich 配置
TORTOISE_ORM = {
    "connections": {
        "default": config.DATABASE_URL,
    },
    "apps": {
        "models": {
            "models": MODELS,
            "default_connection": "default",
        },
    },
}


async def ensure_runtime_schema():
    """轻量补齐运行期新增列；只做 ADD COLUMN，不破坏已有数据。"""
    conn = connections.get("default")
    try:
        await conn.execute_script("""
        ALTER TABLE messages ADD COLUMN IF NOT EXISTS manually_deleted BOOLEAN NOT NULL DEFAULT FALSE;
        ALTER TABLE messages ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMPTZ NULL;
        ALTER TABLE messages ADD COLUMN IF NOT EXISTS delete_reason VARCHAR(255) NULL;
        ALTER TABLE user_bindings ADD COLUMN IF NOT EXISTS auth_email VARCHAR(255) NULL;
        ALTER TABLE user_bindings ADD COLUMN IF NOT EXISTS time_expire VARCHAR(100) NULL;
        ALTER TABLE bans ADD COLUMN IF NOT EXISTS chat_id BIGINT NOT NULL DEFAULT 0;
        UPDATE bans AS b
        SET chat_id = m.chat_id
        FROM (
            SELECT user_id, MAX(chat_id) AS chat_id
            FROM messages
            WHERE chat_id <> 0
            GROUP BY user_id
        ) AS m
        WHERE b.chat_id = 0 AND b.user_id = m.user_id;
        ALTER TABLE actions ADD COLUMN IF NOT EXISTS chat_id BIGINT NOT NULL DEFAULT 0;
        ALTER TABLE verify_queue ADD COLUMN IF NOT EXISTS chat_id BIGINT NOT NULL DEFAULT 0;
        ALTER TABLE warnings ADD COLUMN IF NOT EXISTS chat_id BIGINT NOT NULL DEFAULT 0;
        ALTER TABLE game_scores ADD COLUMN IF NOT EXISTS chat_id BIGINT NOT NULL DEFAULT 0;
        CREATE TABLE IF NOT EXISTS whitelists (
            id BIGSERIAL PRIMARY KEY,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            chat_id BIGINT NOT NULL DEFAULT 0,
            user_id BIGINT NOT NULL,
            username VARCHAR(255) NULL,
            first_name VARCHAR(255) NULL,
            reason VARCHAR(500) NULL,
            operator_id BIGINT NULL,
            enabled BOOLEAN NOT NULL DEFAULT TRUE,
            UNIQUE(chat_id, user_id)
        );
        CREATE INDEX IF NOT EXISTS idx_whitelists_user_enabled ON whitelists(user_id, enabled);
        CREATE TABLE IF NOT EXISTS shop_items (
            id BIGSERIAL PRIMARY KEY,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            chat_id BIGINT NOT NULL DEFAULT 0,
            item_key VARCHAR(50) NOT NULL,
            name VARCHAR(255) NOT NULL,
            cost INT NOT NULL DEFAULT 0,
            stock INT NOT NULL DEFAULT -1,
            "desc" VARCHAR(500) NULL,
            enabled BOOLEAN NOT NULL DEFAULT TRUE,
            UNIQUE(chat_id, item_key)
        );
        ALTER TABLE shop_items ADD COLUMN IF NOT EXISTS stock INT NOT NULL DEFAULT -1;
        CREATE INDEX IF NOT EXISTS idx_shop_items_chat_enabled ON shop_items(chat_id, enabled);
        CREATE TABLE IF NOT EXISTS shop_transactions (
            id BIGSERIAL PRIMARY KEY,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            chat_id BIGINT NOT NULL DEFAULT 0,
            user_id BIGINT NOT NULL,
            username VARCHAR(255) NULL,
            item_key VARCHAR(50) NOT NULL,
            item_name VARCHAR(255) NOT NULL,
            cost INT NOT NULL DEFAULT 0,
            status VARCHAR(30) NOT NULL DEFAULT 'success',
            note VARCHAR(500) NULL
        );
        CREATE INDEX IF NOT EXISTS idx_shop_transactions_chat_created ON shop_transactions(chat_id, created_at);
        CREATE INDEX IF NOT EXISTS idx_shop_transactions_chat_user ON shop_transactions(chat_id, user_id);
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS prize_count INT NOT NULL DEFAULT 1;
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS remaining_count INT NOT NULL DEFAULT 1;
        UPDATE lottery_events SET prize_count = 1 WHERE prize_count IS NULL OR prize_count <= 0;
        UPDATE lottery_events SET remaining_count = CASE WHEN status = 'ended' THEN 0 ELSE GREATEST(1, prize_count) END WHERE remaining_count IS NULL OR remaining_count < 0;
        ALTER TABLE lottery_participants ADD COLUMN IF NOT EXISTS is_winner BOOLEAN NOT NULL DEFAULT FALSE;
        DELETE FROM lottery_participants a
        USING lottery_participants b
        WHERE a.id > b.id AND a.event_id = b.event_id AND a.user_id = b.user_id;
        CREATE UNIQUE INDEX IF NOT EXISTS uq_lottery_participants_event_user ON lottery_participants(event_id, user_id);
        CREATE TABLE IF NOT EXISTS lottery_events (
            id BIGSERIAL PRIMARY KEY,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            chat_id BIGINT NOT NULL,
            message_id BIGINT NULL,
            creator_id BIGINT NOT NULL,
            creator_name VARCHAR(255) NULL,
            prize VARCHAR(500) NOT NULL,
            condition_type VARCHAR(50) NOT NULL DEFAULT 'all',
            condition_value INT NOT NULL DEFAULT 0,
            end_type VARCHAR(50) NOT NULL DEFAULT 'people',
            end_value INT NOT NULL DEFAULT 0,
            status VARCHAR(20) NOT NULL DEFAULT 'active',
            winner_id BIGINT NULL,
            winner_name VARCHAR(255) NULL,
            ended_at TIMESTAMPTZ NULL
        );
        CREATE TABLE IF NOT EXISTS lottery_participants (
            id BIGSERIAL PRIMARY KEY,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            event_id BIGINT NOT NULL,
            user_id BIGINT NOT NULL,
            username VARCHAR(255) NULL,
            first_name VARCHAR(255) NULL
        );
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS title VARCHAR(255) NULL;
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS description TEXT NULL;
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS draw_type VARCHAR(50) NOT NULL DEFAULT 'people';
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS draw_at TIMESTAMPTZ NULL;
        ALTER TABLE lottery_events ADD COLUMN IF NOT EXISTS max_participants INT NOT NULL DEFAULT 0;
        CREATE INDEX IF NOT EXISTS idx_lottery_events_chat_status ON lottery_events(chat_id, status);
        CREATE INDEX IF NOT EXISTS idx_lottery_participants_event_user ON lottery_participants(event_id, user_id);
        """)
    except Exception as e:
        logger.warning(f"⚠️ 自动补齐数据库字段失败: {type(e).__name__}: {e}")


async def init_db():
    await Tortoise.init(
        db_url=config.DATABASE_URL,
        modules={"models": MODELS},
    )
    # 先补齐已有表新增列，避免 generate_schemas 创建索引时缺列失败。
    await ensure_runtime_schema()
    # aerich 用 generate_schemas；正式环境用 aerich migrate
    await Tortoise.generate_schemas()
    await ensure_runtime_schema()
    logger.info("✅ 数据库连接成功")


async def close_db():
    await Tortoise.close_connections()
    logger.info("🛑 数据库连接已关闭")