PostgreSQL index concurrently

See also

Description

django/contrib/postgres/operations.py

See also 

  1from django.contrib.postgres.signals import get_citext_oids
  2from django.contrib.postgres.signals import get_hstore_oids
  3from django.contrib.postgres.signals import register_type_handlers
  4from django.db.migrations import AddIndex
  5from django.db.migrations import RemoveIndex
  6from django.db.migrations.operations.base import Operation
  7from django.db.utils import NotSupportedError
  8
  9
 10class CreateExtension(Operation):
 11    reversible = True
 12
 13    def __init__(self, name):
 14        self.name = name
 15
 16    def state_forwards(self, app_label, state):
 17        pass
 18
 19    def database_forwards(self, app_label, schema_editor, from_state, to_state):
 20        if schema_editor.connection.vendor != "postgresql":
 21            return
 22        schema_editor.execute(
 23            "CREATE EXTENSION IF NOT EXISTS %s" % schema_editor.quote_name(self.name)
 24        )
 25        # Clear cached, stale oids.
 26        get_hstore_oids.cache_clear()
 27        get_citext_oids.cache_clear()
 28        # Registering new type handlers cannot be done before the extension is
 29        # installed, otherwise a subsequent data migration would use the same
 30        # connection.
 31        register_type_handlers(schema_editor.connection)
 32
 33    def database_backwards(self, app_label, schema_editor, from_state, to_state):
 34        schema_editor.execute("DROP EXTENSION %s" % schema_editor.quote_name(self.name))
 35        # Clear cached, stale oids.
 36        get_hstore_oids.cache_clear()
 37        get_citext_oids.cache_clear()
 38
 39    def describe(self):
 40        return "Creates extension %s" % self.name
 41
 42
 43class BloomExtension(CreateExtension):
 44    def __init__(self):
 45        self.name = "bloom"
 46
 47
 48class BtreeGinExtension(CreateExtension):
 49    def __init__(self):
 50        self.name = "btree_gin"
 51
 52
 53class BtreeGistExtension(CreateExtension):
 54    def __init__(self):
 55        self.name = "btree_gist"
 56
 57
 58class CITextExtension(CreateExtension):
 59    def __init__(self):
 60        self.name = "citext"
 61
 62
 63class CryptoExtension(CreateExtension):
 64    def __init__(self):
 65        self.name = "pgcrypto"
 66
 67
 68class HStoreExtension(CreateExtension):
 69    def __init__(self):
 70        self.name = "hstore"
 71
 72
 73class TrigramExtension(CreateExtension):
 74    def __init__(self):
 75        self.name = "pg_trgm"
 76
 77
 78class UnaccentExtension(CreateExtension):
 79    def __init__(self):
 80        self.name = "unaccent"
 81
 82
 83class NotInTransactionMixin:
 84    def _ensure_not_in_transaction(self, schema_editor):
 85        if schema_editor.connection.in_atomic_block:
 86            raise NotSupportedError(
 87                "The %s operation cannot be executed inside a transaction "
 88                "(set atomic = False on the migration)." % self.__class__.__name__
 89            )
 90
 91
 92class AddIndexConcurrently(NotInTransactionMixin, AddIndex):
 93    """Create an index using PostgreSQL's CREATE INDEX CONCURRENTLY syntax."""
 94
 95    atomic = False
 96
 97    def describe(self):
 98        return "Concurrently create index %s on field(s) %s of model %s" % (
 99            self.index.name,
100            ", ".join(self.index.fields),
101            self.model_name,
102        )
103
104    def database_forwards(self, app_label, schema_editor, from_state, to_state):
105        self._ensure_not_in_transaction(schema_editor)
106        model = to_state.apps.get_model(app_label, self.model_name)
107        if self.allow_migrate_model(schema_editor.connection.alias, model):
108            schema_editor.add_index(model, self.index, concurrently=True)
109
110    def database_backwards(self, app_label, schema_editor, from_state, to_state):
111        self._ensure_not_in_transaction(schema_editor)
112        model = from_state.apps.get_model(app_label, self.model_name)
113        if self.allow_migrate_model(schema_editor.connection.alias, model):
114            schema_editor.remove_index(model, self.index, concurrently=True)
115
116
117class RemoveIndexConcurrently(NotInTransactionMixin, RemoveIndex):
118    """Remove an index using PostgreSQL's DROP INDEX CONCURRENTLY syntax."""
119
120    atomic = False
121
122    def describe(self):
123        return "Concurrently remove index %s from %s" % (self.name, self.model_name)
124
125    def database_forwards(self, app_label, schema_editor, from_state, to_state):
126        self._ensure_not_in_transaction(schema_editor)
127        model = from_state.apps.get_model(app_label, self.model_name)
128        if self.allow_migrate_model(schema_editor.connection.alias, model):
129            from_model_state = from_state.models[app_label, self.model_name_lower]
130            index = from_model_state.get_index_by_name(self.name)
131            schema_editor.remove_index(model, index, concurrently=True)
132
133    def database_backwards(self, app_label, schema_editor, from_state, to_state):
134        self._ensure_not_in_transaction(schema_editor)
135        model = to_state.apps.get_model(app_label, self.model_name)
136        if self.allow_migrate_model(schema_editor.connection.alias, model):
137            to_model_state = to_state.models[app_label, self.model_name_lower]
138            index = to_model_state.get_index_by_name(self.name)
139            schema_editor.add_index(model, index, concurrently=True)

