infi.clickhouse_orm/src/infi/clickhouse_orm/migrations.py

import six

from .models import Model, BufferModel
from .fields import DateField, StringField
from .engines import MergeTree
from .utils import escape

from six.moves import zip
from six import iteritems

import logging
logger = logging.getLogger('migrations')


class Operation(object):
    '''
    Base class for migration operations.
    '''

    def apply(self, database):
        raise NotImplementedError()   # pragma: no cover


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())
        if issubclass(self.model_class, BufferModel):
            database.create_table(self.model_class.engine.main_model)
        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())

        # Note that MATERIALIZED and ALIAS fields are always at the end of the DESC,
        # ADD COLUMN ... AFTER doesn't affect it
        table_fields = dict(self._get_table_fields(database))

        # Identify fields that were deleted from the model
        deleted_fields = set(table_fields.keys()) - set(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 iteritems(self.model_class.fields()):
            is_regular_field = not (field.materialized or field.alias)
            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'
                cmd = 'ADD COLUMN %s %s' % (name, field.get_sql())
                if is_regular_field:
                    cmd += ' AFTER %s' % prev_name
                self._alter_table(database, cmd)

            if is_regular_field:
                # ALIAS and MATERIALIZED fields are not stored in the database, and raise DatabaseError
                # (no AFTER column). So we will skip them
                prev_name = name

        # Identify fields whose type was changed
        # The order of class attributes can be changed any time, so we can't count on it
        # Secondly, MATERIALIZED and ALIAS fields are always at the end of the DESC, so we can't expect them to save
        # attribute position. Watch https://github.com/Infinidat/infi.clickhouse_orm/issues/47
        model_fields = {name: field.get_sql(with_default_expression=False)
                        for name, field in iteritems(self.model_class.fields())}
        for field_name, field_sql in self._get_table_fields(database):
            # All fields must have been created and dropped by this moment
            assert field_name in model_fields, 'Model fields and table columns in disagreement'

            if field_sql != model_fields[field_name]:
                logger.info('        Change type of column %s from %s to %s', field_name, field_sql,
                            model_fields[field_name])
                self._alter_table(database, 'MODIFY COLUMN %s %s' % (field_name, model_fields[field_name]))


class AlterTableWithBuffer(Operation):
    '''
    A migration operation for altering a buffer table and its underlying on-disk table.
    The buffer table is dropped, the on-disk table is altered, and then the buffer table
    is re-created.
    '''

    def __init__(self, model_class):
        self.model_class = model_class

    def apply(self, database):
        if issubclass(self.model_class, BufferModel):
            DropTable(self.model_class).apply(database)
            AlterTable(self.model_class.engine.main_model).apply(database)
            CreateTable(self.model_class).apply(database)
        else:
            AlterTable(self.model_class).apply(database)


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.table_name())
        database.drop_table(self.model_class)


class RunPython(Operation):
    '''
    A migration operation that executes given python function on database
    '''
    def __init__(self, func):
        assert callable(func), "'func' parameter must be function"
        self._func = func

    def apply(self, database):
        logger.info('    Executing python operation %s', self._func.__name__)
        self._func(database)


class RunSQL(Operation):
    '''
    A migration operation that executes given SQL on database
    '''

    def __init__(self, sql):
        if isinstance(sql, six.string_types):
            sql = [sql]

        assert isinstance(sql, list), "'sql' parameter must be string or list of strings"
        self._sql = sql

    def apply(self, database):
        logger.info('    Executing raw SQL operations')
        for item in self._sql:
            database.raw(item)


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'