""" 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