dropping constraints and indexes before drop table or column

CHANGES.md

@@ -4,7 +4,9 @@
   - changed `ADD COLUMN DEFAULT NULL` to safe operation for code default
   - changed `ADD COLUMN DEFAULT NOT NULL` to safe operation for `db_default` in django 5.0+
   - added `ZERO_DOWNTIME_MIGRATIONS_KEEP_DEFAULT` settings and changed `ADD COLUMN DEFAULT NOT NULL` with this settings to safe operation for django<5.0
+  - added `ZERO_DOWNTIME_MIGRATIONS_DROP_CONSTRAINS_BEFORE_CASCADE` settings and added dropping constraints and indexes before drop column or table
   - fixed sqlmigrate in idempotent mode
+  - fixed creation unique constraint with include parameter
   - updated unsafe migrations links to documentation
   - updated patched code to latest django version
   - updated test image to ubuntu 24.04

django_zero_downtime_migrations/backends/postgres/schema.py

@@ -5,7 +5,7 @@
 import django
 from django.conf import settings
 from django.contrib.postgres.constraints import ExclusionConstraint
-from django.db.backends.ddl_references import Statement
+from django.db.backends.ddl_references import Statement, Table
 from django.db.backends.postgresql.schema import (
     DatabaseSchemaEditor as PostgresDatabaseSchemaEditor
 )
@@ -417,7 +417,8 @@ class DatabaseSchemaEditorMixin:
     if django.VERSION[:2] >= (5, 0):
         sql_create_unique_index = MultiStatementSQL(
             PGShareUpdateExclusive(
-                "CREATE UNIQUE INDEX CONCURRENTLY %(name)s ON %(table)s (%(columns)s)%(condition)s%(nulls_distinct)s",
+                "CREATE UNIQUE INDEX CONCURRENTLY %(name)s ON %(table)s "
+                "(%(columns)s)%(include)s%(nulls_distinct)s%(condition)s",
                 idempotent_condition=Condition(_sql_index_exists, False),
                 disable_statement_timeout=True,
             ),
@@ -465,6 +466,54 @@ class DatabaseSchemaEditorMixin:
         ),
     )
 
+    _sql_get_table_constraints_introspection = """
+        SELECT
+            c.conname,
+            c.contype,
+            c.conrelid::regclass::text,
+            c.confrelid::regclass::text,
+            array(
+                SELECT attname
+                FROM unnest(c.conkey) WITH ORDINALITY cols(colid, arridx)
+                JOIN pg_attribute AS ca ON cols.colid = ca.attnum
+                WHERE ca.attrelid = c.conrelid
+                ORDER BY cols.arridx
+            ),
+            array(
+                SELECT attname
+                FROM unnest(c.confkey) WITH ORDINALITY cols(colid, arridx)
+                JOIN pg_attribute AS ca ON cols.colid = ca.attnum
+                WHERE ca.attrelid = c.confrelid
+                ORDER BY cols.arridx
+            )
+        FROM pg_constraint AS c
+        WHERE c.conrelid::regclass::text = %s
+        OR c.confrelid::regclass::text = %s
+        ORDER BY c.conrelid::regclass::text, c.conname
+    """
+    _sql_get_index_introspection = """
+        SELECT
+            i.indexrelid::regclass::text,
+            i.indrelid::regclass::text,
+            array(
+                SELECT a.attname
+                FROM (
+                    SELECT unnest(i.indkey)
+                    UNION
+                    SELECT unnest(regexp_matches(i.indexprs::text, ':varattno (\d+)', 'g'))::int
+                    UNION
+                    SELECT unnest(regexp_matches(i.indpred::text, ':varattno (\d+)', 'g'))::int
+                ) cols(varattno)
+                INNER JOIN pg_attribute AS a ON cols.varattno = a.attnum
+                WHERE a.attrelid = i.indrelid
+            )
+        FROM pg_index i
+        LEFT JOIN pg_constraint c ON i.indexrelid = c.conindid 
+        WHERE indrelid::regclass::text = %s
+        AND c.conindid IS NULL
+        ORDER BY i.indrelid::regclass::text, i.indexrelid::regclass::text
+    """
+
     _varchar_type_regexp = re.compile(r'^varchar\((?P<max_length>\d+)\)$')
     _numeric_type_regexp = re.compile(r'^numeric\((?P<precision>\d+), *(?P<scale>\d+)\)$')
 
@@ -508,6 +557,8 @@ def __init__(self, connection, collect_sql=False, atomic=True):
                 DeprecationWarning,
             )
             self.KEEP_DEFAULT = False
