from models import Model from fields import DateField, StringField from engines import MergeTree from utils import escape from itertools import izip import logging logger = logging.getLogger('migrations') class Operation(object): ''' Base class for migration operations. ''' def apply(self, database): raise NotImplementedError() class CreateTable(Operation): ''' A migration operation that creates a table for a given model class. ''' def __init__(self, model_class): self.model_class = model_class def apply(self, database): logger.info(' Create table %s', self.model_class.table_name()) database.create_table(self.model_class) class AlterTable(Operation): ''' A migration operation that compares the table of a given model class to the model's fields, and alters the table to match the model. The operation can: - add new columns - drop obsolete columns - modify column types Default values are not altered by this operation. ''' def __init__(self, model_class): self.model_class = model_class def _get_table_fields(self, database): query = "DESC `%s`.`%s`" % (database.db_name, self.model_class.table_name()) return [(row.name, row.type) for row in database.select(query)] def _alter_table(self, database, cmd): cmd = "ALTER TABLE `%s`.`%s` %s" % (database.db_name, self.model_class.table_name(), cmd) logger.debug(cmd) database._send(cmd) def apply(self, database): logger.info(' Alter table %s', self.model_class.table_name()) table_fields = dict(self._get_table_fields(database)) # Identify fields that were deleted from the model deleted_fields = set(table_fields.keys()) - set(name for name, field in self.model_class._fields) for name in deleted_fields: logger.info(' Drop column %s', name) self._alter_table(database, 'DROP COLUMN %s' % name) del table_fields[name] # Identify fields that were added to the model prev_name = None for name, field in self.model_class._fields: if name not in table_fields: logger.info(' Add column %s', name) assert prev_name, 'Cannot add a column to the beginning of the table' default = field.get_db_prep_value(field.default) cmd = 'ADD COLUMN %s %s DEFAULT %s AFTER %s' % (name, field.db_type, escape(default), prev_name) self._alter_table(database, cmd) prev_name = name # Identify fields whose type was changed model_fields = [(name, field.db_type) for name, field in self.model_class._fields] for model_field, table_field in izip(model_fields, self._get_table_fields(database)): assert model_field[0] == table_field[0], 'Model fields and table columns in disagreement' if model_field[1] != table_field[1]: logger.info(' Change type of column %s from %s to %s', table_field[0], table_field[1], model_field[1]) self._alter_table(database, 'MODIFY COLUMN %s %s' % model_field) class DropTable(Operation): ''' A migration operation that drops the table of a given model class. ''' def __init__(self, model_class): self.model_class = model_class def apply(self, database): logger.info(' Drop table %s', self.model_class.__name__) database.drop_table(self.model_class) class MigrationHistory(Model): ''' A model for storing which migrations were already applied to the containing database. ''' package_name = StringField() module_name = StringField() applied = DateField() engine = MergeTree('applied', ('package_name', 'module_name')) @classmethod def table_name(cls): return 'infi_clickhouse_orm_migrations'