Django/Flask Custom Commands

Custom commands or what many have come to know as management commands or utility commands are one of the essential features that are provided by python frameworks. These commands are quite useful when performing tasks that requires a lot of different methods to complete the task.

Here are some issues I’ve encountered during the course of my work and guess what, i solved them by implementing some of these commands.

• If want to call a third party API and save its data into database
• If want to use a pre-saved database in another environment
• Testing and validating incoming data
• Running pre-made test cases
• Creation of dummy data for testing

That’s not all. There are lots of other issues these commands help me to fix, especially in cases where where I have to access to a function without writing a route for it

The scenario above are some of the issues I encountered during the course of my work and which i was able to fix using Custom commands. The exciting thing about using custom commands it can be expanded and use to accomplish a lot of different functionality. This is because these commands in addition to these frameworks helps create a better environment for developing and testing.

Below, let’s look at a clear example of the implementation of custom commands using one of python’s popular framework, FLASK. For today’s post, i have actually considered introducing you to two of the most famous python frameworks to work with and they are Flask and Django

So if you’re looking to use Flask to implement the custom commands, there are lots of ways to do this. This is because Flask is not a structured framework as compared to Django which has some predefined files and code structure

For the first implementation, you’ll need to have a simple fask app, to this end, you’ll have to install flask.

# create a virtual environment
# install flask
-> pip install Flask

import click
from flask import Flask

app = Flask(__name__)

@app.cli.command("create-user")
@click.argument("name")
def create_user(name):
    print("WORKING for {}".format(name))

In a simple python file, the above code will help you to export the flask app

export FLASK_APP=#name_of_the_pythonfile

#then running the command 

flask create-user admin

This particular code is uniquely designed to execute a command and perform a specific function. Let’s break down some essential point of above code shall we?

#this code will create a command with particular name and register it with our app
@app.cli.command("create-user")

#this code will allow us to pass a argument in the command if needed 
@click.argument("name")

Another implementation that works perfectly is using the blueprints of flask and guess what, they are easy to implement.

Because you are using Flask as framework for your code structure, you’ll not only be working with one file, but with multiple files with complex structures

In one of my examples, I had a JSON of my database collection in mongo and i needed to deploy it in different environments. One of the first option i had was to use a collection dump to import data, however, this had its disadvantage as my data will change every time i make a dump file, not just that, importing can be a lot of work. So what i did was create a file which included my collection JSON so i can easily edit it to meet my needs.

After this, i had to create a command which will deploy all that data in a particular collection. This means i can run the command irrespective of what environment the application is on.

So, to implement it, the command is left only. And while the flask structure is still complex, it is simple as all the flask apps are run from one simple command

flask run

Don’t forget that you’ll have to give flask app location before running the above command. This command is depicted as create_app() function and it is written to kick start your flask application

Back to my structure, I was using blueprint and application factory to structure my code base and make my app. This is a common Flask concepts

def create_app(test_config=None):
    # create and configure the app
    app = Flask(__name__, instance_relative_config=True,static_url_path='')
    app.config.from_mapping()

    CORS(app)
    APP_ROOT = os.path.join(os.path.dirname(__file__), '..')
    dotenv_path = os.path.join(APP_ROOT, '.env')
    load_dotenv(dotenv_path)
    app.config['ENV'] = os.getenv('ENVIRONMENT')

    if test_config is None:
        # load the instance config, if it exists, when not testing
        app.config.from_pyfile('config.py', silent=True)
    else:
        # load the test config if passed in
        app.config.from_mapping(test_config)
    # ensure the instance folder exists
    try:
        os.makedirs(app.instance_path)
    except OSError:
        pass

    @app.errorhandler(400)
    def not_found(error):
        return make_response(jsonify(error='Not found'), 400)

    @app.errorhandler(500)
    def error_500(error):
        return make_response({}, 500)

    db.get_db(mongo=mongo, app=app)
    
    from app.api import notify
    ......

    app.register_blueprint(notify.bp)
    ......

    return app

This was my app structure to run my flask app. What was initially needed was to create a command to fill my database for requirement

Now for that let’s see a new concept Flask-cli as we have our app defined in a different function and it has config set and other things like database initiation and it is not possible to use these with just Flask current app function so we used Flask-cli appcontext which allows us to use our created app and it’s functionality reference around these commands

If wants to learn more about Flask-cli, please do well to check the official document.
That said, let us move to the code part

Note: This approach of appcontext is used majorly when your Flask app has application factory-like structure while also using Flask concepts like Blueprint etc

First import appcontext from flask-cli

