Django recipe: Base management command for running custom SQL

By •

Here is a simple Django management command I use when I want to quickly execute custom SQL in the database.

For me, it comes in handy in cases working with large data sets where Django's tools for interacting with the database can take more time to craft and execute than a well-tailored INSERT command.

Here's the command. Paste it wherever you'd like in your project.

from django.db import transaction, connection
from django.core.management.base import BaseCommand, CommandError


class SimpleSQLCommand(BaseCommand):
    help = "A base class for packaging simple SQL operations as a command"
    # Overriding these attributes is what you will need to do when subclassing
    # this command for use.
    flush = None  # An optional Django database model to be flushed
    sql = ""  # The SQL command to be run

    def handle(self, *args, **options):
        # Validate
        if not self.sql:
            raise CommandError("'sql' attribute must be set")

        # Flush model if it is provided
        if self.flush:
            if options.get("verbosity") >= 1:
                self.stdout.write("- Flushing %s" % self.flush.__name__)
            self.flush_model(self.flush)

        # Run custom sql
        if options.get("verbosity") >= 1:
            self.stdout.write("- Running custom SQL")
        self.execute_sql(self.sql)

    @transaction.atomic
    def flush_model(self, model):
        """
        Flushes the provided model using the lower-level TRUNCATE SQL command.
        """
        cursor = connection.cursor()
        cursor.execute("TRUNCATE %s CASCADE;" % (model._meta.db_table))

    @transaction.atomic
    def execute_sql(self, sql):
        """
        Executes the provided SQL command.
        """
        cursor = connection.cursor()
        cursor.execute(sql)

Then simply import it into wherever you're putting your new custom command and provide the sql attribute. If you want to flush a model prior to running the command, as I often do, you can also provide it to the optional model attribute.

from myapp import models
from mytoolbox import SimpleSQLCommand


class Command(SimpleSQLCommand):
    flush = models.MyModel
    # I might typically have more JOINs and other SQL crap going on
    # but let's keep it simple for this example
    sql = """
        INSERT INTO myapp_mymodel ("group","count")
        SELECT group, COUNT(*)
        FROM myapp_myothermodel
        GROUP BY group
        """

Then, provided you've put your new command in the right spot, running it should be as simple as this.

$ python manage.py mynewcommand
en
333