# pylint: skip-file from datetime import datetime import logging import typing as t import sys import os import re from functools import wraps from functools import cached_property from app.classes.shared.helpers import helper from app.classes.shared.console import console logger = logging.getLogger(__name__) try: import peewee from playhouse.migrate import ( SqliteMigrator, Operation, SQL, SqliteDatabase, make_index_name ) except ModuleNotFoundError as e: logger.critical("Import Error: Unable to load {} module".format( e.name), exc_info=True) console.critical("Import Error: Unable to load {} module".format(e.name)) sys.exit(1) MIGRATE_TABLE = 'migratehistory' MIGRATE_TEMPLATE = '''# Generated by database migrator import peewee def migrate(migrator, db): """ Write your migrations here. """ {migrate} def rollback(migrator, db): """ Write your rollback migrations here. """ {rollback}''' class MigrateHistory(peewee.Model): """ Presents the migration history in a database. """ name = peewee.CharField(unique=True) migrated_at = peewee.DateTimeField(default=datetime.utcnow) # noinspection PyTypeChecker def __unicode__(self) -> str: """ String representation of this migration """ return self.name class Meta: table_name = MIGRATE_TABLE def get_model(method): """ Convert string to model class. """ @wraps(method) def wrapper(migrator, model, *args, **kwargs): if isinstance(model, str): return method(migrator, migrator.table_dict[model], *args, **kwargs) return method(migrator, model, *args, **kwargs) return wrapper # noinspection PyProtectedMember class Migrator(object): def __init__(self, database: t.Union[peewee.Database, peewee.Proxy]): """ Initializes the migrator """ if isinstance(database, peewee.Proxy): database = database.obj self.database: SqliteDatabase = database self.table_dict: t.Dict[str, peewee.Model] = {} self.operations: t.List[t.Union[Operation, callable]] = [] self.migrator = SqliteMigrator(database) def run(self): """ Runs operations. """ for op in self.operations: if isinstance(op, Operation): op.run() else: op() self.clean() def clean(self): """ Cleans the operations. """ self.operations = list() def sql(self, sql: str, *params): """ Executes raw SQL. """ self.operations.append(SQL(sql, *params)) def create_table(self, model: peewee.Model) -> peewee.Model: """ Creates model and table in database. """ self.table_dict[model._meta.table_name] = model model._meta.database = self.database self.operations.append(model.create_table) return model @get_model def drop_table(self, model: peewee.Model): """ Drops model and table from database. """ del self.table_dict[model._meta.table_name] self.operations.append(lambda: model.drop_table(cascade=False)) @get_model def add_columns(self, model: peewee.Model, **fields: peewee.Field) -> peewee.Model: """ Creates new fields. """ for name, field in fields.items(): model._meta.add_field(name, field) self.operations.append(self.migrator.add_column( model._meta.table_name, field.column_name, field)) if field.unique: self.operations.append(self.migrator.add_index( model._meta.table_name, (field.column_name,), unique=True)) return model @get_model def drop_columns(self, model: peewee.Model, names: str) -> peewee.Model: """ Removes fields from model. """ fields = [field for field in model._meta.fields.values() if field.name in names] for field in fields: self.__del_field__(model, field) if field.unique: # Drop unique index index_name = make_index_name( model._meta.table_name, [field.column_name]) self.operations.append(self.migrator.drop_index( model._meta.table_name, index_name)) self.operations.append( self.migrator.drop_column( model._meta.table_name, field.column_name, cascade=False)) return model def __del_field__(self, model: peewee.Model, field: peewee.Field): """ Deletes field from model. """ model._meta.remove_field(field.name) delattr(model, field.name) if isinstance(field, peewee.ForeignKeyField): obj_id_name = field.column_name if field.column_name == field.name: obj_id_name += '_id' delattr(model, obj_id_name) delattr(field.rel_model, field.backref) @get_model def rename_column(self, model: peewee.Model, old_name: str, new_name: str) -> peewee.Model: """ Renames field in model. """ field = model._meta.fields[old_name] if isinstance(field, peewee.ForeignKeyField): old_name = field.column_name self.__del_field__(model, field) field.name = field.column_name = new_name model._meta.add_field(new_name, field) if isinstance(field, peewee.ForeignKeyField): field.column_name = new_name = field.column_name + '_id' self.operations.append(self.migrator.rename_column( model._meta.table_name, old_name, new_name)) return model @get_model def rename_table(self, model: peewee.Model, new_name: str) -> peewee.Model: """ Renames table in database. """ old_name = model._meta.table_name del self.table_dict[model._meta.table_name] model._meta.table_name = new_name self.table_dict[model._meta.table_name] = model self.operations.append(self.migrator.rename_table(old_name, new_name)) return model @get_model def add_index(self, model: peewee.Model, *columns: str, unique=False) -> peewee.Model: """Create indexes.""" model._meta.indexes.append((columns, unique)) columns_ = [] for col in columns: field = model._meta.fields.get(col) if len(columns) == 1: field.unique = unique field.index = not unique if isinstance(field, peewee.ForeignKeyField): col = col + '_id' columns_.append(col) self.operations.append(self.migrator.add_index( model._meta.table_name, columns_, unique=unique)) return model @get_model def drop_index(self, model: peewee.Model, *columns: str) -> peewee.Model: """Drop indexes.""" columns_ = [] for col in columns: field = model._meta.fields.get(col) if not field: continue if len(columns) == 1: field.unique = field.index = False if isinstance(field, peewee.ForeignKeyField): col = col + '_id' columns_.append(col) index_name = make_index_name(model._meta.table_name, columns_) model._meta.indexes = [(cols, _) for ( cols, _) in model._meta.indexes if columns != cols] self.operations.append(self.migrator.drop_index( model._meta.table_name, index_name)) return model @get_model def add_not_null(self, model: peewee.Model, *names: str) -> peewee.Model: """Add not null.""" for name in names: field = model._meta.fields[name] field.null = False self.operations.append(self.migrator.add_not_null( model._meta.table_name, field.column_name)) return model @get_model def drop_not_null(self, model: peewee.Model, *names: str) -> peewee.Model: """Drop not null.""" for name in names: field = model._meta.fields[name] field.null = True self.operations.append(self.migrator.drop_not_null( model._meta.table_name, field.column_name)) return model @get_model def add_default(self, model: peewee.Model, name: str, default: t.Any) -> peewee.Model: """Add default.""" field = model._meta.fields[name] model._meta.defaults[field] = field.default = default self.operations.append(self.migrator.apply_default( model._meta.table_name, name, field)) return model # noinspection PyProtectedMember class MigrationManager(object): filemask = re.compile(r"[\d]+_[^\.]+\.py$") def __init__(self, database: t.Union[peewee.Database, peewee.Proxy]): """ Initializes the migration manager. """ if not isinstance(database, (peewee.Database, peewee.Proxy)): raise RuntimeError('Invalid database: {}'.format(database)) self.database = database @cached_property def model(self) -> t.Type[MigrateHistory]: """ Initialize and cache the MigrationHistory model. """ MigrateHistory._meta.database = self.database MigrateHistory._meta.table_name = 'migratehistory' MigrateHistory._meta.schema = None MigrateHistory.create_table(True) return MigrateHistory @property def done(self) -> t.List[str]: """ Scans migrations in the database. """ return [mm.name for mm in self.model.select().order_by(self.model.id)] @property def todo(self): """ Scans migrations in the file system. """ if not os.path.exists(helper.migration_dir): logger.warning('Migration directory: {} does not exist.'.format( helper.migration_dir)) os.makedirs(helper.migration_dir) return sorted(f[:-3] for f in os.listdir(helper.migration_dir) if self.filemask.match(f)) @property def diff(self) -> t.List[str]: """ Calculates difference between the filesystem and the database. """ done = set(self.done) return [name for name in self.todo if name not in done] @cached_property def migrator(self) -> Migrator: """ Create migrator and setup it with fake migrations. """ migrator = Migrator(self.database) for name in self.done: self.up_one(name, migrator, True) return migrator def compile(self, name, migrate='', rollback=''): """ Compiles a migration. """ name = datetime.utcnow().strftime('%Y%m%d%H%M%S') + '_' + name filename = name + '.py' path = os.path.join(helper.migration_dir, filename) with open(path, 'w') as f: f.write(MIGRATE_TEMPLATE.format( migrate=migrate, rollback=rollback, name=filename)) return name def create(self, name: str = 'auto', auto: bool = False) -> t.Optional[str]: """ Creates a migration. """ migrate = rollback = '' if auto: raise NotImplementedError logger.info('Creating migration "{}"'.format(name)) name = self.compile(name, migrate, rollback) logger.info('Migration has been created as "{}"'.format(name)) return name def clear(self): """Clear migrations.""" self.model.delete().execute() def up(self, name: t.Optional[str] = None): """ Runs all unapplied migrations. """ logger.info('Starting migrations') console.info('Starting migrations') done = [] diff = self.diff if not diff: logger.info('There is nothing to migrate') console.info('There is nothing to migrate') return done migrator = self.migrator for mname in diff: done.append(self.up_one(mname, self.migrator)) if name and name == mname: break return done def read(self, name: str): """ Reads a migration from a file. """ call_params = dict() if helper.is_os_windows() and sys.version_info >= (3, 0): # if system is windows - force utf-8 encoding call_params['encoding'] = 'utf-8' with open(os.path.join(helper.migration_dir, name + '.py'), **call_params) as f: code = f.read() scope = {} code = compile(code, '', 'exec', dont_inherit=True) exec(code, scope, None) return scope.get('migrate', lambda m, d: None), scope.get('rollback', lambda m, d: None) def up_one(self, name: str, migrator: Migrator, fake: bool = False, rollback: bool = False) -> str: """ Runs a migration with a given name. """ try: migrate_fn, rollback_fn = self.read(name) if fake: migrate_fn(migrator, self.database) migrator.clean() return name with self.database.transaction(): if rollback: logger.info('Rolling back "{}"'.format(name)) rollback_fn(migrator, self.database) migrator.run() self.model.delete().where(self.model.name == name).execute() else: logger.info('Migrate "{}"'.format(name)) migrate_fn(migrator, self.database) migrator.run() if name not in self.done: self.model.create(name=name) logger.info('Done "{}"'.format(name)) return name except Exception: self.database.rollback() operation_name = 'Rollback' if rollback else 'Migration' logger.exception('{} failed: {}'.format(operation_name, name)) raise def down(self): """ Rolls back migrations. """ if not self.done: raise RuntimeError('No migrations are found.') name = self.done[-1] migrator = self.migrator self.up_one(name, migrator, False, True) logger.warning('Rolled back migration: {}'.format(name))