56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
|
import os
|
||
|
|
|
||
|
|
from abc import ABCMeta, abstractmethod, abstractproperty
|
||
|
|
from collections import namedtuple
|
||
|
|
from six import add_metaclass
|
||
|
|
|
||
|
|
MigrationPhase = namedtuple('MigrationPhase', ['name', 'alembic_revision', 'flags'])
|
||
|
|
|
||
|
|
|
||
|
|
@add_metaclass(ABCMeta)
|
||
|
|
class DataMigration(object):
|
||
|
|
@abstractproperty
|
||
|
|
def alembic_migration_revision(self):
|
||
|
|
""" Returns the alembic migration revision corresponding to the currently configured phase.
|
||
|
|
"""
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def has_flag(self, flag):
|
||
|
|
""" Returns true if the data migration's current phase has the given flag set. """
|
||
|
|
|
||
|
|
|
||
|
|
class NullDataMigration(DataMigration):
|
||
|
|
@property
|
||
|
|
def alembic_migration_revision(self):
|
||
|
|
return 'head'
|
||
|
|
|
||
|
|
def has_flag(self, flag):
|
||
|
|
raise NotImplementedError()
|
||
|
|
|
||
|
|
|
||
|
|
class DefinedDataMigration(DataMigration):
|
||
|
|
def __init__(self, name, env_var, phases):
|
||
|
|
self.name = name
|
||
|
|
self.phases = {phase.name: phase for phase in phases}
|
||
|
|
|
||
|
|
phase_name = os.getenv(env_var)
|
||
|
|
if phase_name is None:
|
||
|
|
msg = 'Missing env var `%s` for data migration `%s`' % (env_var, self.name)
|
||
|
|
raise Exception(msg)
|
||
|
|
|
||
|
|
current_phase = self.phases.get(phase_name)
|
||
|
|
if current_phase is None:
|
||
|
|
msg = 'Unknown phase `%s` for data migration `%s`' % (phase_name, self.name)
|
||
|
|
raise Exception(msg)
|
||
|
|
|
||
|
|
self.current_phase = current_phase
|
||
|
|
|
||
|
|
@property
|
||
|
|
def alembic_migration_revision(self):
|
||
|
|
assert self.current_phase
|
||
|
|
return self.current_phase.alembic_revision
|
||
|
|
|
||
|
|
def has_flag(self, flag):
|
||
|
|
assert self.current_phase
|
||
|
|
return flag in self.current_phase.flags
|