tests/postgres_tests/test_operations.py

See also 

  1import unittest
  2
  3from django.db import connection
  4from django.db import models
  5from django.db.models import Index
  6from django.db.utils import NotSupportedError
  7from django.test import modify_settings
  8from migrations.test_base import OperationTestBase
  9
 10try:
 11    from django.contrib.postgres.operations import (
 12        AddIndexConcurrently,
 13        RemoveIndexConcurrently,
 14    )
 15    from django.contrib.postgres.indexes import BrinIndex, BTreeIndex
 16except ImportError:
 17    pass
 18
 19
 20@unittest.skipUnless(connection.vendor == "postgresql", "PostgreSQL specific tests.")
 21@modify_settings(INSTALLED_APPS={"append": "migrations"})
 22class AddIndexConcurrentlyTests(OperationTestBase):
 23    app_label = "test_add_concurrently"
 24
 25    def test_requires_atomic_false(self):
 26        project_state = self.set_up_test_model(self.app_label)
 27        new_state = project_state.clone()
 28        operation = AddIndexConcurrently(
 29            "Pony",
 30            models.Index(fields=["pink"], name="pony_pink_idx"),
 31        )
 32        msg = (
 33            "The AddIndexConcurrently operation cannot be executed inside "
 34            "a transaction (set atomic = False on the migration)."
 35        )
 36        with self.assertRaisesMessage(NotSupportedError, msg):
 37            with connection.schema_editor(atomic=True) as editor:
 38                operation.database_forwards(
 39                    self.app_label, editor, project_state, new_state
 40                )
 41
 42    def test_add(self):
 43        project_state = self.set_up_test_model(self.app_label, index=False)
 44        table_name = "%s_pony" % self.app_label
 45        index = Index(fields=["pink"], name="pony_pink_idx")
 46        new_state = project_state.clone()
 47        operation = AddIndexConcurrently("Pony", index)
 48        self.assertEqual(
 49            operation.describe(),
 50            "Concurrently create index pony_pink_idx on field(s) pink of " "model Pony",
 51        )
 52        operation.state_forwards(self.app_label, new_state)
 53        self.assertEqual(
 54            len(new_state.models[self.app_label, "pony"].options["indexes"]), 1
 55        )
 56        self.assertIndexNotExists(table_name, ["pink"])
 57        # Add index.
 58        with connection.schema_editor(atomic=False) as editor:
 59            operation.database_forwards(
 60                self.app_label, editor, project_state, new_state
 61            )
 62        self.assertIndexExists(table_name, ["pink"])
 63        # Reversal.
 64        with connection.schema_editor(atomic=False) as editor:
 65            operation.database_backwards(
 66                self.app_label, editor, new_state, project_state
 67            )
 68        self.assertIndexNotExists(table_name, ["pink"])
 69        # Deconstruction.
 70        name, args, kwargs = operation.deconstruct()
 71        self.assertEqual(name, "AddIndexConcurrently")
 72        self.assertEqual(args, [])
 73        self.assertEqual(kwargs, {"model_name": "Pony", "index": index})
 74
 75    def test_add_other_index_type(self):
 76        project_state = self.set_up_test_model(self.app_label, index=False)
 77        table_name = "%s_pony" % self.app_label
 78        new_state = project_state.clone()
 79        operation = AddIndexConcurrently(
 80            "Pony",
 81            BrinIndex(fields=["pink"], name="pony_pink_brin_idx"),
 82        )
 83        self.assertIndexNotExists(table_name, ["pink"])
 84        # Add index.
 85        with connection.schema_editor(atomic=False) as editor:
 86            operation.database_forwards(
 87                self.app_label, editor, project_state, new_state
 88            )
 89        self.assertIndexExists(table_name, ["pink"], index_type="brin")
 90        # Reversal.
 91        with connection.schema_editor(atomic=False) as editor:
 92            operation.database_backwards(
 93                self.app_label, editor, new_state, project_state
 94            )
 95        self.assertIndexNotExists(table_name, ["pink"])
 96
 97    def test_add_with_options(self):
 98        project_state = self.set_up_test_model(self.app_label, index=False)
 99        table_name = "%s_pony" % self.app_label