+        self.DROP_CONSTRAINS_BEFORE_CASCADE = getattr(
+            settings, "ZERO_DOWNTIME_MIGRATIONS_DROP_CONSTRAINS_BEFORE_CASCADE", True)
 
     def execute(self, sql, params=()):
         if sql is DUMMY_SQL:
@@ -599,13 +650,50 @@ def _flush_deferred_sql(self):
                 self.execute(sql)
             self.deferred_sql.clear()
 
+    def _get_constraints(self, cursor, model):
+        cursor.execute(self._sql_get_table_constraints_introspection, [model._meta.db_table, model._meta.db_table])
+        for constraint, kind, table, table_ref, columns, columns_ref in cursor.fetchall():
+            yield constraint, kind, table, table_ref, columns, columns_ref
+
+    def _get_indexes(self, cursor, model):
+        cursor.execute(self._sql_get_index_introspection, [model._meta.db_table])
+        for index, table, columns in cursor.fetchall():
+            yield index, table, columns
+
+    def _drop_collect_sql_introspection_related_duplicates(self, drop_constraint_queries):
+        # django internals use introspection to find related constraints
+        # that can duplicated explicit constraints dropping in this method
+        # so just clean up collected sql for this duplicates
+        if self.collect_sql:
+            handled_queries = set()
+            deduplicated_collected_sql = []
+            for query in self.collected_sql:
+                if query not in handled_queries:
+                    deduplicated_collected_sql.append(query)
+                    if query in drop_constraint_queries:
+                        handled_queries.add(query)
+            self.collected_sql = deduplicated_collected_sql
+
     def create_model(self, model):
         super().create_model(model)
         self._flush_deferred_sql()
 
     def delete_model(self, model):
+        drop_constraint_queries = set()
+        if self.DROP_CONSTRAINS_BEFORE_CASCADE:
+            with self.connection.cursor() as cursor:
+                for constraint, kind, table, table_ref, columns, columns_ref in self._get_constraints(cursor, model):
+                    if kind == "f":
+                        self.execute(Statement(
+                            self.sql_delete_fk,
+                            table=Table(table, self.quote_name),
+                            name=self.quote_name(constraint),
+                        ))
+                        if self.collect_sql:
+                            drop_constraint_queries.add(self.collected_sql[-1])
         super().delete_model(model)
         self._flush_deferred_sql()
+        self._drop_collect_sql_introspection_related_duplicates(drop_constraint_queries)
 
     def alter_index_together(self, model, old_index_together, new_index_together):
         super().alter_index_together(model, old_index_together, new_index_together)
@@ -641,8 +729,52 @@ def add_field(self, model, field):
         self._flush_deferred_sql()
 
     def remove_field(self, model, field):
+        drop_constraint_queries = set()
+        if self.DROP_CONSTRAINS_BEFORE_CASCADE:
+            with self.connection.cursor() as cursor:
+                # as foreign key can have index as dependent object it important to drop all foreign keys first
+                for constraint, kind, table, table_ref, columns, columns_ref in self._get_constraints(cursor, model):
+                    # drop foreign key for current model columns
+                    if kind == "f" and table == model._meta.db_table and field.column in columns:
+                        self.execute(Statement(
+                            self.sql_delete_fk,
+                            table=Table(table, self.quote_name),
+                            name=self.quote_name(constraint),
+                        ))
+                        if self.collect_sql:
+                            drop_constraint_queries.add(self.collected_sql[-1])
+                    # drop foreign key for target model columns, i.e. backrefs
+                    if kind == "f" and table_ref == model._meta.db_table and field.column in columns_ref:
+                        self.execute(Statement(
+                            self.sql_delete_fk,
+                            table=Table(table, self.quote_name),
+                            name=self.quote_name(constraint),
+                        ))
+                        if self.collect_sql:
+                            drop_constraint_queries.add(self.collected_sql[-1])
+                for constraint, kind, table, table_ref, columns, columns_ref in self._get_constraints(cursor, model):
+                    # drop unique constraints for current model columns
+                    if kind == "u" and table == model._meta.db_table and field.column in columns:
+                        self.execute(Statement(
+                            self.sql_delete_unique,
+                            table=Table(table, self.quote_name),
+                            name=self.quote_name(constraint),
+                        ))
+                        if self.collect_sql:
+                            drop_constraint_queries.add(self.collected_sql[-1])
+                for index, table, columns in self._get_indexes(cursor, model):
+                    # drop indexes for current model columns
+                    if table == model._meta.db_table and field.column in columns:
+                        self.execute(Statement(
+                            self.sql_delete_index_concurrently,
+                            table=Table(table, self.quote_name),
+                            name=self.quote_name(index),
+                        ))
+                        if self.collect_sql:
+                            drop_constraint_queries.add(self.collected_sql[-1])
         super().remove_field(model, field)
         self._flush_deferred_sql()
