Flask-Login 使用多用户表登录的代码实现优化

Steak Overcooked
,
1

关于Flask-Login多用户表登陆

​

为什么发帖

  • 课设要求做一个面向多种用户类型的网站，且每个用户类型的操作界面、属性、权限都不一样。
  • 想过使用 一个 User 表，附上role属性。但每个角色的属性都不一样而且有比较大的差别，便作罢了
  • 我用了三个用户表，完成编写后，觉得不太简洁，想问一下大家有没有别的思路
  • 直接对session进行操作会不会很粗暴
  • 看了叫做flask-login-multi的扩展，18年就不再维护了

数据库模型

class BaseUser(db.Model, UserMixin):
    __abstract__ = True

    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(30))
    password_hash = db.Column(db.String(128))

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)


class Student(BaseUser):
    stu_id = db.Column(db.String(32), unique=True, index=True)
    role = db.Column(db.Integer, default=0)
    # ...


class Teacher(BaseUser):
    tea_id = db.Column(db.String(32), unique=True, index=True)
    role = db.Column(db.Integer, default=1)
    # ...


class Admin(BaseUser):
    adm_id = db.Column(db.String(32), unique=True, index=True)
    role = db.Column(db.Integer, default=2)
    # ...

登录表单

class LoginForm(FlaskForm):
    id = StringField('ID', validators=[DataRequired()])
    password = PasswordField(u'密码', validators=[DataRequired()])
    role = RadioField(u'身份', choices=[('student', u'学生'), ('teacher', u'老师'), ('admin', u'管理员')],
                      validators=[DataRequired()])
    remember = BooleanField(u'记住我')
    submit = SubmitField(u'登录')

登录视图

@bp.route('/login', methods=['GET', 'POST'])
def login():
    if current_user.is_authenticated:
        return redirect(url_for('main.index'))
    form = LoginForm()
    if form.validate_on_submit():
        id = form.id.data
        role = form.role.data
        if role == 'student':
            user = Student.query.filter_by(stu_id=id).first()
            if user is not None and user.check_password(form.password.data):
                login_user(user, remember=form.remember.data)
                session['role'] = 'student'
            else:
                flash(u'用户名或密码错误', 'warning')
                return redirect(url_for('.login'))
        if role == 'teacher':
            user = Teacher.query.filter_by(tea_id=id).first()
            if user is not None and user.check_password(form.password.data):
                login_user(user, remember=form.remember.data)
                session['role'] = 'teacher'
            else:
                flash(u'用户名或密码错误', 'warning')
                return redirect(url_for('.login'))
        if role == 'admin':
            user = Admin.query.filter_by(adm_id=id).first()
            if user is not None and user.check_password(form.password.data):
                login_user(user, remember=form.remember.data)
                session['role'] = 'admin'
            else:
                flash(u'用户名或密码错误', 'warning')
                return redirect(url_for('.login'))
        flash(u'{}登陆成功'.format(user.name), 'success')
        return redirect(url_for('main.index'))
    return render_template('auth/login.html', form=form)

login_manager 的 userloader

@login.user_loader
def load_user(user_id):
    if session.get('role') == 'student':
        return Student.query.get(int(user_id))
    elif session.get('role') == 'teacher':
        return Teacher.query.get(int(user_id))
    else:
        return Admin.query.get(int(user_id))

用户过滤装饰器

def admin_required(func):
    @wraps(func)
    def decorated_function(*args, **kwargs):
        if not current_user.role == 2:
            return abort(403)
        return func(*args, **kwargs)
    return decorated_function


def student_required(func):
    @wraps(func)
    def decorated_function(*args, **kwargs):
        if not current_user.role == 0:
            return abort(403)
        return func(*args, **kwargs)
    return decorated_function


def teacher_required(func):
    @wraps(func)
    def decorated_function(*args, **kwargs):
        if not current_user.role == 1:
            return abort(403)
        return func(*args, **kwargs)
    return decorated_function

​ 谢谢你可以看到这里，如果你有更好的解决方法，麻烦告诉我，再次感谢！

1 个赞
2

简单想到几条改进和优化建议：

模型类定义

  • 角色（role）在数据库里与其用数字不如用英文字符表示比较直观。方便随时识别，而不用翻到模型定义才想起来教师的 role 值是 1。而且这可以和你 session 的设置值、登陆表单选项值相一致。
  • 既然 BaseUser 类里已经定义了主键 id，三个子类里的 id 列似乎没必要？而且三个子类的 role 定义是不是可以放到 BaseUser 基类里？

装饰器合并

三个装饰器可以简化成一个接受 role 参数的装饰器：

def role_required(role_name):
    def decorator(func):
        @wraps(func)
        def decorated_function(*args, **kwargs):
            if current_user.role != role_name:
                abort(403)
            return func(*args, **kwargs)
        return decorated_function
    return decorator

假设改成了英文表示 role，用法类似：

@app.route('/dashboard')
@role_required('admin')
def dashboard():
    # ...

使用字典替代重复代码和 if 判断

load_user 函数那里可以考虑用字典匹配来替代重复的代码和多个 if 判断：

role_mapping = {
    'student': Student,
    'teacher': Teacher,
    'admin': Admin
}

@login.user_loader
def load_user(user_id):
    role = session.get('role', 'admin')
    return role_mapping[role].query.get(int(user_id))

类似，登陆视图也可以做同样的简化处理（不过前提是统一使用 BaseUser 的 id 列作为主键）：

@bp.route('/login', methods=['GET', 'POST'])
def login():
    if current_user.is_authenticated:
        return redirect(url_for('main.index'))
    form = LoginForm()
    if form.validate_on_submit():
        id = form.id.data
        role = form.role.data
        user = role_mapping[role].query.get_or_404(id)
        if user is not None and user.check_password(form.password.data):
            login_user(user, remember=form.remember.data)
            session['role'] = role
        else:
            flash(u'用户名或密码错误', 'warning')
            return redirect(url_for('.login'))
        flash(u'{}登陆成功'.format(user.name), 'success')
        return redirect(url_for('main.index'))
    return render_template('auth/login.html', form=form)

直接用 session 我个人认为没问题。代码没有实际测试，以上供参考。因为没有相关经验，暂时想到这些，如果有其他的想法再来补充。

1 个赞
3

谢谢您的回答，我已经根据回答修改了一部分代码。

我之所以给每个不同的角色设立单独的（stu_id, tea_id, adm_id），是因为每个人都有不同的学号（或工号）我想让其通过学号登录，而数据库的主键ID是依次增加的，不符合每个人学号的规律。

所以我就中和了一下，在BaseUser 中加上了额外了 uid 列，登录视图如下：

@bp.route('/login', methods=['GET', 'POST'])
def login():
    if current_user.is_authenticated:
        return redirect(url_for('main.index'))
    form = LoginForm()
    if form.validate_on_submit():
        uid = form.id.data
        role = form.role.data
        # user = role_mapping[role].query.get_or_404(id)
        user = role_mapping[role].query.filter_by(uid=uid).first_or_404()
        if user is not None and user.check_password(form.password.data):
            login_user(user, remember=form.remember.data)
            session['role'] = role
        else:
            flash(u'用户名或密码错误', 'warning')
            return redirect(url_for('.login'))
        flash(u'{}登陆成功'.format(user.name), 'success')
        return redirect(url_for('main.index'))
    return render_template('auth/login.html', form=form)

PS：催更一下 api :stuck_out_tongue_closed_eyes:

1 个赞
4

这样也很好。API 书在努力写作中 : P

1 个赞