79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
"""
|
||
Auth Worker 插件 — 鉴权管理（登录 / 注册 / 资料 / 改密）
|
||
|
||
基于 Plugin 基类统一生命周期：数据库由基类初始化（auth.db），
|
||
on_init 中执行 admin 幂等种子逻辑，on_cleanup 无需重复 dispose。
|
||
"""
|
||
import logging
|
||
from datetime import datetime
|
||
|
||
from chongming_worker.plugin import Plugin
|
||
from chongming_password import hash_password
|
||
|
||
from .router import * # noqa: F401,F403 触发 @rpc_method 注册
|
||
|
||
# 模块加载时即注册 SQLModel 元数据 —— 基类 _bootstrap 建表（create_all）只创建
|
||
# 已导入的表，若拖到 on_init 才导入会导致全新库缺少 users 表（no such table）。
|
||
from .models.auth import Users # noqa: F401
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class AuthPlugin(Plugin):
|
||
"""Auth 插件 — 统一生命周期。"""
|
||
|
||
type = "auth"
|
||
db_name = "auth.db" # 自动在 data_dir 下建 SQLite 库
|
||
auto_bind = True # 首次初始化时自动 auto_bind_models("auth")
|
||
|
||
async def on_init(self):
|
||
"""首次 create_workers 时调用一次：初始化 admin 账号（幂等）。"""
|
||
from .models.auth import Users
|
||
from sqlalchemy import select
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
db = self.db
|
||
|
||
# 初始化管理员（多进程并发启动时必须幂等，避免 UNIQUE 约束冲突）
|
||
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()
|
||
|
||
# 写入全局变量供 router 使用（数据库由基类在 cleanup 时统一 dispose）
|
||
self.set_var("db", db)
|
||
|
||
async def on_cleanup(self):
|
||
"""cleanup 时调用（数据库由基类处理）。"""
|
||
logger.info("[auth] 清理完成")
|
||
|
||
|
||
# 暴露给 app lifespan 的模块级入口（bound method，签名与旧写法完全一致）
|
||
plugin = AuthPlugin()
|
||
create_workers = plugin.create_workers
|
||
cleanup = plugin.cleanup
|