+        self._drop_collect_sql_introspection_related_duplicates(drop_constraint_queries)
 
     def alter_field(self, model, old_field, new_field, strict=False):
         super().alter_field(model, old_field, new_field, strict)

tests/apps/good_flow_drop_column_with_constraints/migrations/0001_initial.py

@@ -0,0 +1,142 @@
+from django.db import migrations, models
+import django.db.models.deletion
+import django.db.models.functions.math
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = []
+
+    operations = [
+        migrations.CreateModel(
+            name="TestTableParent",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+            ],
+            options={
+                "db_table": "drop_col_test_table_parent",
+            },
+        ),
+        migrations.CreateModel(
+            name="TestTableMain",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("main_id", models.IntegerField(null=True, unique=True)),
+                ("field_u1", models.IntegerField(null=True)),
+                ("field_u2", models.IntegerField(null=True)),
+                ("field_u3", models.IntegerField(null=True)),
+                ("field_u4", models.IntegerField(null=True)),
+                ("field_i1", models.IntegerField(null=True)),
+                ("field_i2", models.IntegerField(null=True)),
+                ("field_i3", models.IntegerField(null=True)),
+                ("field_i4", models.IntegerField(null=True)),
+                (
+                    "parent",
+                    models.OneToOneField(
+                        null=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="good_flow_drop_column_with_constraints.testtableparent",
+                    ),
+                ),
+            ],
+            options={
+                "db_table": "drop_col_test_table_main",
+            },
+        ),
+        migrations.CreateModel(
+            name="TestTableChild",
+            fields=[
+                (
+                    "id",
+                    models.AutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                (
+                    "main",
+                    models.ForeignKey(
+                        null=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="good_flow_drop_column_with_constraints.testtablemain",
+                        to_field="main_id",
+                    ),
+                ),
+            ],
+            options={
+                "db_table": "drop_col_test_table_child",
+            },
+        ),
+        migrations.AddIndex(
+            model_name="testtablemain",
+            index=models.Index(
+                models.F("parent"), models.F("field_i1"), name="drop_col_i1"
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="testtablemain",
+            index=models.Index(
+                django.db.models.functions.math.Abs("field_i2"), name="drop_col_i2"
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="testtablemain",
+            index=models.Index(
+                models.F("parent"),
+                condition=models.Q(("field_i3__gt", 0)),
+                name="drop_col_i3",
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="testtablemain",
+            index=models.Index(
+                models.F("parent"), include=("field_i4",), name="drop_col_i4"
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="testtablemain",
+            constraint=models.UniqueConstraint(
+                models.F("parent"), models.F("field_u1"), name="drop_col_u1"
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="testtablemain",
+            constraint=models.UniqueConstraint(
+                django.db.models.functions.math.Abs("field_u2"), name="drop_col_u2"
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="testtablemain",
+            constraint=models.UniqueConstraint(
+                models.F("parent"),
+                condition=models.Q(("field_u3__gt", 0)),
+                name="drop_col_u3",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="testtablemain",
+            constraint=models.UniqueConstraint(
+                models.F("parent"), include=("field_u4",), name="drop_col_u4"
+            ),
+        ),
+    ]

tests/apps/good_flow_drop_column_with_constraints/migrations/0002_remove_testtablemain_drop_col_u1_and_more.py

@@ -0,0 +1,52 @@
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("good_flow_drop_column_with_constraints", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            database_operations=[],
+            state_operations=[
+                migrations.RemoveConstraint(
+                    model_name="testtablemain",
+                    name="drop_col_u1",
+                ),
+                migrations.RemoveConstraint(
+                    model_name="testtablemain",
+                    name="drop_col_u2",
+                ),
+                migrations.RemoveConstraint(
+                    model_name="testtablemain",
+                    name="drop_col_u3",
+                ),
+                migrations.RemoveConstraint(
+                    model_name="testtablemain",
+                    name="drop_col_u4",
+                ),
+                migrations.RemoveIndex(
+                    model_name="testtablemain",
+                    name="drop_col_i1",
+                ),
+                migrations.RemoveIndex(
+                    model_name="testtablemain",
+                    name="drop_col_i2",
+                ),
+                migrations.RemoveIndex(
+                    model_name="testtablemain",
+                    name="drop_col_i3",
+                ),
+                migrations.RemoveIndex(
+                    model_name="testtablemain",
+                    name="drop_col_i4",
+                ),
+                migrations.RemoveField(
+                    model_name="testtablechild",
+                    name="main",
+                ),
+            ],
+        )
+    ]