ngfchl
84 lines
3.0 KiB
Python
84 lines
3.0 KiB
Python
"""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;
|
|
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
|
|
);
|
|
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},
|
|
)
|
|
# 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("🛑 数据库连接已关闭")
|