123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137 |
- 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):
- reversible = True
- def __init__(self, name):
- self.name = name
- def state_forwards(self, app_label, state):
- pass
- def database_forwards(self, app_label, schema_editor, from_state, to_state):
- if schema_editor.connection.vendor != 'postgresql':
- return
- schema_editor.execute("CREATE EXTENSION IF NOT EXISTS %s" % schema_editor.quote_name(self.name))
- # Clear cached, stale oids.
- get_hstore_oids.cache_clear()
- get_citext_oids.cache_clear()
- # Registering new type handlers cannot be done before the extension is
- # installed, otherwise a subsequent data migration would use the same
- # connection.
- register_type_handlers(schema_editor.connection)
- def database_backwards(self, app_label, schema_editor, from_state, to_state):
- schema_editor.execute("DROP EXTENSION %s" % schema_editor.quote_name(self.name))
- # Clear cached, stale oids.
- get_hstore_oids.cache_clear()
- get_citext_oids.cache_clear()
- def describe(self):
- return "Creates extension %s" % self.name
- class BtreeGinExtension(CreateExtension):
- def __init__(self):
- self.name = 'btree_gin'
- class BtreeGistExtension(CreateExtension):
- def __init__(self):
- self.name = 'btree_gist'
- class CITextExtension(CreateExtension):
- def __init__(self):
- self.name = 'citext'
- class CryptoExtension(CreateExtension):
- def __init__(self):
- self.name = 'pgcrypto'
- class HStoreExtension(CreateExtension):
- def __init__(self):
- self.name = 'hstore'
- class TrigramExtension(CreateExtension):
- def __init__(self):
- self.name = 'pg_trgm'
- 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)
|