How can I use flask_login with multiple user classes and different database entries?

Question

If someone were to have a flask web app with several "types" of user, for example in a school setting there may be "teacher", "student" and "admin". Each user needs to have a way to log in to the app, but Flask-Login only provides one @login.user_loader. My question is similar to this question, about multiple user classes, but the difference is that each type of user does not have the same model structure as the others. For example, if one were to use flask-sqlalchemy and declared the following models:

class Teacher(db.Model, UserMixin):
    id = db.Column(db.Integer,primary_key=True)
    title = db.Column(db.String(32))
    pwd_hash = db.Column(db.String(128))
    first_name = db.Column(db.String(128))
    last_name = db.Column(db.String(128))
    students = db.relationship('Student',backref='teacher',lazy='dynamic') #because a teacher will teach multiple students
    
class Student(db.Model, UserMixin):
    id = db.Column(db.Integer,primary_key=True)
    username = db.Column(db.String(64),unique=True)
    teacher_id = db.Column(db.ForeignKey('teacher.id'),db.Integer) #for the ease of things, the student is only ever taught by one teacher

class Admin(db.Model, UserMixin):
    email = db.Column(db.String(128),unique=True)
    name = db.Column(db.String(128))

The user_loader is normally used like this, with the User class here being a generic user class:

@login.user_loader
def load_user(id):
    return User.query.get(int(id))

This means that you have to have just one User class that can be logged in. The normal approach would be to condense the Student, Teacher and Admin class into one User class, with a role field in the model which tells you what level of access they are. But this brings up multiple problems. For example, not all students have an email at their disposal yet the administrative account for the school needs one. Similarly, the relationship between students and teachers would fail because there is no model for the teachers to have students in their class etc.

So how can I implement a way of there being multiple user models with different attributes and model fields that can each be logged in separately with flask_login?

Answer 1

Your User needs only an id and a relationship to the tables which that user corresponds to. You can create this by creating a table for roles and roles-relationships:

class User(db.Model, UserMixin):
    """"""
    __tablename__ = "user"
    # Core
    id = db.Column(db.Integer,primary_key=True)
    roles = db.relationship('Role', secondary=roles_users,
                            backref=db.backref('users', lazy='dynamic'))

class Role(db.Model, RoleMixin):
    __tablename__ = "role"
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(80), unique=True)

roles_users = db.Table('roles_users',
        db.Column('user_id', db.Integer(), db.ForeignKey('user.id')),
        db.Column('role_id', db.Integer(), db.ForeignKey('role.id')))

You can check for a given role using the following:

class User(db.Model, UserMixin):
    # all the args go here...
    ...
    def has_role(self, role_name):
        """Does this user have this permission?"""
        my_role = Role.query.filter_by(name=role_name).first()
        if my_role in self.roles:
            return True
        else:
            return False

Next, create a decorator which checks for a role

from flask import redirect
from flask_login import current_user
from functools import wraps


def require_role(role):
    """make sure user has this role"""
    def decorator(func):
        @wraps(func)
        def wrapped_function(*args, **kwargs):
            if not current_user.has_role(role):
                return redirect("/")
            else:
                return func(*args, **kwargs)
        return wrapped_function
    return decorator

Step 3: Apply the decorator in operations:

@app.route("/")
@login_required
@require_role(role="teacher")
def teachers():
    """Only allow teachers"""
    # Otherwise, proceed
    return render_template('teachers.html')