DRF框架使用djangorestframework-simplejwt实现自定义用户类的登录逻辑

1. token校验

生成并校验token，参考上一篇文章：https://blog.csdn.net/u013071014/article/details/144562716?spm=1001.2014.3001.5502

2. 新建role、menu子模块

1. 新建role、menu这两个子模块

   在PyCharm -> Tools -> Run manage.py Task 中执行命令

   # 创建role子模块
   startapp role
   # 创建menu子模块
   startapp menu
   

2. 在role子模块的models.py中定义角色模型

   from rest_framework import serializers
   
   # Create your models here.
   from user.models import SysUser
   
   
   # Create your models here.
   # 系统角色表
   class SysRole(models.Model):
       id = models.AutoField(primary_key=True)
       name = models.CharField(max_length=255, null=True, verbose_name="角色名称")
       code = models.CharField(max_length=255, null=True, verbose_name="角色权限字符串")
       create_time = models.DateField(null=True, verbose_name="创建时间")
       update_time = models.DateField(null=True, verbose_name="更新时间")
       remark = models.CharField(max_length=1024, null=True, verbose_name="备注")
   
       class Meta:
           db_table = 'sys_role'
   
   
   class SysRoleSerializer(serializers.ModelSerializer):
       class Meta:
           model = SysRole
           fields = '__all__'
   
   
   # 用户和角色关联表
   class SysUserRole(models.Model):
       id = models.AutoField(primary_key=True)
       role = models.ForeignKey(SysRole, on_delete=models.PROTECT)
       user = models.ForeignKey(SysUser, on_delete=models.PROTECT)
   
       class Meta:
           db_table = 'sys_user_role'
   

3. 在menu子模块的models.py中定义菜单模型

   from rest_framework import serializers
   from role.models import SysRole
   
   # 菜单表
   class SysMenu(models.Model):
       id = models.AutoField(primary_key=True)
       name = models.CharField(max_length=64, unique=True, verbose_name="菜单名称")
       icon = models.CharField(max_length=255, null=True, verbose_name="菜单图标")
       parent_id = models.IntegerField(null=True, verbose_name="父菜单ID")
       order_num = models.IntegerField(null=True, verbose_name="显示排序")
       path = models.CharField(max_length=255, null=True, verbose_name="路由地址")
       component = models.CharField(max_length=255, null=True, verbose_name="组件路径")
       menu_type = models.CharField(max_length=1, null=True, verbose_name="菜单类型(M目录 C菜单 F按钮)")
       perms = models.CharField(max_length=255, null=True, verbose_name="权限标识")
       create_time = models.DateField(null=True, verbose_name="创建时间")
       update_time = models.DateField(null=True, verbose_name="更新时间")
       remark = models.CharField(max_length=1024, null=True, verbose_name="备注")
   
       # # 比较排序编号大小
       # def __lt__(self, other):
       #     return self.order_num < other.order_num
   
       class Meta:
           db_table = 'sys_menu'
   
   
   class SysMenuSerializer(serializers.ModelSerializer):
       class Meta:
           model = SysMenu
           fields = '__all__'
   
   
   # 序列化对象，带子节点
   class SysMenuTreeSerializer(serializers.ModelSerializer):
       children = serializers.SerializerMethodField()
   
       def get_children(self, obj):
           if hasattr(obj, 'children'):
               serializerMenuList: list[SysMenu] = list()
               for sysMenu in obj.children:
                   serializerMenuList.append(SysMenuSerializer(sysMenu).data)
               return serializerMenuList
   
       class Meta:
           model = SysMenu
           fields = '__all__'
   
   
   # 角色菜单关联表
   class SysRoleMenu(models.Model):
       id = models.AutoField(primary_key=True)
       role = models.ForeignKey(SysRole, on_delete=models.PROTECT)
       menu = models.ForeignKey(SysMenu, on_delete=models.PROTECT)
   
       class Meta:
           db_table = 'sys_role_menu'
   
   
   class SysRoleMenuSerializer(serializers.ModelSerializer):
       class Meta:
           model = SysRoleMenu
           fields = '__all__'
   

