Fixed #21039 -- Added AddIndexConcurrently/RemoveIndexConcurrently operations for PostgreSQL.

Thanks to Simon Charettes for review.

Co-Authored-By: Daniel Tao <daniel.tao@gmail.com>

AUTHORS

@@ -214,6 +214,7 @@ answer newbie questions, and generally made Django that much better:
     Daniel Poelzleithner <https://poelzi.org/>
     Daniel Pyrathon <pirosb3@gmail.com>
     Daniel Roseman <http://roseman.org.uk/>
+    Daniel Tao <https://philosopherdeveloper.com/>
     Daniel Wiesmann <daniel.wiesmann@gmail.com>
     Danilo Bargen
     Dan Johnson <danj.py@gmail.com>

django/contrib/postgres/indexes.py

@@ -18,9 +18,9 @@ class PostgresIndex(Index):
         # indexes.
         return Index.max_name_length - len(Index.suffix) + len(self.suffix)
 
-    def create_sql(self, model, schema_editor, using=''):
+    def create_sql(self, model, schema_editor, using='', **kwargs):
         self.check_supported(schema_editor)
-        statement = super().create_sql(model, schema_editor, using=' USING %s' % self.suffix)
+        statement = super().create_sql(model, schema_editor, using=' USING %s' % self.suffix, **kwargs)
         with_params = self.get_with_params()
         if with_params:
             statement.parts['extra'] = 'WITH (%s) %s' % (

django/contrib/postgres/operations.py

@@ -1,7 +1,9 @@
 from django.contrib.postgres.signals import (
     get_citext_oids, get_hstore_oids, register_type_handlers,
 )
+from django.db.migrations import AddIndex, RemoveIndex
 from django.db.migrations.operations.base import Operation
+from django.db.utils import NotSupportedError
 
 
 class CreateExtension(Operation):
@@ -75,3 +77,61 @@ class UnaccentExtension(CreateExtension):
 
     def __init__(self):
         self.name = 'unaccent'
+
+
+class NotInTransactionMixin:
+    def _ensure_not_in_transaction(self, schema_editor):
+        if schema_editor.connection.in_atomic_block:
+            raise NotSupportedError(
+                'The %s operation cannot be executed inside a transaction '
+                '(set atomic = False on the migration).'
+                % self.__class__.__name__
+            )
+
+
+class AddIndexConcurrently(NotInTransactionMixin, AddIndex):
+    """Create an index using PostgreSQL's CREATE INDEX CONCURRENTLY syntax."""
+    atomic = False
+
+    def describe(self):
+        return 'Concurrently create index %s on field(s) %s of model %s' % (
+            self.index.name,
+            ', '.join(self.index.fields),
+            self.model_name,
+        )
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        self._ensure_not_in_transaction(schema_editor)
+        model = to_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            schema_editor.add_index(model, self.index, concurrently=True)
+
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
+        self._ensure_not_in_transaction(schema_editor)
+        model = from_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            schema_editor.remove_index(model, self.index, concurrently=True)
+
+
+class RemoveIndexConcurrently(NotInTransactionMixin, RemoveIndex):
+    """Remove an index using PostgreSQL's DROP INDEX CONCURRENTLY syntax."""
+    atomic = False
+
+    def describe(self):
+        return 'Concurrently remove index %s from %s' % (self.name, self.model_name)
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        self._ensure_not_in_transaction(schema_editor)
+        model = from_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            from_model_state = from_state.models[app_label, self.model_name_lower]
+            index = from_model_state.get_index_by_name(self.name)
+            schema_editor.remove_index(model, index, concurrently=True)
+
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
+        self._ensure_not_in_transaction(schema_editor)
+        model = to_state.apps.get_model(app_label, self.model_name)
+        if self.allow_migrate_model(schema_editor.connection.alias, model):
+            to_model_state = to_state.models[app_label, self.model_name_lower]
+            index = to_model_state.get_index_by_name(self.name)
+            schema_editor.add_index(model, index, concurrently=True)

django/db/backends/base/schema.py

@@ -959,9 +959,9 @@ class BaseDatabaseSchemaEditor:
             condition=(' WHERE ' + condition) if condition else '',
         )
 
-    def _delete_index_sql(self, model, name):
+    def _delete_index_sql(self, model, name, sql=None):
         return Statement(
-            self.sql_delete_index,
+            sql or self.sql_delete_index,
             table=Table(model._meta.db_table, self.quote_name),
             name=self.quote_name(name),
         )

