Tokens from Complex Objects

A very common setup is to have your users information (usernames, passwords, roles, etc) stored in a database. Now, lets pretend that we want to create an access tokens where the tokens identity is a username, and we also want to store a users roles as an additional claim in the token. We can do this with the user_claims_loader() decorator, discussed in the previous section. However, if we pass the username to the user_claims_loader(), we would end up needing to query this user from the database two times. The first time would be when login endpoint is hit and we need to verify a username and password. The second time would be in the user_claims_loader() function, because we need to query the roles for this user. This isn’t a huge deal, but obviously it could be more efficient.

This extension provides the ability to pass any object to the create_access_token() function, which will then be passed as is to the user_claims_loader(). This allows us access the database only once, but introduces a new issue that needs to be addressed. We still need to pull the username out of the object, so that we can have the username be the identity for the new token. We have a second decorator we can use for this, user_identity_loader(), which lets you take any object passed in to create_access_token() and return a json serializable identity from that object.

Here is an example of this in action:

from flask import Flask, jsonify, request
from flask_jwt_extended import (
    JWTManager, jwt_required, create_access_token,
    get_jwt_identity, get_jwt_claims
)

app = Flask(__name__)

app.config['JWT_SECRET_KEY'] = 'super-secret'  # Change this!
jwt = JWTManager(app)


# This is an example of a complex object that we could build
# a JWT from. In practice, this will likely be something
# like a SQLAlchemy instance.
class UserObject:
    def __init__(self, username, roles):
        self.username = username
        self.roles = roles


# Create a function that will be called whenever create_access_token
# is used. It will take whatever object is passed into the
# create_access_token method, and lets us define what custom claims
# should be added to the access token.
@jwt.user_claims_loader
def add_claims_to_access_token(user):
    return {'roles': user.roles}


# Create a function that will be called whenever create_access_token
# is used. It will take whatever object is passed into the
# create_access_token method, and lets us define what the identity
# of the access token should be.
@jwt.user_identity_loader
def user_identity_lookup(user):
    return user.username


@app.route('/login', methods=['POST'])
def login():
    username = request.json.get('username', None)
    password = request.json.get('password', None)
    if username != 'test' or password != 'test':
        return jsonify({"msg": "Bad username or password"}), 401

    # Create an example UserObject
    user = UserObject(username='test', roles=['foo', 'bar'])

    # We can now pass this complex object directly to the
    # create_access_token method. This will allow us to access
    # the properties of this object in the user_claims_loader
    # function, and get the identity of this object from the
    # user_identity_loader function.
    access_token = create_access_token(identity=user)
    ret = {'access_token': access_token}
    return jsonify(ret), 200


@app.route('/protected', methods=['GET'])
@jwt_required
def protected():
    ret = {
        'current_identity': get_jwt_identity(),  # test
        'current_roles': get_jwt_claims()  # ['foo', 'bar']
    }
    return jsonify(ret), 200


if __name__ == '__main__':
    app.run()