ngfchl
55 lines
1.6 KiB
Python
55 lines
1.6 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;
|
|
""")
|
|
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("🛑 数据库连接已关闭")
|