65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
from pathlib import Path
|
|
from datetime import datetime
|
|
|
|
from chongming_worker import auto_bind_models
|
|
from chongming_worker.runtime import create_workers as _create_workers
|
|
from chongming_password import hash_password
|
|
|
|
from .router import *
|
|
|
|
|
|
auto_bind_models("auth")
|
|
|
|
|
|
async def create_workers(session, config, count=1, worker_manager=None):
|
|
from chongming_database import Database
|
|
from .global_var import set_global_var
|
|
from .models.auth import Users
|
|
from sqlalchemy import select
|
|
|
|
db_path = Path(config.get("data_dir", "./data")) / "auth.db"
|
|
db = Database(db_path)
|
|
await db.initialize()
|
|
|
|
# 初始化管理员（多进程并发启动时必须幂等，避免 UNIQUE 约束冲突）
|
|
from sqlalchemy.exc import IntegrityError
|
|
|
|
async with db.get_session() as sess:
|
|
select_stmt = select(Users) \
|
|
.where(Users.username == "admin") # type: ignore
|
|
result = await sess.execute(select_stmt)
|
|
user: Users | None = result.scalar_one_or_none()
|
|
|
|
# 按用户名未命中时，再按邮箱兜底，避免历史数据中已存在同邮箱用户
|
|
if user is None:
|
|
select_stmt = select(Users) \
|
|
.where(Users.email == "admin@example.com") # type: ignore
|
|
result = await sess.execute(select_stmt)
|
|
user = result.scalar_one_or_none()
|
|
|
|
if user is None:
|
|
user = Users(
|
|
username="admin",
|
|
hashed_password=hash_password("admin"),
|
|
email="admin@example.com",
|
|
created_at=datetime.now(),
|
|
updated_at=datetime.now(),
|
|
is_superuser=True
|
|
)
|
|
sess.add(user)
|
|
try:
|
|
await sess.commit()
|
|
except IntegrityError:
|
|
# 其他并发进程已抢先创建 admin，当前进程忽略即可
|
|
await sess.rollback()
|
|
|
|
set_global_var("db", db)
|
|
|
|
return _create_workers(session, config, count, worker_manager)
|
|
|
|
async def cleanup(worker_item: str):
|
|
from .global_var import db
|
|
|
|
if db is not None:
|
|
await db.dispose()
|
|
print(f"{worker_item} 清理完毕") |