executor.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. from __future__ import unicode_literals
  2. from django.db import migrations
  3. from django.apps.registry import apps as global_apps
  4. from .loader import MigrationLoader
  5. from .recorder import MigrationRecorder
  6. class MigrationExecutor(object):
  7. """
  8. End-to-end migration execution - loads migrations, and runs them
  9. up or down to a specified set of targets.
  10. """
  11. def __init__(self, connection, progress_callback=None):
  12. self.connection = connection
  13. self.loader = MigrationLoader(self.connection)
  14. self.recorder = MigrationRecorder(self.connection)
  15. self.progress_callback = progress_callback
  16. def migration_plan(self, targets):
  17. """
  18. Given a set of targets, returns a list of (Migration instance, backwards?).
  19. """
  20. plan = []
  21. applied = set(self.loader.applied_migrations)
  22. for target in targets:
  23. # If the target is (app_label, None), that means unmigrate everything
  24. if target[1] is None:
  25. for root in self.loader.graph.root_nodes():
  26. if root[0] == target[0]:
  27. for migration in self.loader.graph.backwards_plan(root):
  28. if migration in applied:
  29. plan.append((self.loader.graph.nodes[migration], True))
  30. applied.remove(migration)
  31. # If the migration is already applied, do backwards mode,
  32. # otherwise do forwards mode.
  33. elif target in applied:
  34. backwards_plan = self.loader.graph.backwards_plan(target)[:-1]
  35. # We only do this if the migration is not the most recent one
  36. # in its app - that is, another migration with the same app
  37. # label is in the backwards plan
  38. if any(node[0] == target[0] for node in backwards_plan):
  39. for migration in backwards_plan:
  40. if migration in applied:
  41. plan.append((self.loader.graph.nodes[migration], True))
  42. applied.remove(migration)
  43. else:
  44. for migration in self.loader.graph.forwards_plan(target):
  45. if migration not in applied:
  46. plan.append((self.loader.graph.nodes[migration], False))
  47. applied.add(migration)
  48. return plan
  49. def migrate(self, targets, plan=None, fake=False):
  50. """
  51. Migrates the database up to the given targets.
  52. """
  53. if plan is None:
  54. plan = self.migration_plan(targets)
  55. for migration, backwards in plan:
  56. if not backwards:
  57. self.apply_migration(migration, fake=fake)
  58. else:
  59. self.unapply_migration(migration, fake=fake)
  60. def collect_sql(self, plan):
  61. """
  62. Takes a migration plan and returns a list of collected SQL
  63. statements that represent the best-efforts version of that plan.
  64. """
  65. statements = []
  66. for migration, backwards in plan:
  67. with self.connection.schema_editor(collect_sql=True) as schema_editor:
  68. project_state = self.loader.project_state((migration.app_label, migration.name), at_end=False)
  69. if not backwards:
  70. migration.apply(project_state, schema_editor, collect_sql=True)
  71. else:
  72. migration.unapply(project_state, schema_editor, collect_sql=True)
  73. statements.extend(schema_editor.collected_sql)
  74. return statements
  75. def apply_migration(self, migration, fake=False):
  76. """
  77. Runs a migration forwards.
  78. """
  79. if self.progress_callback:
  80. self.progress_callback("apply_start", migration, fake)
  81. if not fake:
  82. # Test to see if this is an already-applied initial migration
  83. if self.detect_soft_applied(migration):
  84. fake = True
  85. else:
  86. # Alright, do it normally
  87. with self.connection.schema_editor() as schema_editor:
  88. project_state = self.loader.project_state((migration.app_label, migration.name), at_end=False)
  89. migration.apply(project_state, schema_editor)
  90. # For replacement migrations, record individual statuses
  91. if migration.replaces:
  92. for app_label, name in migration.replaces:
  93. self.recorder.record_applied(app_label, name)
  94. else:
  95. self.recorder.record_applied(migration.app_label, migration.name)
  96. # Report progress
  97. if self.progress_callback:
  98. self.progress_callback("apply_success", migration, fake)
  99. def unapply_migration(self, migration, fake=False):
  100. """
  101. Runs a migration backwards.
  102. """
  103. if self.progress_callback:
  104. self.progress_callback("unapply_start", migration, fake)
  105. if not fake:
  106. with self.connection.schema_editor() as schema_editor:
  107. project_state = self.loader.project_state((migration.app_label, migration.name), at_end=False)
  108. migration.unapply(project_state, schema_editor)
  109. # For replacement migrations, record individual statuses
  110. if migration.replaces:
  111. for app_label, name in migration.replaces:
  112. self.recorder.record_unapplied(app_label, name)
  113. else:
  114. self.recorder.record_unapplied(migration.app_label, migration.name)
  115. # Report progress
  116. if self.progress_callback:
  117. self.progress_callback("unapply_success", migration, fake)
  118. def detect_soft_applied(self, migration):
  119. """
  120. Tests whether a migration has been implicitly applied - that the
  121. tables it would create exist. This is intended only for use
  122. on initial migrations (as it only looks for CreateModel).
  123. """
  124. project_state = self.loader.project_state((migration.app_label, migration.name), at_end=True)
  125. apps = project_state.render()
  126. found_create_migration = False
  127. # Bail if the migration isn't the first one in its app
  128. if [name for app, name in migration.dependencies if app == migration.app_label]:
  129. return False
  130. # Make sure all create model are done
  131. for operation in migration.operations:
  132. if isinstance(operation, migrations.CreateModel):
  133. model = apps.get_model(migration.app_label, operation.name)
  134. if model._meta.swapped:
  135. # We have to fetch the model to test with from the
  136. # main app cache, as it's not a direct dependency.
  137. model = global_apps.get_model(model._meta.swapped)
  138. if model._meta.db_table not in self.connection.introspection.get_table_list(self.connection.cursor()):
  139. return False
  140. found_create_migration = True
  141. # If we get this far and we found at least one CreateModel migration,
  142. # the migration is considered implicitly applied.
  143. return found_create_migration