Custom Django Migration Operations with Index

This feature is super hacky. On a scale of 1 to hacky, this is on the edge of “you probably shouldn’t do this.”

This feature is the ability to generate custom migration operations using the indexes attribute on a model. Then in those custom operations, you can change the database schema however you would like.

From a DX perspective, it’s the ability to do:

from django.db import models

class CustomMigrationOperation(models.Index):
    # ... to be filled out later
    pass

class MyModel(models.Model):
    field = models.TextField()

    class Meta:
        indexes = [CustomMigrationOperation()]

Which will cause the following migration to be auto-generated in manage.py makemigrations.

# `manage.py makemigrations` will generate:
class Migration(migrations.Migration):
    operations = [
        CreateModel(
            name="MyModel",
            fields=[...],
            options={"indexes": [CustomMigrationOperation()]}
        ),
    ]

With this being the SQL that is run when running manage.py migrate or shown by manage.py sqlmigrate.

BEGIN;
-- Create model MyModel
CREATE TABLE "myapp_mymodel" (
    "id" integer NOT NULL PRIMARY KEY GENERATED BY
        DEFAULT AS IDENTITY,
    "field" text NOT NULL
);

-- Custom SQL would go here!
ALTER TABLE "myapp_mymodel" AND FIELDS WITH CUSTOM OPERATION;

COMMIT;

The reason this all works out of the box is because the Index class in Django has two methods to control the SQL to manage the database schema. The methods are create_sql and remove_sql.

from django.db import models

class CustomMigrationOperation(models.Index):
    def create_sql(self, model, schema_editor, using="", **kwargs):
        # Don't do this exactly! It's not properly escaped.
        return f'ALTER TABLE "{model._meta.db_table}" AND ' \
                'FIELDS WITH CUSTOM OPERATION;'

    def remove_sql(self, model, schema_editor, **kwargs):
        return f'DROP CUSTOM OPERATION FROM "{model._meta.db_table}"'

This SQL can be just about anything. You could add table comments or set up triggers with this mechanism. The best use cases would be things related to creating or editing schema that is related to the table as a whole.

A more comprehensive example

Where I used this was in the project django-security-label to create security labels to configure anonymization rules from PostgreSQL Anonymizer. The benefit here is that I’m making the anonymization rules a first-class priority rather than hiding them in custom migrations. A developer has to determine how should a field be anonymized when they are defining the model. Then as the models change, Django’s migration system determine what needs to be added/removed to the schema automatically.

from django.db import models
from django_security_label import labels

class MaskedColumn(models.Model):
    confidential = models.TextField()

    class Meta:
        indexes = [
            labels.AnonymizeColumn(
                fields=["confidential"],
                string_literal="MASKED WITH VALUE $$CONFIDENTIAL$$",
            ),
        ]

That project was able to use the indexes attribute to generate custom SQL and manage it at the models layer to reduce developer overhead.

So yeah. Definitely hacky, but I hope it gets you thinking about how you can extend Django to do more in the database.

If you have thoughts, comments, or questions, please let me know. You can set up a meeting with me, or find me on the Fediverse, Django Discord server or use email.