django/db/backends/postgresql/schema.py

@@ -13,7 +13,11 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
     sql_set_sequence_owner = 'ALTER SEQUENCE %(sequence)s OWNED BY %(table)s.%(column)s'
 
     sql_create_index = "CREATE INDEX %(name)s ON %(table)s%(using)s (%(columns)s)%(extra)s%(condition)s"
+    sql_create_index_concurrently = (
+        "CREATE INDEX CONCURRENTLY %(name)s ON %(table)s%(using)s (%(columns)s)%(extra)s%(condition)s"
+    )
     sql_delete_index = "DROP INDEX IF EXISTS %(name)s"
+    sql_delete_index_concurrently = "DROP INDEX CONCURRENTLY IF EXISTS %(name)s"
 
     sql_create_column_inline_fk = 'REFERENCES %(to_table)s(%(to_column)s)%(deferrable)s'
     # Setting the constraint to IMMEDIATE runs any deferred checks to allow
@@ -157,3 +161,24 @@ class DatabaseSchemaEditor(BaseDatabaseSchemaEditor):
         if opclasses:
             return IndexColumns(table, columns, self.quote_name, col_suffixes=col_suffixes, opclasses=opclasses)
         return super()._index_columns(table, columns, col_suffixes, opclasses)
+
+    def add_index(self, model, index, concurrently=False):
+        self.execute(index.create_sql(model, self, concurrently=concurrently), params=None)
+
+    def remove_index(self, model, index, concurrently=False):
+        self.execute(index.remove_sql(model, self, concurrently=concurrently))
+
+    def _delete_index_sql(self, model, name, sql=None, concurrently=False):
+        sql = self.sql_delete_index_concurrently if concurrently else self.sql_delete_index
+        return super()._delete_index_sql(model, name, sql)
+
+    def _create_index_sql(
+        self, model, fields, *, name=None, suffix='', using='',
+        db_tablespace=None, col_suffixes=(), sql=None, opclasses=(),
+        condition=None, concurrently=False,
+    ):
+        sql = self.sql_create_index if not concurrently else self.sql_create_index_concurrently
+        return super()._create_index_sql(
+            model, fields, name=name, suffix=suffix, using=using, db_tablespace=db_tablespace,
+            col_suffixes=col_suffixes, sql=sql, opclasses=opclasses, condition=condition,
+        )

django/db/models/indexes.py

@@ -49,17 +49,18 @@ class Index:
         # it's handled outside of that class, the work is done here.
         return sql % tuple(map(schema_editor.quote_value, params))
 
-    def create_sql(self, model, schema_editor, using=''):
+    def create_sql(self, model, schema_editor, using='', **kwargs):
         fields = [model._meta.get_field(field_name) for field_name, _ in self.fields_orders]
         col_suffixes = [order[1] for order in self.fields_orders]
         condition = self._get_condition_sql(model, schema_editor)
         return schema_editor._create_index_sql(
             model, fields, name=self.name, using=using, db_tablespace=self.db_tablespace,
             col_suffixes=col_suffixes, opclasses=self.opclasses, condition=condition,
+            **kwargs,
         )
 
-    def remove_sql(self, model, schema_editor):
-        return schema_editor._delete_index_sql(model, self.name)
+    def remove_sql(self, model, schema_editor, **kwargs):
+        return schema_editor._delete_index_sql(model, self.name, **kwargs)
 
     def deconstruct(self):
         path = '%s.%s' % (self.__class__.__module__, self.__class__.__name__)

docs/ref/contrib/postgres/operations.txt

@@ -98,3 +98,33 @@ run the query ``CREATE EXTENSION IF NOT EXISTS hstore;``.
 .. class:: UnaccentExtension()
 
     Installs the ``unaccent`` extension.