100        new_state = project_state.clone()
101        index = BTreeIndex(fields=["pink"], name="pony_pink_btree_idx", fillfactor=70)
102        operation = AddIndexConcurrently("Pony", index)
103        self.assertIndexNotExists(table_name, ["pink"])
104        # Add index.
105        with connection.schema_editor(atomic=False) as editor:
106            operation.database_forwards(
107                self.app_label, editor, project_state, new_state
108            )
109        self.assertIndexExists(table_name, ["pink"], index_type="btree")
110        # Reversal.
111        with connection.schema_editor(atomic=False) as editor:
112            operation.database_backwards(
113                self.app_label, editor, new_state, project_state
114            )
115        self.assertIndexNotExists(table_name, ["pink"])
116
117
118@unittest.skipUnless(connection.vendor == "postgresql", "PostgreSQL specific tests.")
119@modify_settings(INSTALLED_APPS={"append": "migrations"})
120class RemoveIndexConcurrentlyTests(OperationTestBase):
121    app_label = "test_rm_concurrently"
122
123    def test_requires_atomic_false(self):
124        project_state = self.set_up_test_model(self.app_label, index=True)
125        new_state = project_state.clone()
126        operation = RemoveIndexConcurrently("Pony", "pony_pink_idx")
127        msg = (
128            "The RemoveIndexConcurrently operation cannot be executed inside "
129            "a transaction (set atomic = False on the migration)."
130        )
131        with self.assertRaisesMessage(NotSupportedError, msg):
132            with connection.schema_editor(atomic=True) as editor:
133                operation.database_forwards(
134                    self.app_label, editor, project_state, new_state
135                )
136
137    def test_remove(self):
138        project_state = self.set_up_test_model(self.app_label, index=True)
139        table_name = "%s_pony" % self.app_label
140        self.assertTableExists(table_name)
141        new_state = project_state.clone()
142        operation = RemoveIndexConcurrently("Pony", "pony_pink_idx")
143        self.assertEqual(
144            operation.describe(),
145            "Concurrently remove index pony_pink_idx from Pony",
146        )
147        operation.state_forwards(self.app_label, new_state)
148        self.assertEqual(
149            len(new_state.models[self.app_label, "pony"].options["indexes"]), 0
150        )
151        self.assertIndexExists(table_name, ["pink"])
152        # Remove index.
153        with connection.schema_editor(atomic=False) as editor:
154            operation.database_forwards(
155                self.app_label, editor, project_state, new_state
156            )
157        self.assertIndexNotExists(table_name, ["pink"])
158        # Reversal.
159        with connection.schema_editor(atomic=False) as editor:
160            operation.database_backwards(
161                self.app_label, editor, new_state, project_state
162            )
163        self.assertIndexExists(table_name, ["pink"])
164        # Deconstruction.
165        name, args, kwargs = operation.deconstruct()
166        self.assertEqual(name, "RemoveIndexConcurrently")
167        self.assertEqual(args, [])
168        self.assertEqual(kwargs, {"model_name": "Pony", "name": "pony_pink_idx"})