Deferable unique constraints in django (for postgresql) by creating a specific django db backend

__init__.py

"""A DB backend, based on ``django.db.backends.postgresql_psycopg2`` to cover a specific needs.

The need is to add the "DEFERRABLE INITIALLY DEFERRED" when defining a unique constraint
on `sort_order` fields, to make the ``Orderable`` base model works correctly.

"""

base.py

"""Defines the ``DatabaseWrapper`` class of the ``mtp_db_backend``"""


from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper as BaseDatabaseWrapper

from .schema import DatabaseSchemaEditor


class DatabaseWrapper(BaseDatabaseWrapper):
    """Simple override to use our own schema editor."""

    SchemaEditorClass = DatabaseSchemaEditor

schema.py

"""Defines the ``DatabaseSchemaEditor`` class of the ``mtp_db_backend``"""

from django.db.backends.postgresql_psycopg2.schema import (
    DatabaseSchemaEditor as BaseDatabaseSchemaEditor
)


class SqlUniqueDeferrableForSortOrder(str):
    """String with interpolation interception to manage "sort_order" constraints.

    The goal is, for this field, to have the following SQL:

        ``UNIQUE (%(columns)s) DEFERRABLE INITIALLY DEFERRED``

    instead of simply

        ``UNIQUE (%(columns)s)``

    """

    deferrable_sql = '%(default_sql)s DEFERRABLE INITIALLY DEFERRED'

    def __mod__(self, data):
        """Override string interpolation to intercept `sort_order` and update the constraint.

        First the default interpolation is done to get the normal "unique" sql constraint.
        Then we check if the ``columns`` entry of the ``data`` dict contains one named "sort_order"
        and if True, we use ``deferrable_sql`` to compose the final sql.

        """

        # Get the default sql
        result = super().__mod__(data)

        # We need to read the colums to check for "sort_order"
        # Columns is a string with column name separated by comma, all part being encapsulated in "
        # Example: '"sort_order", "training_id"'
        columns = data['columns'][1:-1].split('", "')

        # If "sort_order" is in the columns, we add the "DEFERRABLE INITIALLY DEFERRED" part
        if 'sort_order' in columns:
            result = self.deferrable_sql % {
                'default_sql': result
            }

        # And we return the final sql
        return result


class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
    """Override of the default shema editor to manage unique constraints for the "sort_order" field.

    For this we use ``SqlUniqueDeferrableForSortOrder`` that will intercept the interpolation and
    if "sort_order" is in the given columm, will add the "DEFERRABLE INITIALLY DEFERRED" arguments
    to the constraint.

    It's applied to two sql strings:
    - ``sql_create_table_unique``: used when creating a table from scratch, without migration
    - ``sql_create_unique``: used when adding a constraint to an existing table

    """

    sql_create_table_unique = SqlUniqueDeferrableForSortOrder(
        BaseDatabaseSchemaEditor.sql_create_table_unique
    )

    sql_create_unique = SqlUniqueDeferrableForSortOrder(
        BaseDatabaseSchemaEditor.sql_create_unique
    )

tests.py

from unittest import TestCase

from .schema import DatabaseSchemaEditor


class SqlUniqueDeferrableForSortOrderTestCase(TestCase):

    @staticmethod
    def quote_name(column):
        """Simulate the ``quote_name`` from the default schema editor."""
        return '"%s"' % column

    def test_default_interpolation_for_create_table_unique_with_one_column(self):
        """Test that the it returns the default result with one column (not "sort_order")."""

        sql = DatabaseSchemaEditor.sql_create_table_unique % {
            'columns': ", ".join(self.quote_name(column) for column in ['foo']),
        }
        self.assertEqual(sql, 'UNIQUE ("foo")')

    def test_default_interpolation_for_create_table_unique_with_many_columns(self):
        """Test that the it returns the default result with many columns (not "sort_order")."""

        sql = DatabaseSchemaEditor.sql_create_table_unique % {
            'columns': ", ".join(self.quote_name(column) for column in ['foo', 'bar']),
        }
        self.assertEqual(sql, 'UNIQUE ("foo", "bar")')

    def test_default_interpolation_for_create_unique_with_one_column(self):
        """Test that the it returns the correct result with one column ("sort_order")."""

        sql = DatabaseSchemaEditor.sql_create_unique % {
            'table': self.quote_name('my_table'),
            'name': self.quote_name('my_index'),
            'columns': ", ".join(self.quote_name(column) for column in ['sort_order']),
        }
        self.assertEqual(sql, 'ALTER TABLE "my_table" ADD CONSTRAINT "my_index" '
                              'UNIQUE ("sort_order") DEFERRABLE INITIALLY DEFERRED')

    def test_default_interpolation_for_create_unique_with_many_columns(self):
        """Test that the it returns the correct result with many columns (one is "sort_order")."""

        sql = DatabaseSchemaEditor.sql_create_unique % {
            'table': self.quote_name('my_table'),
            'name': self.quote_name('my_index'),
            'columns': ", ".join(self.quote_name(column) for column in ['foo', 'sort_order']),
        }
        self.assertEqual(sql, 'ALTER TABLE "my_table" ADD CONSTRAINT "my_index" '
                              'UNIQUE ("foo", "sort_order") DEFERRABLE INITIALLY DEFERRED')