123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335 |
- from django.contrib.postgres.signals import (
- get_citext_oids,
- get_hstore_oids,
- register_type_handlers,
- )
- from django.db import NotSupportedError, router
- from django.db.migrations import AddConstraint, AddIndex, RemoveIndex
- from django.db.migrations.operations.base import Operation
- from django.db.models.constraints import CheckConstraint
-
-
- 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" or not router.allow_migrate(
- schema_editor.connection.alias, app_label
- ):
- return
- if not self.extension_exists(schema_editor, self.name):
- 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):
- if not router.allow_migrate(schema_editor.connection.alias, app_label):
- return
- if self.extension_exists(schema_editor, self.name):
- schema_editor.execute(
- "DROP EXTENSION IF EXISTS %s" % schema_editor.quote_name(self.name)
- )
- # Clear cached, stale oids.
- get_hstore_oids.cache_clear()
- get_citext_oids.cache_clear()
-
- def extension_exists(self, schema_editor, extension):
- with schema_editor.connection.cursor() as cursor:
- cursor.execute(
- "SELECT 1 FROM pg_extension WHERE extname = %s",
- [extension],
- )
- return bool(cursor.fetchone())
-
- def describe(self):
- return "Creates extension %s" % self.name
-
- @property
- def migration_name_fragment(self):
- return "create_extension_%s" % self.name
-
-
- class BloomExtension(CreateExtension):
- def __init__(self):
- self.name = "bloom"
-
-
- 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)
-
-
- class CollationOperation(Operation):
- def __init__(self, name, locale, *, provider="libc", deterministic=True):
- self.name = name
- self.locale = locale
- self.provider = provider
- self.deterministic = deterministic
-
- def state_forwards(self, app_label, state):
- pass
-
- def deconstruct(self):
- kwargs = {"name": self.name, "locale": self.locale}
- if self.provider and self.provider != "libc":
- kwargs["provider"] = self.provider
- if self.deterministic is False:
- kwargs["deterministic"] = self.deterministic
- return (
- self.__class__.__qualname__,
- [],
- kwargs,
- )
-
- def create_collation(self, schema_editor):
- if self.deterministic is False and not (
- schema_editor.connection.features.supports_non_deterministic_collations
- ):
- raise NotSupportedError(
- "Non-deterministic collations require PostgreSQL 12+."
- )
- args = {"locale": schema_editor.quote_name(self.locale)}
- if self.provider != "libc":
- args["provider"] = schema_editor.quote_name(self.provider)
- if self.deterministic is False:
- args["deterministic"] = "false"
- schema_editor.execute(
- "CREATE COLLATION %(name)s (%(args)s)"
- % {
- "name": schema_editor.quote_name(self.name),
- "args": ", ".join(
- f"{option}={value}" for option, value in args.items()
- ),
- }
- )
-
- def remove_collation(self, schema_editor):
- schema_editor.execute(
- "DROP COLLATION %s" % schema_editor.quote_name(self.name),
- )
-
-
- class CreateCollation(CollationOperation):
- """Create a collation."""
-
- def database_forwards(self, app_label, schema_editor, from_state, to_state):
- if schema_editor.connection.vendor != "postgresql" or not router.allow_migrate(
- schema_editor.connection.alias, app_label
- ):
- return
- self.create_collation(schema_editor)
-
- def database_backwards(self, app_label, schema_editor, from_state, to_state):
- if not router.allow_migrate(schema_editor.connection.alias, app_label):
- return
- self.remove_collation(schema_editor)
-
- def describe(self):
- return f"Create collation {self.name}"
-
- @property
- def migration_name_fragment(self):
- return "create_collation_%s" % self.name.lower()
-
-
- class RemoveCollation(CollationOperation):
- """Remove a collation."""
-
- def database_forwards(self, app_label, schema_editor, from_state, to_state):
- if schema_editor.connection.vendor != "postgresql" or not router.allow_migrate(
- schema_editor.connection.alias, app_label
- ):
- return
- self.remove_collation(schema_editor)
-
- def database_backwards(self, app_label, schema_editor, from_state, to_state):
- if not router.allow_migrate(schema_editor.connection.alias, app_label):
- return
- self.create_collation(schema_editor)
-
- def describe(self):
- return f"Remove collation {self.name}"
-
- @property
- def migration_name_fragment(self):
- return "remove_collation_%s" % self.name.lower()
-
-
- class AddConstraintNotValid(AddConstraint):
- """
- Add a table constraint without enforcing validation, using PostgreSQL's
- NOT VALID syntax.
- """
-
- def __init__(self, model_name, constraint):
- if not isinstance(constraint, CheckConstraint):
- raise TypeError(
- "AddConstraintNotValid.constraint must be a check constraint."
- )
- super().__init__(model_name, constraint)
-
- def describe(self):
- return "Create not valid constraint %s on model %s" % (
- self.constraint.name,
- self.model_name,
- )
-
- def database_forwards(self, app_label, schema_editor, from_state, to_state):
- model = from_state.apps.get_model(app_label, self.model_name)
- if self.allow_migrate_model(schema_editor.connection.alias, model):
- constraint_sql = self.constraint.create_sql(model, schema_editor)
- if constraint_sql:
- # Constraint.create_sql returns interpolated SQL which makes
- # params=None a necessity to avoid escaping attempts on
- # execution.
- schema_editor.execute(str(constraint_sql) + " NOT VALID", params=None)
-
- @property
- def migration_name_fragment(self):
- return super().migration_name_fragment + "_not_valid"
-
-
- class ValidateConstraint(Operation):
- """Validate a table NOT VALID constraint."""
-
- def __init__(self, model_name, name):
- self.model_name = model_name
- self.name = name
-
- def describe(self):
- return "Validate constraint %s on model %s" % (self.name, self.model_name)
-
- def database_forwards(self, app_label, schema_editor, from_state, to_state):
- model = from_state.apps.get_model(app_label, self.model_name)
- if self.allow_migrate_model(schema_editor.connection.alias, model):
- schema_editor.execute(
- "ALTER TABLE %s VALIDATE CONSTRAINT %s"
- % (
- schema_editor.quote_name(model._meta.db_table),
- schema_editor.quote_name(self.name),
- )
- )
-
- def database_backwards(self, app_label, schema_editor, from_state, to_state):
- # PostgreSQL does not provide a way to make a constraint invalid.
- pass
-
- def state_forwards(self, app_label, state):
- pass
-
- @property
- def migration_name_fragment(self):
- return "%s_validate_%s" % (self.model_name.lower(), self.name.lower())
-
- def deconstruct(self):
- return (
- self.__class__.__name__,
- [],
- {
- "model_name": self.model_name,
- "name": self.name,
- },
- )
|