4. 在PyCharm -> Tools -> Run manage.py Task 中执行命令，数据迁移

   makemigrations
   migrate
   

5. 在数据库中新建测试数据

   • sys_user

     INSERT INTO `sys_user` VALUES (4, 'admin', 'pbkdf2_sha256$870000$OMbZBGLrBvBuSPTDwCqppS$dGBRvnXNkUsRw0K1jUqGuHqM71ZcvdtX4ldNDkzrhfU=', NULL, '', NULL, 0, NULL, '2024-12-18 07:16:56.702681', '', 1, 0, 0, NULL, '');
     INSERT INTO `sys_user` VALUES (5, 'user1', 'pbkdf2_sha256$870000$q3ltFEG3YlhQmSeYSLjsN4$pR21wvrDTr0lsN6+SG9767UAZ0L0MznsB9Yld6SMfu0=', NULL, '', NULL, 0, NULL, '2024-12-18 07:30:24.686198', '', 1, 0, 0, NULL, '');
     INSERT INTO `sys_user` VALUES (6, 'user2', 'pbkdf2_sha256$870000$nu9rnC9WZesakp9VRsABcM$GyXKI0Rty4XxO3UH7gmfughxgUIHxAj+K1Cp+o+6xWE=', NULL, '', NULL, 0, NULL, '2024-12-18 07:37:00.372880', '', 1, 0, 0, NULL, '');
     

   • sys_role

     INSERT INTO `sys_role` VALUES (1, '超级管理员', 'admin', NULL, NULL, '系统最高权限');
     INSERT INTO `sys_role` VALUES (2, '普通角色', 'common', NULL, NULL, '普通角色');
     INSERT INTO `sys_role` VALUES (3, '测试角色', 'test', NULL, NULL, '测试角色');
     

   • sys_menu

     INSERT INTO `sys_menu` VALUES (1, '系统管理', 'system', 0, 1, '/sys', '', 'M', '', NULL, NULL, '系统管理目录');
     INSERT INTO `sys_menu` VALUES (2, '业务管理', 'monitor', 0, 2, '/bsns', '', 'M', '', NULL, NULL, '业务管理目录');
     INSERT INTO `sys_menu` VALUES (3, '用户管理', 'user', 1, 1, '/sys/user', '', 'C', 'system:user:list', NULL, NULL, '用户管理菜单');
     INSERT INTO `sys_menu` VALUES (4, '角色管理', 'peoples', 1, 2, '/sys/role', '', 'C', 'system:role:list', NULL, NULL, '角色管理菜单');
     INSERT INTO `sys_menu` VALUES (5, '菜单管理', 'tree-table', 1, 3, '/sys/menu', '', 'C', 'system:menu:list', NULL, NULL, '菜单管理菜单');
     INSERT INTO `sys_menu` VALUES (6, '部门管理', 'tree', 2, 1, '/bsns/department', '', 'C', '', NULL, NULL, '部门管理菜单');
     INSERT INTO `sys_menu` VALUES (7, '岗位管理', 'post', 2, 2, '/bsns/post', '', 'C', '', NULL, NULL, '岗位管理菜单');
     

   • sys_user_role

     INSERT INTO `sys_user_role` VALUES (1, 1, 4);
     INSERT INTO `sys_user_role` VALUES (2, 2, 4);
     INSERT INTO `sys_user_role` VALUES (3, 2, 5);
     INSERT INTO `sys_user_role` VALUES (4, 3, 6);
     

   • sys_role_menu

     INSERT INTO `sys_role_menu` VALUES (1, 2, 2);
     INSERT INTO `sys_role_menu` VALUES (2, 6, 2);
     INSERT INTO `sys_role_menu` VALUES (3, 7, 2);
     INSERT INTO `sys_role_menu` VALUES (4, 1, 1);
     INSERT INTO `sys_role_menu` VALUES (5, 3, 1);
     INSERT INTO `sys_role_menu` VALUES (6, 4, 1);
     INSERT INTO `sys_role_menu` VALUES (7, 5, 1);
     INSERT INTO `sys_role_menu` VALUES (8, 2, 1);
     INSERT INTO `sys_role_menu` VALUES (9, 6, 1);
     INSERT INTO `sys_role_menu` VALUES (10, 7, 1);
     