+
+Index concurrent operations
+===========================
+
+.. versionadded:: 3.0
+
+PostgreSQL supports the ``CONCURRENTLY`` option to ``CREATE INDEX`` and
+``DROP INDEX`` statements to add and remove indexes without locking out writes.
+This option is useful for adding or removing an index in a live production
+database.
+
+.. class:: AddIndexConcurrently(model_name, index)
+
+    Like :class:`~django.db.migrations.operations.AddIndex`, but creates an
+    index with the ``CONCURRENTLY`` option. This has a few caveats to be aware
+    of when using this option, see `the PostgreSQL documentation of building
+    indexes concurrently <https://www.postgresql.org/docs/current/
+    sql-createindex.html#SQL-CREATEINDEX-CONCURRENTLY>`_.
+
+.. class:: RemoveIndexConcurrently(model_name, name)
+
+    Like :class:`~django.db.migrations.operations.RemoveIndex`, but removes the
+    index with the ``CONCURRENTLY`` option. This has a few caveats to be aware
+    of when using this option, see `the PostgreSQL documentation
+    <https://www.postgresql.org/docs/current/sql-dropindex.html>`_.
+
+.. note::
+
+    The ``CONCURRENTLY`` option is not supported inside a transaction (see
+    :ref:`non-atomic migration <non-atomic-migrations>`).

docs/releases/3.0.txt

@@ -153,6 +153,10 @@ Minor features
 * The new :class:`~django.contrib.postgres.fields.RangeBoundary` expression
   represents the range boundaries.
 
+* The new :class:`~django.contrib.postgres.operations.AddIndexConcurrently`
+  and :class:`~django.contrib.postgres.operations.RemoveIndexConcurrently`
+  classes allow creating and dropping indexes ``CONCURRENTLY`` on PostgreSQL.
+
 :mod:`django.contrib.redirects`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

tests/migrations/test_base.py

@@ -13,8 +13,6 @@ from django.test import TransactionTestCase
 from django.test.utils import extend_sys_path
 from django.utils.module_loading import module_dir
 
