Flask Login with Multiple User Classes and SSO

user_11036104 0
Viewed 744

Flask beginner here. I have three user classes for my flask-login. GoogleUser and FacebookUser utilize SSO and does not require the password field:

from InstaFlix import db, login_manager
from datetime import datetime
from flask_login import UserMixin

# Get a user by their ID
@login_manager.user_loader
def load_user(user_id):
    return User.query.get(user_id)

class User(db.Model, UserMixin):
    id = db.Column(db.Integer, primary_key=True) # primary keys are required by SQLAlchemy
    email = db.Column(db.String(100), unique =True, nullable=False)
    password = db.Column(db.String(60), nullable=False)
    name = db.Column(db.String(100), nullable=True)
    username = db.Column(db.String(20), unique=True, nullable=False)
    profile_pic = db.Column(db.String(100), unique=False, nullable=False, default='default.jpg')

class GoogleUser(db.Model, UserMixin):
    id = db.Column(db.String(20), primary_key=True) 
    email = db.Column(db.String(100), unique =True, nullable=False)
    name = db.Column(db.String(40), unique=True, nullable=False)
    profile_pic = db.Column(db.String(100), unique=False, nullable=True)

class FacebookUser(db.Model, UserMixin):
    id = db.Column(db.String(20), primary_key=True) 
    email = db.Column(db.String(100), unique =True, nullable=False)
    name = db.Column(db.String(40), unique=True, nullable=False)
    profile_pic = db.Column(db.String(100), unique=False, nullable=True)

I want to load a specific class depending on which route the user goes too. For example https://127.0.0.1:5000/fb would allow the user to log in. Then call the class FacebookUser to save their information into.

However, I can only return one class at a time.

How do you a call a specific user class with the load_user() function, depending on the route?

As of right now, only one class can be loaded depending on which one I query.

flask-login flask-sqlalchemy flask python
1 Answers

IMO, there is no need to create multiple models for users with different type/role, you can create a new column to store the user type. Besides, the columns in your GoogleUser and FacebookUser model are completely duplicate with the User model. I would recommend just use one User class for them. It's more elegant and convenient.

If you want to distinguish different user type, just add a new column like this:

class User(db.Model, UserMixin):
    # ...
    login_type = db.Column(db.String(20))

If you really need to use multiple user models, you can use session to store the login type, for example:

from flask import session

@oauth_bp.route('/callback/<provider_name>')
def oauth_callback(provider_name):
    if provider_name not in providers.keys():
        abort(404)

    provider = providers[provider_name]
    response = provider.authorized_response()
    session['login_type'] = provider_name

Then read it in load_user, use different user model depend on its value:

from flask import session

@login_manager.user_loader
def load_user(user_id):
    login_type = session.get('login_type')
    if login_type == 'google':
        return GoogleUser.get(user_id)
    elif login_type == 'facebook':
        return FacebookUser.get(user_id)
    else:
        return User.query.get(user_id)
user_55118490
Related