44 lines
1.4 KiB
Python
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()
|