your-workflow/backend/app/core/db.py

44 lines
1.4 KiB
Python

from tortoise.contrib.fastapi import Tortoise
from .logger import logger
from pathlib import Path
from models.db.user import User
import hashlib
async def init_tortoise(
db_type: str, db_config: dict, admin_password: str | None = None
):
if db_type == "mysql":
db_url = (
f"mysql://{db_config['db_user']}:{db_config['db_password']}"
f"@{db_config['db_host']}:{db_config['db_port']}/{db_config['db_name']}"
)
elif db_type == "sqlite":
BASE_DIR = Path(__file__).parent.parent.resolve() / "sqlite"
BASE_DIR.mkdir(parents=True, exist_ok=True)
db_file = BASE_DIR / db_config["db_name"]
db_url = f"sqlite://{db_file}"
else:
raise ValueError("Unsupported db_type")
try:
await Tortoise.init(
db_url=db_url,
modules={"models": ["models.db"]}, # 这里写你模型所在路径
)
await Tortoise.generate_schemas()
if admin_password:
admin_user = await User.filter(username="admin").first()
if not admin_user:
await User(
password=hashlib.md5(admin_password.encode()).hexdigest(), username="admin", level=0
).save()
return True
except Exception as e:
logger.error(e)
raise e
async def close_tortoise():
await Tortoise.close_connections()