Three weeks ago our development team was chatting. We were launching a new app and we have decided to manage pricing plans through the Django admin interface. Our problem was that we needed at least a FREE plan in the database to enable the user to sign up. The production environment was already set up and everything worked. But when we added the development environment, we suddenly couldn’t sign up. We have checked the logs and found out that none can sign up due to missing pricing plans. So the question was: How to populate the database with general pricing plans on deploy if they do not exist?

Custom Django command

After some research, we’ve decided to register a custom Django command and define pricing plans in the YAML file. After successful deploy command is run. So let’s take a look at how we can do that.

Basic set up

Create a new Django project called django_populate_db_with_command:

django-admin.py startproject django_populate_db_with_command

After that, we will create a new app called pricing:

django-admin.py startapp pricing

Now we can add our pricing plan model:

# pricing/models.py
from django.db import models


# Create your models here.
class PricingPlan(models.Model):
    name = models.CharField('Name of plan', max_length=50, unique=True)
    amount = models.DecimalField('Price for plan', max_digits=4, decimal_places=2, default=0)

Create migrations, migrate and create superuser:

python manage.py makemigrations
python manage.py migrate
python manage.py createsuperuser

We will also register the model to admin. Therefore we will be able to see if plans were created.

# pricing/admin.py
from django.contrib import admin


# Register your models here.
from pricing.models import PricingPlan


class PricingPlanAdmin(admin.ModelAdmin):
    pass


admin.site.register(PricingPlan, PricingPlanAdmin)

Don’t forget to register app in settings.py:

# ... other code ...
INSTALLED_APPS = [
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'pricing'
]
# ... other code ...

Start the development server,

python manage.py runserver

go to http://localhost:8000 and check that there are no pricing plans.

Register command

To register command into manage.py we need to create a new directory inside our pricing app called management. Inside that, we need to create a directory called commands. And inside that, we create a new python file called createpricingplans.py. The structure of pricing app should look like this:

  • pricing
    • management
      • commands
        • createpricingplans.py
    • migrations
    • __init__.py
    • admin.py
    • apps.py
    • models.py
    • tests.py
    • views.py

Our createpricingplans.py should contain the following code:

# pricing/management(commands/createpricingplans.py 
import yaml
from django.core.management.base import BaseCommand
from pricing.models import PricingPlan


class Command(BaseCommand):
    help = 'Create pricing plans from yaml'

    def add_arguments(self, parser):
        parser.add_argument('--path', type=str)

    def handle(self, *args, **options):
        with open(options['path'], 'r') as stream:
            plans = yaml.safe_load(stream)
            for plan in plans:
                pricing_plan, _ = PricingPlan.objects.get_or_create(name=plan['name'])
                pricing_plan.amount = plan['amount']
                pricing_plan.save()

        self.stdout.write(self.style.SUCCESS('Successfully created plans'))

What is done here? Inside add_arguments method, an argument –path is registered to command to tell it which YAML file to use. The method called handle actually executes the command. Firstly, it reads the YAML file from provided path. Then it tries to get the existing plan with the name defined in the YAML file. If it does not exists it creates a new one. After that, it updates it with the amount defined in the YAML file.

To successfully run the command, we need to add a YAML file. Create a new YAML file inside pricing app that looks like this:

- name: Free
  amount: 0

Now we are ready to run our command:

python manage.py createpricingplans --path pricing/plans.yaml

Open http://localhost:8000 and check that our pricing plan is added now.

The code can be found here on Github.

Source: https://docs.djangoproject.com/en/3.0/howto/custom-management-commands/

That’s all – happy coding!

Categories: Django