operations.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344
  1. from django.contrib.postgres.signals import (
  2. get_citext_oids,
  3. get_hstore_oids,
  4. register_type_handlers,
  5. )
  6. from django.db import NotSupportedError, router
  7. from django.db.migrations import AddConstraint, AddIndex, RemoveIndex
  8. from django.db.migrations.operations.base import Operation, OperationCategory
  9. from django.db.models.constraints import CheckConstraint
  10. class CreateExtension(Operation):
  11. reversible = True
  12. category = OperationCategory.ADDITION
  13. def __init__(self, name):
  14. self.name = name
  15. def state_forwards(self, app_label, state):
  16. pass
  17. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  18. if schema_editor.connection.vendor != "postgresql" or not router.allow_migrate(
  19. schema_editor.connection.alias, app_label
  20. ):
  21. return
  22. if not self.extension_exists(schema_editor, self.name):
  23. schema_editor.execute(
  24. "CREATE EXTENSION IF NOT EXISTS %s"
  25. % schema_editor.quote_name(self.name)
  26. )
  27. # Clear cached, stale oids.
  28. get_hstore_oids.cache_clear()
  29. get_citext_oids.cache_clear()
  30. # Registering new type handlers cannot be done before the extension is
  31. # installed, otherwise a subsequent data migration would use the same
  32. # connection.
  33. register_type_handlers(schema_editor.connection)
  34. if hasattr(schema_editor.connection, "register_geometry_adapters"):
  35. schema_editor.connection.register_geometry_adapters(
  36. schema_editor.connection.connection, True
  37. )
  38. def database_backwards(self, app_label, schema_editor, from_state, to_state):
  39. if not router.allow_migrate(schema_editor.connection.alias, app_label):
  40. return
  41. if self.extension_exists(schema_editor, self.name):
  42. schema_editor.execute(
  43. "DROP EXTENSION IF EXISTS %s" % schema_editor.quote_name(self.name)
  44. )
  45. # Clear cached, stale oids.
  46. get_hstore_oids.cache_clear()
  47. get_citext_oids.cache_clear()
  48. def extension_exists(self, schema_editor, extension):
  49. with schema_editor.connection.cursor() as cursor:
  50. cursor.execute(
  51. "SELECT 1 FROM pg_extension WHERE extname = %s",
  52. [extension],
  53. )
  54. return bool(cursor.fetchone())
  55. def describe(self):
  56. return "Creates extension %s" % self.name
  57. @property
  58. def migration_name_fragment(self):
  59. return "create_extension_%s" % self.name
  60. class BloomExtension(CreateExtension):
  61. def __init__(self):
  62. self.name = "bloom"
  63. class BtreeGinExtension(CreateExtension):
  64. def __init__(self):
  65. self.name = "btree_gin"
  66. class BtreeGistExtension(CreateExtension):
  67. def __init__(self):
  68. self.name = "btree_gist"
  69. class CITextExtension(CreateExtension):
  70. def __init__(self):
  71. self.name = "citext"
  72. class CryptoExtension(CreateExtension):
  73. def __init__(self):
  74. self.name = "pgcrypto"
  75. class HStoreExtension(CreateExtension):
  76. def __init__(self):
  77. self.name = "hstore"
  78. class TrigramExtension(CreateExtension):
  79. def __init__(self):
  80. self.name = "pg_trgm"
  81. class UnaccentExtension(CreateExtension):
  82. def __init__(self):
  83. self.name = "unaccent"
  84. class NotInTransactionMixin:
  85. def _ensure_not_in_transaction(self, schema_editor):
  86. if schema_editor.connection.in_atomic_block:
  87. raise NotSupportedError(
  88. "The %s operation cannot be executed inside a transaction "
  89. "(set atomic = False on the migration)." % self.__class__.__name__
  90. )
  91. class AddIndexConcurrently(NotInTransactionMixin, AddIndex):
  92. """Create an index using PostgreSQL's CREATE INDEX CONCURRENTLY syntax."""
  93. atomic = False
  94. category = OperationCategory.ADDITION
  95. def describe(self):
  96. return "Concurrently create index %s on field(s) %s of model %s" % (
  97. self.index.name,
  98. ", ".join(self.index.fields),
  99. self.model_name,
  100. )
  101. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  102. self._ensure_not_in_transaction(schema_editor)
  103. model = to_state.apps.get_model(app_label, self.model_name)
  104. if self.allow_migrate_model(schema_editor.connection.alias, model):
  105. schema_editor.add_index(model, self.index, concurrently=True)
  106. def database_backwards(self, app_label, schema_editor, from_state, to_state):
  107. self._ensure_not_in_transaction(schema_editor)
  108. model = from_state.apps.get_model(app_label, self.model_name)
  109. if self.allow_migrate_model(schema_editor.connection.alias, model):
  110. schema_editor.remove_index(model, self.index, concurrently=True)
  111. class RemoveIndexConcurrently(NotInTransactionMixin, RemoveIndex):
  112. """Remove an index using PostgreSQL's DROP INDEX CONCURRENTLY syntax."""
  113. atomic = False
  114. category = OperationCategory.REMOVAL
  115. def describe(self):
  116. return "Concurrently remove index %s from %s" % (self.name, self.model_name)
  117. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  118. self._ensure_not_in_transaction(schema_editor)
  119. model = from_state.apps.get_model(app_label, self.model_name)
  120. if self.allow_migrate_model(schema_editor.connection.alias, model):
  121. from_model_state = from_state.models[app_label, self.model_name_lower]
  122. index = from_model_state.get_index_by_name(self.name)
  123. schema_editor.remove_index(model, index, concurrently=True)
  124. def database_backwards(self, app_label, schema_editor, from_state, to_state):
  125. self._ensure_not_in_transaction(schema_editor)
  126. model = to_state.apps.get_model(app_label, self.model_name)
  127. if self.allow_migrate_model(schema_editor.connection.alias, model):
  128. to_model_state = to_state.models[app_label, self.model_name_lower]
  129. index = to_model_state.get_index_by_name(self.name)
  130. schema_editor.add_index(model, index, concurrently=True)
  131. class CollationOperation(Operation):
  132. def __init__(self, name, locale, *, provider="libc", deterministic=True):
  133. self.name = name
  134. self.locale = locale
  135. self.provider = provider
  136. self.deterministic = deterministic
  137. def state_forwards(self, app_label, state):
  138. pass
  139. def deconstruct(self):
  140. kwargs = {"name": self.name, "locale": self.locale}
  141. if self.provider and self.provider != "libc":
  142. kwargs["provider"] = self.provider
  143. if self.deterministic is False:
  144. kwargs["deterministic"] = self.deterministic
  145. return (
  146. self.__class__.__qualname__,
  147. [],
  148. kwargs,
  149. )
  150. def create_collation(self, schema_editor):
  151. args = {"locale": schema_editor.quote_name(self.locale)}
  152. if self.provider != "libc":
  153. args["provider"] = schema_editor.quote_name(self.provider)
  154. if self.deterministic is False:
  155. args["deterministic"] = "false"
  156. schema_editor.execute(
  157. "CREATE COLLATION %(name)s (%(args)s)"
  158. % {
  159. "name": schema_editor.quote_name(self.name),
  160. "args": ", ".join(
  161. f"{option}={value}" for option, value in args.items()
  162. ),
  163. }
  164. )
  165. def remove_collation(self, schema_editor):
  166. schema_editor.execute(
  167. "DROP COLLATION %s" % schema_editor.quote_name(self.name),
  168. )
  169. class CreateCollation(CollationOperation):
  170. """Create a collation."""
  171. category = OperationCategory.ADDITION
  172. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  173. if schema_editor.connection.vendor != "postgresql" or not router.allow_migrate(
  174. schema_editor.connection.alias, app_label
  175. ):
  176. return
  177. self.create_collation(schema_editor)
  178. def database_backwards(self, app_label, schema_editor, from_state, to_state):
  179. if not router.allow_migrate(schema_editor.connection.alias, app_label):
  180. return
  181. self.remove_collation(schema_editor)
  182. def describe(self):
  183. return f"Create collation {self.name}"
  184. @property
  185. def migration_name_fragment(self):
  186. return "create_collation_%s" % self.name.lower()
  187. class RemoveCollation(CollationOperation):
  188. """Remove a collation."""
  189. category = OperationCategory.REMOVAL
  190. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  191. if schema_editor.connection.vendor != "postgresql" or not router.allow_migrate(
  192. schema_editor.connection.alias, app_label
  193. ):
  194. return
  195. self.remove_collation(schema_editor)
  196. def database_backwards(self, app_label, schema_editor, from_state, to_state):
  197. if not router.allow_migrate(schema_editor.connection.alias, app_label):
  198. return
  199. self.create_collation(schema_editor)
  200. def describe(self):
  201. return f"Remove collation {self.name}"
  202. @property
  203. def migration_name_fragment(self):
  204. return "remove_collation_%s" % self.name.lower()
  205. class AddConstraintNotValid(AddConstraint):
  206. """
  207. Add a table constraint without enforcing validation, using PostgreSQL's
  208. NOT VALID syntax.
  209. """
  210. category = OperationCategory.ADDITION
  211. def __init__(self, model_name, constraint):
  212. if not isinstance(constraint, CheckConstraint):
  213. raise TypeError(
  214. "AddConstraintNotValid.constraint must be a check constraint."
  215. )
  216. super().__init__(model_name, constraint)
  217. def describe(self):
  218. return "Create not valid constraint %s on model %s" % (
  219. self.constraint.name,
  220. self.model_name,
  221. )
  222. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  223. model = from_state.apps.get_model(app_label, self.model_name)
  224. if self.allow_migrate_model(schema_editor.connection.alias, model):
  225. constraint_sql = self.constraint.create_sql(model, schema_editor)
  226. if constraint_sql:
  227. # Constraint.create_sql returns interpolated SQL which makes
  228. # params=None a necessity to avoid escaping attempts on
  229. # execution.
  230. schema_editor.execute(str(constraint_sql) + " NOT VALID", params=None)
  231. @property
  232. def migration_name_fragment(self):
  233. return super().migration_name_fragment + "_not_valid"
  234. class ValidateConstraint(Operation):
  235. """Validate a table NOT VALID constraint."""
  236. category = OperationCategory.ALTERATION
  237. def __init__(self, model_name, name):
  238. self.model_name = model_name
  239. self.name = name
  240. def describe(self):
  241. return "Validate constraint %s on model %s" % (self.name, self.model_name)
  242. def database_forwards(self, app_label, schema_editor, from_state, to_state):
  243. model = from_state.apps.get_model(app_label, self.model_name)
  244. if self.allow_migrate_model(schema_editor.connection.alias, model):
  245. schema_editor.execute(
  246. "ALTER TABLE %s VALIDATE CONSTRAINT %s"
  247. % (
  248. schema_editor.quote_name(model._meta.db_table),
  249. schema_editor.quote_name(self.name),
  250. )
  251. )
  252. def database_backwards(self, app_label, schema_editor, from_state, to_state):
  253. # PostgreSQL does not provide a way to make a constraint invalid.
  254. pass
  255. def state_forwards(self, app_label, state):
  256. pass
  257. @property
  258. def migration_name_fragment(self):
  259. return "%s_validate_%s" % (self.model_name.lower(), self.name.lower())
  260. def deconstruct(self):
  261. return (
  262. self.__class__.__name__,
  263. [],
  264. {
  265. "model_name": self.model_name,
  266. "name": self.name,
  267. },
  268. )