from flask.cli import with_appcontext

Now let’s make a simple command where we will fill our database with some pre-made JSON data

@click.command("seed_db")
@with_appcontext
def seed_db():
    mail_template = mongo.db.mail_template.insert_many(templates)
    mail_variable_exist = mongo.db.mail_variables.insert_many(variables)

Let’s break the above down so you understand exactly what i’m trying to say. For this example, I have used MongoDB for my database and here we are using app reference so I made a command directly with @click, with this, there will be no need to create command from the app, rather we are using app reference here. Moving on, we have used @with_appcontext for the reason I have explained above. The last part of the code is simply inserting a JSON Array into particular collections.

Now register the created command with app

app.cli.add_command(seed_db)

Here is exactly how the structure of the file will look

def create_app(test_config=None):
    # create and configure the app
    app = Flask(__name__, instance_relative_config=True,static_url_path='')
    app.config.from_mapping()

    CORS(app)
    APP_ROOT = os.path.join(os.path.dirname(__file__), '..')
    dotenv_path = os.path.join(APP_ROOT, '.env')
    load_dotenv(dotenv_path)
    app.config['ENV'] = os.getenv('ENVIRONMENT')

    if test_config is None:
        # load the instance config, if it exists, when not testing
        app.config.from_pyfile('config.py', silent=True)
    else:
        # load the test config if passed in
        app.config.from_mapping(test_config)
    # ensure the instance folder exists
    try:
        os.makedirs(app.instance_path)
    except OSError:
        pass

    @app.errorhandler(400)
    def not_found(error):
        return make_response(jsonify(error='Not found'), 400)

    @app.errorhandler(500)
    def error_500(error):
        return make_response({}, 500)

    db.get_db(mongo=mongo, app=app)
    
    from app.api import notify
    ......

    app.register_blueprint(notify.bp)
    ......
    app.cli.add_command(seed_db)

    return app

@click.command("seed_db")
@with_appcontext
def seed_db():
    mail_template = mongo.db.mail_template.insert_many(templates)
    mail_variable_exist = mongo.db.mail_variables.insert_many(variables)

Note: I have declared the command just below the app for this example, even though it’s better to write different commands in different files and just register in the app for better code structure and understanding.

With this, we are able to make a custom command with flask application factory structure. Now below are some point/mistakes which I made during the initial implementation of custom commands

• I was trying to register the command outside the flask app which will never work besides the operation function you want to use. So it’s always important to define it inside the app as it won’t work outside the flask app context
• Also, I was trying to return the app with the command, which should also not to be done as the app should perform just the function you have created it for. This is because, when you set the flask app, it is expected that the app would be returned from there,2 instances of an app will not work.
• I also tried the current_app concept which allows us to access app config and other data we initiated at the time of creating the app. This again is wrong as the app is not set to run at that time. To this end, it won’t work as the current_app is set up when we run the command
  flask run
  Note: The current_app is also part of the Flask context concept. This is useful because it is designed to help you access routes views functions etc. Again, it is not possible to do this with Flask Blueprint and application factory concepts, this is why they created application context which will hold these essential values.
  Consequently, you’ll have to use app_context as it uses the context at the time when we set the flask app

Now we will see Django Custom command.

Django Custom command: Django comes with a variety of command-line utilities that can be worked using Django-admin.py or by using manage.py script. One good thing about this is that you can also add your own custom commands. More so, it can serve as an interface to execute cron jobs. In this tutorial, you are going to learn how to code your own commands.

Introduction of Django custom commands:

In Django there are some pre-built commands like startprojects, runserver or collectstatic, I’m sure you are already familiar with this.

The name of the command file will be used to invoke using the command line utility. For example, if our command was called my_custom_command.py, then we will be able to execute it by using this custom command.

 python manage.py my_custom_command 

Let’s explore this using a simple example.

Below, is a basic example of what the custom command should look like:

What you are doing below is simply generating an otp with the help of a in-built random function

custom_command /OTP.py

from django.core.management.base import BaseCommand
import random

class Command(BaseCommand):
    help = 'custom command'

    def handle(self, *args, **kwargs):
        output = random.randint(10000000,90000000)
        self.stdout.write("OTP is", output)

Django custom command is composed of a class named Command which is inherited from BaseCommand. The command code should be defined inside the handle() function.

Here, we will run our custom command by using “ python manage.py OTP“

 Output:   OTP is 87353589 

Now we have a better understanding of what this is and how it is different from a regular Python script and its benefits.

Handling Arguments:

