iTi-Flask-IOT/iti/applications/models/sys_user.py

from marshmallow import post_dump
from sqlalchemy import distinct, select
from iti.applications.extensions import db, jwt, cache_simple
from sqlalchemy.ext.hybrid import hybrid_property
from werkzeug.security import generate_password_hash, check_password_hash
import datetime
from .sys_rel_role_menu import sys_role_menu
from .sys_rel_user_role import sys_user_role
from .sys_menu import SysMenu
from sqlalchemy.orm import joinedload
from iti.applications.common.enums import GenderEnum, StatusEnum
from iti.applications.common.utils import BaseSchema
from apiflask.fields import String, DateTime, Enum, Nested, List
from iti.applications.common.crud import BaseModelMixin
from .sys_role import Role


@jwt.user_lookup_loader
def user_lookup_loader(header, payload):
    """
    用户查找加载器
    """
    identity = payload.get("sub", None)
    # 过期时间
    exp = payload.get("exp", None)

    return load_with_cache(identity=identity, exp=exp)


def load_with_cache(identity, exp):
    if identity is None:
        return None
    cached = cache_simple.get(key=f"user_{identity}")
    if cached is not None:
        return cached
    dbUser = db.session.scalar(
        select(User)
        .filter_by(id=identity)
        .options(joinedload(User.roles).noload(Role.menus), joinedload(User.depts))
    )
    if dbUser is None:
        return None
    dbUser.permissions

    # 计算需缓存时长
    if exp is None:
        return None
    expTime = datetime.datetime.fromtimestamp(exp)
    now = datetime.datetime.now()

    cache_simple.set(
        key=f"user_{identity}", value=dbUser, timeout=(expTime - now).seconds
    )
    return dbUser


class User(BaseModelMixin):
    """
    用户表
    """

    __tablename__ = "sys_user"
    username = db.Column(db.String(64), nullable=False, comment="用户名")
    phone = db.Column(db.String(13), nullable=True, comment="手机号")
    email = db.Column(db.String(255), nullable=True, comment="邮箱")
    _password = db.Column("password", db.String(255), nullable=False, comment="密码")
    realname = db.Column(db.String(32), nullable=True, comment="真实姓名")
    desc = db.Column(db.Text, nullable=True, comment="描述")
    avatar = db.Column(db.String(255), nullable=True, comment="头像")
    gender = db.Column(
        db.Enum(GenderEnum, values_callable=lambda x: [e.value for e in x]),
        nullable=False,
        default=GenderEnum.SECURE.value,
        comment="性别",
    )
    status = db.Column(
        db.Enum(StatusEnum, values_callable=lambda x: [e.value for e in x]),
        nullable=False,
        default=StatusEnum.ENABLED.value,
        comment="状态",
    )
    # 关系
    roles = db.relationship(
        "Role",
        secondary="sys_user_role",
        primaryjoin="User.id == sys_user_role.c.user_id",
        secondaryjoin="and_(Role.id == sys_user_role.c.role_id, Role.status == 'enabled')",
        back_populates="users",
    )
    depts = db.relationship(
        "SysDept",
        secondary="sys_user_dept",
        primaryjoin="User.id == sys_user_dept.c.user_id",
        secondaryjoin="and_(SysDept.id == sys_user_dept.c.dept_id, SysDept.status == 'enabled')",
        back_populates="users",
    )

    @hybrid_property
    def password(self):
        return self._password

    @password.setter
    def password(self, value):
        if value is not None:
            self._password = generate_password_hash(value, method="pbkdf2:sha256")

    def check_password(self, value) -> bool:
        return check_password_hash(self._password, value)

    _permissions = []

    @hybrid_property
    def permissions(self):
        if len(self._permissions) == 0:
            self._permissions = self.get_permissions()
        return self._permissions

    @permissions.setter
    def permissions(self, value):
        self._permissions = value

    def get_permissions(self):
        permissions = db.session.scalars(
            select(distinct(SysMenu.auth_code))
            .join(sys_role_menu, SysMenu.id == sys_role_menu.c.menu_id)
            .join(sys_user_role, sys_role_menu.c.role_id == sys_user_role.c.role_id)
            .filter(
                sys_user_role.c.user_id == self.id,
                SysMenu.status == StatusEnum.ENABLED.value,
                SysMenu.auth_code.isnot(None),
            )
            .order_by(SysMenu.auth_code.asc())
        ).all()
        return permissions


class UserSchema(BaseSchema):
    def __init__(self, *args, **kwargs):
        self.roles_type = kwargs.pop("roles_type", "code")  # code | id
        super().__init__(*args, **kwargs)

    id = String()
    username = String()
    phone = String()
    email = String()
    password = String(load_only=True)
    realname = String()
    avatar = String()
    gender = Enum(GenderEnum, by_value=True)
    status = Enum(StatusEnum, by_value=True)
    desc = String()
    created_at = DateTime(data_key="createdAt", format="%Y-%m-%d %H:%M:%S")
    updated_at = DateTime(data_key="updatedAt", format="%Y-%m-%d %H:%M:%S")

    # 关系
    roles = Nested("RoleSchema", many=True, dump_only=True, exclude=["users"])
    depts = Nested(
        "SysDeptSchema",
        many=True,
        dump_only=True,
        exclude=["users", "children", "parent"],
    )
    permissions = List(String())

    @post_dump
    def patch_roles(self, data, **kwargs):
        """
        角色code列表
        """
        if "roles" in data:
            if self.roles_type == "code":
                role_codes = [role["code"] for role in data["roles"]]
            else:
                role_codes = [role["id"] for role in data["roles"]]
            data["roles"] = role_codes
        # 部门id列表
        if "depts" in data:
            dept_ids = [dept["id"] for dept in data["depts"]]
            data["depts"] = dept_ids
        return data