How to use multiple User Classes in flask login?

Question

I am creating a Flask app in which 2 types of users can register and log in. In the @login.user_loader function, how can I configure it to check different tables based on the type of user that is logged in?

Here's my models.py code:-

    from app import db, login
from werkzeug.security import generate_password_hash, check_password_hash
from flask_login import UserMixin

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

class Patient(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    full_name = db.Column(db.String(64), index=True)
    city = db.Column(db.String(20))
    email = db.Column(db.String(120), index=True, unique=True)
    password_hash = db.Column(db.String(120))

    def __repr__(self):
        return '<Patient {}>'.format(self.full_name)

    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 Doctor(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    full_name = db.Column(db.String(64), index=True)
    city = db.Column(db.String(20))
    qual = db.Column(db.String(20))
    fees = db.Column(db.Integer)
    phone = db.Column(db.Integer)
    address = db.Column(db.String(120))
    email = db.Column(db.String(120), index=True, unique=True)
    password_hash = db.Column(db.String(120))

    def __repr__(self):
        return '<Doctor {}>'.format(self.full_name)

    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)

Currently, it is only configured for the Patient table. How to set it in such a way that whenever a Doctor is logged in, it extracts information from Doctor's table?

P.S:- I do have a choice input which takes the type of user that is logging in on the Login page. Can I use that choice data somewhere in the user_loader function?

Answer 1

Try using Joined Table Inheritance - create a single User class from which the Doctor and Patient classes inherit from. This way they share common properties (like email and password) but have their own distinct properties too. This type of polymorphism in SQLAlchemy uses independent tables for each class, but there are other examples on the page that use different methods.

An example:

class User(UserMixin, db.Model):
    email = ...
    password_hash = ...

    __mapper_args__ = {
        'polymorphic_identity':'user',
        'polymorphic_on':type
    }

class Doctor(User):
    fees = ...

    __mapper_args__ = {
        'polymorphic_identity':'user'
    }

Then your user_loader function can look like this:

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