Django makes use of the argparse, which is part of Python’s standard library. To handle arguments in our custom command, we will have to define a method named add_arguments.

Positional Arguments: The next example is a command that creates random user instances. It takes a mandatory argument named total, which will define the number of users that will be created by the command.

custom_command /create_users.py

from django.contrib.auth.models import User
from django.core.management.base import BaseCommand
from django.utils.crypto import get_random_string

classCommand(BaseCommand):
    help = 'Create random users'

    defadd_arguments(self, parser):
        parser.add_argument('total', type=int, help='Indicates the number of users to be created')

    defhandle(self, *args, **kwargs):
        total = kwargs['total']
        for i in range(total):
            User.objects.create_user(username=get_random_string(), password='123')

we will run custom command like: python manage.py create_users 10

This command will create 10 random string users.

Optional Arguments:

The optional (and named) arguments can be passed in any order. In the example below you will find the argument named “prefix”, which will be used to compose the username field.

custom_command /create_users.py

from django.contrib.auth.models import User
from django.core.management.base import BaseCommand
from django.utils.crypto import get_random_string

classCommand(BaseCommand):
    help = 'Create random users'

    def add_arguments(self, parser):
        parser.add_argument('total', type=int, help='Indicates the number of users to be created')

        parser.add_argument('-p', '--prefix', type=str, help='Define a username prefix')

    def handle(self, *args, **kwargs):
        total = kwargs['total']
        prefix = kwargs['prefix']

        for i in range(total):
            if prefix:
                username = '{prefix}_{random_string}'.format(prefix=prefix, random_string=get_random_string())
            else:
                username = get_random_string()
            User.objects.create_user(username=username, email='', password='123')

we run custom command in this case like:
 python manage.py create_users 10 --prefix custom_user 
 python manage.py create_users 10 -p custom_user 

If the prefix is used, the username field will be created as custom_user_oYwxtt4vNHR. If not, it will be created simply as  oYwxtt4vNHR – a random string.

Flag Arguments: Another type of optional arguments are flags and they are used to handle boolean values. This is super important when we want to add an -admin flag to instruct our command to create a superuser or to create a regular user if the flag is not present.

custom_command /create_users.py

from django.contrib.auth.models import User
from django.core.management.base import BaseCommand
from django.utils.crypto import get_random_string

classCommand(BaseCommand):
    help = 'Create random users'

    defadd_arguments(self, parser):
        parser.add_argument('total', type=int, help='Indicates the number of users to be created')
        parser.add_argument('-p', '--prefix', type=str, help='Define a username prefix')
        parser.add_argument('-a', '--admin', action='store_true', help='Create an admin account')

    defhandle(self, *args, **kwargs):
        total = kwargs['total']
        prefix = kwargs['prefix']
        admin = kwargs['admin']

        for i in range(total):
            if prefix:
                username = '{prefix}_{random_string}'.format(prefix=prefix, random_string=get_random_string())
            else:
                username = get_random_string()

            if admin:
                User.objects.create_superuser(username=username, email='', password='123')
            else:
                User.objects.create_user(username=username, email='', password='123')

we will run these custom commands like:
 python manage.py create_users 2 --admin 
 python manage.py create_users 2 -a 

This will create a user with admin flag.

Arbitrary List of Arguments: Now we will create a new command named delete_users. In this new command, we will be able to pass a list of user ids and the command will delete those users from the database.

custom_command/delete_users.py

from django.contrib.auth.models import User
from django.core.management.base import BaseCommand

classCommand(BaseCommand):
    help = 'Delete users'

    defadd_arguments(self, parser):
        parser.add_argument('user_id', nargs='+', type=int, help='User ID')

    defhandle(self, *args, **kwargs):
        users_ids = kwargs['user_id']

        for user_id in users_ids:
            try:
                user = User.objects.get(pk=user_id)
                user.delete()
                self.stdout.write('User "%s (%s)" deleted with success!' % (user.username, user_id))
            except User.DoesNotExist:
                self.stdout.write('User with id "%s" does not exist.' % user_id)

we will run this by  python manage.py delete_users 1 

We can also pass a number of ids separated by spaces, so the command will delete the users in a single call:

python manage.py delete_users 1 2 3 4

Output:

User with id "1" does not exist.
User "9teHR4Y7Bz4q (2)" deleted with success!
User "ABdSgmBtfO2t (3)" deleted with success!
User "BsDxOO8Uxgvo (4)" deleted with success!

Thanks to all of the exciting information in today’s post, we now know the overall utility of custom command in python especially when working with both Django and Flask frameworks.