76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
from .loader import MigrationLoader
|
|
from .recorder import MigrationRecorder
|
|
|
|
|
|
class MigrationExecutor(object):
|
|
"""
|
|
End-to-end migration execution - loads migrations, and runs them
|
|
up or down to a specified set of targets.
|
|
"""
|
|
|
|
def __init__(self, connection, progress_callback=None):
|
|
self.connection = connection
|
|
self.loader = MigrationLoader(self.connection)
|
|
self.loader.load_disk()
|
|
self.recorder = MigrationRecorder(self.connection)
|
|
self.progress_callback = progress_callback
|
|
|
|
def migration_plan(self, targets):
|
|
"""
|
|
Given a set of targets, returns a list of (Migration instance, backwards?).
|
|
"""
|
|
plan = []
|
|
applied = self.recorder.applied_migrations()
|
|
for target in targets:
|
|
# If the migration is already applied, do backwards mode,
|
|
# otherwise do forwards mode.
|
|
if target in applied:
|
|
for migration in self.loader.graph.backwards_plan(target)[:-1]:
|
|
if migration in applied:
|
|
plan.append((self.loader.graph.nodes[migration], True))
|
|
applied.remove(migration)
|
|
else:
|
|
for migration in self.loader.graph.forwards_plan(target):
|
|
if migration not in applied:
|
|
plan.append((self.loader.graph.nodes[migration], False))
|
|
applied.add(migration)
|
|
return plan
|
|
|
|
def migrate(self, targets, plan=None):
|
|
"""
|
|
Migrates the database up to the given targets.
|
|
"""
|
|
if plan is None:
|
|
plan = self.migration_plan(targets)
|
|
for migration, backwards in plan:
|
|
if not backwards:
|
|
self.apply_migration(migration)
|
|
else:
|
|
self.unapply_migration(migration)
|
|
|
|
def apply_migration(self, migration):
|
|
"""
|
|
Runs a migration forwards.
|
|
"""
|
|
if self.progress_callback:
|
|
self.progress_callback("apply_start", migration)
|
|
with self.connection.schema_editor() as schema_editor:
|
|
project_state = self.loader.graph.project_state((migration.app_label, migration.name), at_end=False)
|
|
migration.apply(project_state, schema_editor)
|
|
self.recorder.record_applied(migration.app_label, migration.name)
|
|
if self.progress_callback:
|
|
self.progress_callback("apply_success", migration)
|
|
|
|
def unapply_migration(self, migration):
|
|
"""
|
|
Runs a migration backwards.
|
|
"""
|
|
if self.progress_callback:
|
|
self.progress_callback("unapply_start", migration)
|
|
with self.connection.schema_editor() as schema_editor:
|
|
project_state = self.loader.graph.project_state((migration.app_label, migration.name), at_end=False)
|
|
migration.unapply(project_state, schema_editor)
|
|
self.recorder.record_unapplied(migration.app_label, migration.name)
|
|
if self.progress_callback:
|
|
self.progress_callback("unapply_success", migration)
|