3. 修改登录视图

1. 自定义序列化

   在user/views.py中自定义序列化器，继承 TokenObtainPairSerializer 并重写 validate 方法

   from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
   from menu.models import SysMenu, SysMenuTreeSerializer
   from role.models import SysRole
   from user.models import SysUserSerializer, SysUser
   
   
   class MyTokenObtainPairSerilizer(TokenObtainPairSerializer):
       @classmethod
       def get_token(cls, user):
           token = super().get_token(user)
           # 可在此添加用户相关信息
           token['id'] = user.id
           token['username'] = user.username
           token['email'] = user.email
           return token
   
       def buildTreeMenu(self, sysMenuList):
           resultMenuList: list[SysMenu] = list()
           for menu in sysMenuList:
               # 判断父节点
               if menu.parent_id == 0:
                   resultMenuList.append(menu)
   
               # 寻找子节点
               for e in sysMenuList:
                   if e.parent_id == menu.id:
                       if not hasattr(menu, 'children'):
                           menu.children = list()
                       menu.children.append(e)
   
           return resultMenuList
   
       def validate(self, attrs):
           data = super().validate(attrs)
           # 获取登录用户的角色信息
           user = self.user
           roleList = SysRole.objects.raw(
               "SELECT id, name FROM sys_role WHERE id IN(SELECT role_id FROM sys_user_role WHERE user_id = " + str(
                   user.id) + ")")
           # 有多个角色，可能存在重复的菜单，使用set去重
           menuSet: set[SysMenu] = set()
           for role in roleList:
               # print(role.id, role.name)
               menuList = SysMenu.objects.raw(
                   "SELECT * FROM sys_menu WHERE id IN(SELECT menu_id FROM sys_role_menu WHERE role_id = " + str(
                       role.id) + ")")
               for menu in menuList:
                   # print(menu.id, menu.name)
                   menuSet.add(menu)
   
           # 根据order_num排序
           menuList: list[SysMenu] = list(menuSet)
           sortedMenuList = sorted(menuList, key=lambda menu: menu.order_num)
           # 构造树
           sysMenuList: list[SysMenu] = self.buildTreeMenu(sortedMenuList)
           # 序列化
           serializerMenuList = SysMenuTreeSerializer(sysMenuList, many=True).data
           data['menuList'] = serializerMenuList
           # 返回用户的部分属性
           new_user = {
               "id":user.id,
               "username":user.username,
               "email":user.email,
               "avatar":user.avatar,
               "phone_number":user.phone_number
           }
           data['userInfo'] = new_user
           return data
   

2. 修改LoginView视图

   将serializer_class修改为自定义序列化器

   serializer_class = MyTokenObtainPairSerilizer
   

   class LoginView(TokenObtainPairView):
       serializer_class = MyTokenObtainPairSerilizer
       def post(self, request, *args, **kwargs):
           serializer = self.get_serializer(data=request.data)
           try:
               serializer.is_valid(raise_exception=True)
           except TokenError as e:
               raise InvalidToken(e.args[0])
   
           return Response(serializer.validated_data, status=status.HTTP_200_OK)
   

3. 修改user/urls.py的path

   将登录接口的path修改为LoginView

   # 修改前
   path('login', TokenObtainPairView.as_view(), name='token_obtain_pair'),
   # 修改后
   path('login', LoginView.as_view(), name='token_obtain_pair'),
   

4. 验证接口

登录接口返回了token信息、用户信息、菜单树