flask_jwt identity_callback feature

[rayrapetyan]

Hello.

flask_jwt has a really useful feature - identity_callback. It is calling every time jwt_required decorator is triggered. In this user-defined call-back you may e.g. initialize a g.auth_user variable and use it from the protected endpoint.

In flask_jwt_extended to achieve the same I have to manually call get_jwt_identity() from every protected endpoint and\or add an additional decorator to every function, which is not very handy having hundreds of such points in the existing code.

Is there a way to trigger some functionality from the every protected endpoint automatically?

Thanks!

[vimalloc]

When originally designing this application I was against this idea, as I don't want to force user to take the database hit on every endpoint if they didn't need to. Thinking about it more now, I think we could make this work. Stream of thought incoming.

Instead of requiring the identity callback like what is done in flask-jwt, we can have a default callback which simply returns None. This doesn't add any substantial overhead or cause breaking changes for current users. Then we could add a way to override that callback to so that you could load/access whatever user objects you wanted in your endpoints.

We would need to figure out how to handle errors in that callback, for example if the database was unreachable, or if the JWT was for a user that was no longer in the database. In Flask-JWT, you return None to indicate an error, which wouldn't work for us if we are returning None already by default. We could try to do it with exceptions instead, but would need to hook it into some app.errorhandler magic. Perhaps the decorators could call the identity callback, catch all exceptions, and re-raise the exception under a new exception type that has an error handler associated with it (that could also be overridden like the other error handlers in this extension).

Any thoughts on any of this? Think this would fill your requirements?

(Adding this link, as it's a related issue #30)

[rayrapetyan]

Why not just make "requre_jwt" decorator to check if there is a callback defined (callback-pointer variable is initialized), then call it. In this case returning None from callback can be an indication of an error, like in flask-jwt, is this inappropriate? Such approach will not break anything for existing users, or I'm missing something?

[vimalloc]

That sounds reasonable. I'll try to get that setup this week.

Cheers.

[vimalloc]

I've been working on this in the user_loader branch, if it looks good to you I think it's ready to be merged.

It will look like this:

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

app = Flask(__name__)
app.secret_key = 'super-secret'  # Change this!
jwt = JWTManager(app)


# Just an example here, would likely be sqlalchemy or some
# such in a real application
class UserObject:
    def __init__(self, username, roles):
        self.username = username
        self.roles = roles

users = {
    'foo': ['admin'],
    'bar': ['peasant'],
    'baz': ['peasant']
}


@jwt.user_loader_callback_loader
def user_loader_callback(identity):
    if identity not in users:
        return None
    return UserObject(username=identity, roles=users[identity])


@app.route('/login', methods=['POST'])
def login():
    username = request.json.get('username', None)
    access_token = create_access_token(identity=username)
    ret = {'access_token': access_token}
    return jsonify(ret), 200


@app.route('/admin-only', methods=['GET'])
@jwt_required
def protected():
    # Could alternately use get_current_user() here, instead of the LocalProxy
    if 'admin' not in current_user.roles:
        return jsonify({"msg": "Forbidden"}), 403
    return jsonify({"secret_msg": "don't forget to drink your ovaltine"})

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

And it will allow you to customize the error messages when the user loader callback returns None if so desired.

@jwt.user_loader_error_loader
def custom_user_loader_error(identity):
    return jsonify({"msg": "User not found"}), 404

The default error handler will return:

jsonify({'msg': "Error loading the user {}".format(identity)}), 401

[rayrapetyan]

Looks excellent, thank you so much!

[vimalloc]

Wonderful 👍

This is released in 2.4.0, which should be hitting pypi shortly. Here it the link to the documentation. Let me know if you run into any problems with this.

Cheers!

[zedrdave]

This is a very helpful feature…

But am I missing something, or this only updates flask_jwt_extended.current_user, not flask_login.current_user?

Any reason the latter is not also set? Sounds like it might be useful for compatibility with other modules…

[vimalloc]

That’s simply not how things work. This isn’t flask login, doesn’t rely on flask login, and doesn’t provide a flask_login package or import, thus there is 0 reason why it should be setting a flask_login value.

[zedrdave]

@vimalloc Indeed, I misread flask-jwt's doc and thought it did set flask_login (when it in fact only uses its own module, just like flask-jwt-extended). That seems to be consistent then.