Implementing Flask-Login with multiple User Classes

Question

I am writing an app that has multiple classes that function as Users (for example, a School Account and a Staff account). I'm trying to use Flask-Login to make this easy but I'm not quite sure how to make it, so that when a user logs in I can have my app check to see whether or not the username belongs to a School account or Staff account, and then log in appropriately.

I know how to figure out which type of account it belongs to (since all usernames must be unique). But after that I'm not sure how to tell the app that I want it to login that specific user.

Right now, I only have one universal login page. Is it easier if I make separate login pages for Staff accounts and School accounts? I'm using a MySQL database through Flask-SQLAlchemy.

Answer 1

You can define each User with a specific role. For example, user 'x' can be SCHOOL while user 'y' can be 'STAFF'.

class User(db.Model):      __tablename__ = 'User'     id = db.Column(db.Integer,primary_key=True)     username = db.Column(db.String(80),unique=True)     pwd_hash = db.Column(db.String(200))     email = db.Column(db.String(256),unique=True)     is_active = db.Column(db.Boolean,default=False)     urole = db.Column(db.String(80))       def __init__(self,username,pwd_hash,email,is_active,urole):             self.username = username             self.pwd_hash = pwd_hash             self.email = email             self.is_active = is_active             self.urole = urole      def get_id(self):             return self.id     def is_active(self):             return self.is_active     def activate_user(self):             self.is_active = True              def get_username(self):             return self.username     def get_urole(self):             return self.urole 

Flask-login however does not have the concept of user roles yet and I wrote my own version of login_required decorator to override that. So you might want to use something like:

def login_required(role="ANY"):     def wrapper(fn):         @wraps(fn)         def decorated_view(*args, **kwargs):              if not current_user.is_authenticated():                return current_app.login_manager.unauthorized()             urole = current_app.login_manager.reload_user().get_urole()             if ( (urole != role) and (role != "ANY")):                 return current_app.login_manager.unauthorized()                   return fn(*args, **kwargs)         return decorated_view     return wrapper 

Then, you can use this decorator on a view function like:

@app.route('/school/') @login_required(role="SCHOOL") def restricted_view_for_school():     pass 

Answer 2

@codegeek i found this very useful, thanks. I had to modify the code a bit to get it working for me, so i figured i'd drop it here in case it can help anyone else:

from functools import wraps

login_manager = LoginManager()

...

def login_required(role="ANY"):
    def wrapper(fn):
        @wraps(fn)
        def decorated_view(*args, **kwargs):
            if not current_user.is_authenticated():
              return login_manager.unauthorized()
            if ((current_user.role != role) and (role != "ANY")):
                return login_manager.unauthorized()
            return fn(*args, **kwargs)
        return decorated_view
    return wrapper