-from .models import FoodManager, FoodQuerySet
-
 
 class MigrationTestBase(TransactionTestCase):
     """
@@ -57,14 +55,14 @@ class MigrationTestBase(TransactionTestCase):
     def assertColumnNotNull(self, table, column, using='default'):
         self.assertEqual(self._get_column_allows_null(table, column, using), False)
 
-    def assertIndexExists(self, table, columns, value=True, using='default'):
+    def assertIndexExists(self, table, columns, value=True, using='default', index_type=None):
         with connections[using].cursor() as cursor:
             self.assertEqual(
                 value,
                 any(
                     c["index"]
                     for c in connections[using].introspection.get_constraints(cursor, table).values()
-                    if c['columns'] == list(columns)
+                    if c['columns'] == list(columns) and (index_type is None or c['type'] == index_type)
                 ),
             )
 
@@ -266,6 +264,7 @@ class OperationTestBase(MigrationTestBase):
                 bases=['%s.Pony' % app_label],
             ))
         if manager_model:
+            from .models import FoodManager, FoodQuerySet
             operations.append(migrations.CreateModel(
                 'Food',
                 fields=[

tests/postgres_tests/test_operations.py

@@ -0,0 +1,144 @@
+import unittest
+
+from migrations.test_base import OperationTestBase
+
+from django.db import connection, models
+from django.db.models import Index
+from django.db.utils import NotSupportedError
+from django.test import modify_settings
+
+try:
+    from django.contrib.postgres.operations import (
+        AddIndexConcurrently, RemoveIndexConcurrently,
+    )
+    from django.contrib.postgres.indexes import BrinIndex, BTreeIndex
+except ImportError:
+    pass
+
+
+@unittest.skipUnless(connection.vendor == 'postgresql', 'PostgreSQL specific tests.')
+@modify_settings(INSTALLED_APPS={'append': 'migrations'})
+class AddIndexConcurrentlyTests(OperationTestBase):
+    app_label = 'test_add_concurrently'
+
+    def test_requires_atomic_false(self):
+        project_state = self.set_up_test_model(self.app_label)
+        new_state = project_state.clone()
+        operation = AddIndexConcurrently(
+            'Pony',
+            models.Index(fields=['pink'], name='pony_pink_idx'),
+        )
+        msg = (
+            'The AddIndexConcurrently operation cannot be executed inside '
+            'a transaction (set atomic = False on the migration).'
+        )
+        with self.assertRaisesMessage(NotSupportedError, msg):
+            with connection.schema_editor(atomic=True) as editor:
+                operation.database_forwards(self.app_label, editor, project_state, new_state)
+
+    def test_add(self):
+        project_state = self.set_up_test_model(self.app_label, index=False)
+        table_name = '%s_pony' % self.app_label
+        index = Index(fields=['pink'], name='pony_pink_idx')
+        new_state = project_state.clone()
+        operation = AddIndexConcurrently('Pony', index)
+        self.assertEqual(
+            operation.describe(),
+            'Concurrently create index pony_pink_idx on field(s) pink of '
+            'model Pony'
+        )
+        operation.state_forwards(self.app_label, new_state)
+        self.assertEqual(len(new_state.models[self.app_label, 'pony'].options['indexes']), 1)
+        self.assertIndexNotExists(table_name, ['pink'])
+        # Add index.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        self.assertIndexExists(table_name, ['pink'])
+        # Reversal.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_backwards(self.app_label, editor, new_state, project_state)
+        self.assertIndexNotExists(table_name, ['pink'])
+        # Deconstruction.
+        name, args, kwargs = operation.deconstruct()
+        self.assertEqual(name, 'AddIndexConcurrently')
+        self.assertEqual(args, [])
+        self.assertEqual(kwargs, {'model_name': 'Pony', 'index': index})
+
+    def test_add_other_index_type(self):
+        project_state = self.set_up_test_model(self.app_label, index=False)
+        table_name = '%s_pony' % self.app_label
+        new_state = project_state.clone()
+        operation = AddIndexConcurrently(
+            'Pony',
+            BrinIndex(fields=['pink'], name='pony_pink_brin_idx'),
+        )
+        self.assertIndexNotExists(table_name, ['pink'])
+        # Add index.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        self.assertIndexExists(table_name, ['pink'], index_type='brin')
+        # Reversal.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_backwards(self.app_label, editor, new_state, project_state)
+        self.assertIndexNotExists(table_name, ['pink'])
+
+    def test_add_with_options(self):
+        project_state = self.set_up_test_model(self.app_label, index=False)
+        table_name = '%s_pony' % self.app_label
+        new_state = project_state.clone()
+        index = BTreeIndex(fields=['pink'], name='pony_pink_btree_idx', fillfactor=70)
+        operation = AddIndexConcurrently('Pony', index)
+        self.assertIndexNotExists(table_name, ['pink'])
+        # Add index.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        self.assertIndexExists(table_name, ['pink'], index_type='btree')
+        # Reversal.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_backwards(self.app_label, editor, new_state, project_state)
+        self.assertIndexNotExists(table_name, ['pink'])
+
+
+@unittest.skipUnless(connection.vendor == 'postgresql', 'PostgreSQL specific tests.')
+@modify_settings(INSTALLED_APPS={'append': 'migrations'})
+class RemoveIndexConcurrentlyTests(OperationTestBase):
+    app_label = 'test_rm_concurrently'
+
+    def test_requires_atomic_false(self):
+        project_state = self.set_up_test_model(self.app_label, index=True)
+        new_state = project_state.clone()
+        operation = RemoveIndexConcurrently('Pony', 'pony_pink_idx')
+        msg = (
+            'The RemoveIndexConcurrently operation cannot be executed inside '
+            'a transaction (set atomic = False on the migration).'
+        )
+        with self.assertRaisesMessage(NotSupportedError, msg):
+            with connection.schema_editor(atomic=True) as editor:
+                operation.database_forwards(self.app_label, editor, project_state, new_state)
+
+    def test_remove(self):
+        project_state = self.set_up_test_model(self.app_label, index=True)
+        table_name = '%s_pony' % self.app_label
+        self.assertTableExists(table_name)
+        new_state = project_state.clone()
+        operation = RemoveIndexConcurrently('Pony', 'pony_pink_idx')
+        self.assertEqual(
+            operation.describe(),
+            'Concurrently remove index pony_pink_idx from Pony',
+        )
+        operation.state_forwards(self.app_label, new_state)
+        self.assertEqual(len(new_state.models[self.app_label, 'pony'].options['indexes']), 0)
+        self.assertIndexExists(table_name, ['pink'])
+        # Remove index.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_forwards(self.app_label, editor, project_state, new_state)
+        self.assertIndexNotExists(table_name, ['pink'])
+        # Reversal.
+        with connection.schema_editor(atomic=False) as editor:
+            operation.database_backwards(self.app_label, editor, new_state, project_state)
+        self.assertIndexExists(table_name, ['pink'])
+        # Deconstruction.
+        name, args, kwargs = operation.deconstruct()
+        self.assertEqual(name, 'RemoveIndexConcurrently')
+        self.assertEqual(args, [])
+        self.assertEqual(kwargs, {'model_name': 'Pony', 'name': 'pony_pink_idx'})