chongming-plugins/auth/__init__.py

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