Build a blog with Django: Add a custom management command to seed posts

Let's make it really easy to generate seed data for posts by writing a custom management command to do the job for us.

Previously, we had to go into the shell and do the following:

>>> from posts.models import Post
>>> Post.objects.all().delete()

>>> from posts.seed import create_posts
>>> create_posts(1, 5)

But, as I mentioned at the end of that post, it would be so much simpler if we could just do python manage.py createposts --unpublished 1 --published 5 instead.

In this post, I'll show you how.

Let's get started.

Add a createposts management command

As usual, Django has excellent documentation on writing custom commands. So I'll just give you a quick overview and refer you to the docs for further details.

Applications can create their own commands to use with manage.py. To do so simply add a management/commands directory to the application and Django will automatically register a manage.py command for each Python module in that directory whose name doesn't begin with an underscore.

Hence, for our createposts command we will have to create a createposts.py module in the management/commands directory like so:

posts  
├── ...
├── management
│   ├── commands
│   │   ├── createposts.py
│   │   └── __init__.py
│   └── __init__.py
│...

If you now do the following:

(venv) $ cd src && python manage.py

You'd see your new command listed among all the other possible commands:

Type 'manage.py help <subcommand>' for help on a specific subcommand.

Available subcommands:

[auth]
    changepassword
    createsuperuser

[django]
    check
    compilemessages
    createcachetable
    dbshell
    diffsettings
    dumpdata
    flush
    inspectdb
    loaddata
    makemessages
    makemigrations
    migrate
    sendtestemail
    shell
    showmigrations
    sqlflush
    sqlmigrate
    sqlsequencereset
    squashmigrations
    startapp
    startproject
    test
    testserver

[posts]
    createposts

[sessions]
    clearsessions

[staticfiles]
    collectstatic
    findstatic
    runserver

Now we just need to put some code into the module for our specific needs.

Django only has one requirement when it comes to the contents of the module. It MUST have a class Command that extends BaseCommand or one of its subclasses.

Here's the structure of the module:

from django.core.management.base import BaseCommand

# ...

class Command(BaseCommand):  
    help = '...'

    def add_arguments(self, parser):
        # ...

    def handle(self, *args, **options):
        # ...
  1. The help class attribute should contain a short description of the command. It would be printed in the help message when the user runs the command python manage.py help createposts. (see docs)

  2. You override the add_arguments method to add both positional and optional arguments that will be accepted by the command. (see docs)

  3. Subclasses MUST implement the handle method in order to provide the actual logic of the command. (see docs)

Let's fill it all out.

Add a help message

This one is simple. Just add the following:

help = 'Seeds the database with a given number of unpublished and published posts.'  

Add optional arguments

The parser argument that's passed to the add_arguments method is an instance of CommandParser. But, CommandParser is just a subclass of ArgumentParser. Hence, to understand how to work with the parser argument it's helpful to read up on Python's argparse module.

We'll be needing to use the add_argument method of parser.

Here's what we want to do: We want to add two optional arguments called unpublished and published. They both take integer arguments and default to 0 when they aren't provided.

And, we say that as follows:

def add_arguments(self, parser):  
        parser.add_argument('--unpublished',
            default=0,
            type=int,
            help='The number of unpublished posts to create.')

        parser.add_argument('--published',
            default=0,
            type=int,
            help='The number of published posts to create.')

Add the actual logic of the command

Finally, we need to add the logic to do the actual work.

Here's what we want to happen: Firstly, ask the user if they do want to continue with the operation since we'd be deleting all posts before adding the new ones. If they say yes then we want to delete the posts and create the new posts. Otherwise, we ignore the operation.

And, here it is in code:

from django.db import transaction

from posts.models import Post  
from posts.seed import create_posts

# ...

def handle(self, *args, **options):  
        confirm = input("""This operation will IRREVERSIBLY DESTROY all posts currently in the database.
Are you sure you want to do this?

Type 'yes' to continue, or 'no' to cancel: """)

        if confirm == 'yes':
            with transaction.atomic():
                Post.objects.all().delete()

                unpublished = max(0, options['unpublished'])
                published = max(0, options['published'])
                create_posts(unpublished, published)
        else:
            self.stdout.write('Operation cancelled.\n')

Pretty straightforward, except maybe for transaction.atomic.

You see, we want all the changes we're making to the database to happen in one transaction. A transaction is a sequence of operations performed as a single logical unit of work.

Suppose that after deleting all the posts from the database the program crashes due to a bug in create_posts, i.e. all the previous posts are gone but no new posts have been created. That sucks for the user (though it shouldn't be so bad in this case since it's probably all fake data anyway). With the operations wrapped in a transaction that will never happen since any database changes will be rolled back in the event of a failure, i.e. if create_posts crashes then the deleted posts are "restored" in the database. The user is notified of the error (in this case via a stack trace on the command-line) and the database looks as it did before running the command.

I suggest you try out the command and see for yourself how it all works.

Wrap up

We're almost done with the overall "add posts" feature.

Completed adding a custom management command to seed posts

Here's the commit that tracks what we did today.

See you next time when I show you how to handle the display of published posts. We'd be touching on managers, get_absolute_url, URL patterns, views and templates.

I hope you enjoy seeing how all the features of Django come together to provide a seamless web development experience.

P.S. Let me know in the comments below about your experience using custom management commands. Do you like them? Any tips? I'd love to get your feedback.

P.S.S. Subscribe to my newsletter if you're interested in getting exclusive Django content.