85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
"""
|
|
Tortoise-ORM 数据库配置
|
|
"""
|
|
from typing import Optional
|
|
from tortoise import Tortoise
|
|
from tortoise.contrib.fastapi import register_tortoise
|
|
from .config import get_settings
|
|
|
|
async def init_db():
|
|
"""初始化数据库连接"""
|
|
settings = get_settings()
|
|
|
|
db_config = {
|
|
"connections": {
|
|
"default": settings.database_url
|
|
},
|
|
"apps": {
|
|
"models": {
|
|
"models": ["app.models", "aerich.models"],
|
|
"default_connection": "default",
|
|
}
|
|
},
|
|
"use_tz": False,
|
|
"timezone": "UTC",
|
|
}
|
|
|
|
await Tortoise.init(db_config)
|
|
|
|
# 根据配置决定是否生成表结构
|
|
if settings.generate_schemas:
|
|
print("🗃️ 生成数据库表结构...")
|
|
await Tortoise.generate_schemas()
|
|
print("✅ 数据库表结构生成完成")
|
|
else:
|
|
print("⏭️ 跳过数据库表结构生成（建议使用迁移工具）")
|
|
|
|
print(f"✅ 数据库连接成功: {settings.database_url}")
|
|
|
|
async def close_db():
|
|
"""关闭数据库连接"""
|
|
await Tortoise.close_connections()
|
|
print("✅ 数据库连接已关闭")
|
|
|
|
def register_db(app):
|
|
"""将数据库注册到 FastAPI 应用"""
|
|
settings = get_settings()
|
|
|
|
register_tortoise(
|
|
app,
|
|
db_url=settings.database_url,
|
|
modules={"models": ["app.models"]},
|
|
generate_schemas=settings.generate_schemas,
|
|
add_exception_handlers=True,
|
|
)
|
|
|
|
return app
|
|
|
|
# 在 app/core/database.py 中添加
|
|
from aerich import Command
|
|
|
|
async def run_migrations():
|
|
"""运行数据库迁移"""
|
|
try:
|
|
command = Command(
|
|
tortoise_config={
|
|
"connections": {
|
|
"default": get_settings().database_url
|
|
},
|
|
"apps": {
|
|
"models": {
|
|
"models": ["app.models", "aerich.models"],
|
|
"default_connection": "default",
|
|
}
|
|
},
|
|
},
|
|
app="models"
|
|
)
|
|
|
|
# 检查并应用迁移
|
|
await command.init()
|
|
await command.upgrade()
|
|
print("✅ 数据库迁移完成")
|
|
except Exception as e:
|
|
print(f"⚠️ 数据库迁移失败: {e}")
|
|
|