Creating Somewhat Middleware using Decorators in Flask

For a while now I’ve been doing some study on constructing and consuming microservices using Flask and a combination of tools. This has led to a lot of discoveries along the way and improvements overall in the way I write Python code.

In my most recent experiment, while trying to create a user service, I found a nifty way to use decorators to intercept requests to routes especially for performing generic actions like authentication, error handling, blanket resource loading (loading objects in request object for a group of routes/URL patterns) etc. before the main controller logic runs.

In this write-up, I will be assuming you have adequate knowledge of Flask and how it works. I will also be making use of my goto Flask wrapper of choice, [Flask-RESTful](flask-restful-cn.readthedocs.org/en/0.3.4/).

 Why do I need this?

I don’t seek to impose this on you and so you may actually not need to use it. The reason for this is, coming from programming using NodeJS for a while, I find the concept of Express middleware very simple, easy to implement and an elegant way to handle operations before the actual controller logic (This is available also in Django as Middleware or by overriding the dispatch() function in CBVs). It was interesting trying to replicate a similar thing in Flask.

In Express, you could write something like this:

app.use('/user/:id', function (req, res, next) {
  console.log('Performing this operation before the main operation');
  next();
});

app.get('/user/:id', function (req, res, next) {
  res.send('Performing the GET user operation');
});

Shown above are two callback functions. The first one executes for any request to the route pattern /user/:id i.e GET, POST, PUT, DELETE, PATCH. This is the middleware and it always contains a call to next() which indicates that express should move down the chain of execution for this request. To make it a bit clearer, executing GET /user/:id will first perform the generic middleware task for the route pattern, in this case logging some text on the console, before the actual operation specified for GET is performed. Here, GET /user/:id will print two statements.

I will be showing you a way to achieve a similar thing using Flask.

 Flask UserService

Let’s start by writing out the GET endpoint for a UserService in Flask:

from flask import Flask
from flask_restful import Api, Resource, fields

from .models import User

app = Flask(__name__)
api = Api(app)

user_serializer = {
    'id': fields.Integer,
    'username': fields.String,
    'email': fields.String
}


class UserService(Resource):
    @marshal_with(user_serializer)
    def get(self, user_id):
        users = User.query.get(user_id)
        return user

api.add_resource(UserService, '/users/<int:user_id>')

This example assumes we have already created a User model. Navigating to the /users/9 route should return information about the user with id 9.

In this example, we obviously cannot see so much repetition. Let’s extend this to include PUT and DELETE operations.

from flask import Flask
from flask_restful import Api, Resource, fields, reqparse

from .models import User

app = Flask(__name__)
api = Api(app)

user_serializer = {
    'id': fields.Integer,
    'username': fields.String,
    'email': fields.String
}


class UserService(Resource):
    @marshal_with(user_serializer)
    def get(self, user_id):
        user = User.query.get(user_id)
        return user

    @marshal_with(user_serializer)
    def put(self, user_id):
        user = User.query.get(user_id)
        parser = reqparse.RequestParser()
        parser.add_argument('username', type=str)
        parser.add_argument('email', type=str)
        req = parse.parse_args()
        user.username = req.get('username', user.username)
        user.email = req.get('email', user.email)
        db.session.add(user)
        db.session.commit()
        return user

    def delete(self, user_id):
        user = User.query.get(user_id)
        db.session.delete(user)
        db.session.commit()
        return "Deleted user", 204

api.add_resource(UserService, '/users/<int:user_id>')

Some repetition is beginning to happen here. In this case, not even as much as what we’re trying to prevent but should also be sufficient to illustrate the technique. For every request here to the same URL pattern, /users/<int:user_id>, we need to get the user object and then do something to it. Using decorators, we can abstract this logic and more away.

from functools import wraps

from flask import g

from .models import User

def get_user(f):
    @wraps(f)
    def func_wrapper(*args, **kwargs):
        user_id = kwargs.get('user_id')
        user = User.query.get(user_id)
        if user is None:
            abort(404, message="User not found")
        g.user = user
        return f(*args, **kwargs)
    return func_wrapper

We have created a @get_user decorator here to try to get the user and store it in the g object before proceeding to call the function. This way the user object is available for all subsequent operations depending on the request method. As you have noticed here we’ve packed some other logic in here illustrating the benefits of using this technique. Now, an error is thrown if the user doesn’t exist.

We can proceed to decorate all our functions like so:

from functools import wraps

from flask import Flask, g
from flask_restful import Api, Resource, fields, reqparse

from .models import User

app = Flask(__name__)
api = Api(app)

user_serializer = {
    'id': fields.Integer,
    'username': fields.String,
    'email': fields.String
}

def get_user(f):
    @wraps(f)
    def func_wrapper(*args, **kwargs):
        user_id = kwargs.get('user_id')
        user = User.query.get(user_id)
        if user is None:
            abort(404, message="User not found")
        g.user = user
        return f(*args, **kwargs)
    return func_wrapper

class UserService(Resource):
    method_decorators = [get_user]

    @marshal_with(user_serializer)
    def get(self, user_id):
        return g.user

    @marshal_with(user_serializer)
    def put(self, user_id):
        user = g.user
        parser = reqparse.RequestParser()
        parser.add_argument('username', type=str)
        parser.add_argument('email', type=str)
        req = parse.parse_args()
        user.username = req.get('username', user.username)
        user.email = req.get('email', user.email)
        db.session.add(user)
        db.session.commit()
        return user

    def delete(self, user_id):
        user = g.user
        db.session.delete(user)
        db.session.commit()
        return "Deleted user", 204

api.add_resource(UserService, '/users/<int:user_id>')

Here we have made use of the method_decorators parameter of the Resource class available when using Flask-RESTful. This applies all the decorators specified in the list on all the methods defined.

As mentioned earlier, we could also extend this example to include code that ensures that the user accessing the object actually owns it for instance.

We have not done so much in terms of shortening the code but the structure here is a lot more elegant and for more complex scenarios, this might be a lot easier to maintain.

We can also try out the example of preloading some data before running the main view logic.

from .models import Profile

def get_profile_data(f):
    @wraps(f)
    def func_wrapper(*args, **kwargs):
        user_id = kwargs['user_id']
        profile = Profile.query.filter(Profile.user_id == user_id)
        if profile is None:
            abort(400, message="An error occured while trying to retrieve profile")
        g.profile = profile
        return f(*args, **kwargs):
    return func_wrapper

Here, we are able to load the user’s profile and make it accessible to subsequent function calls executed while processing the request. The benefit of this is that it can be used across multiple resources so long as the required parameters are available.

I do hope this was helpful.

 
42
Kudos
 
42
Kudos

Now read this

Demystifying Token-Based Authentication using Django REST Framework

Authentication is one of those things which have now been considered a rote and repetitive task when doing web development. Most applications you will ever develop almost always need to have some form of user authentication to allow... Continue →