Fix revert. Somehow migrations file was destroyed

app/classes/shared/migration.py

@@ -0,0 +1,450 @@
+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, '<string>', '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))

main.py

@@ -19,7 +19,6 @@ else:
 from app.classes.shared.main_models import installer, database
 from app.classes.shared.tasks import TasksManager
 from app.classes.shared.main_controller import Controller
-# pylint: disable=no-name-in-module
 from app.classes.shared.migration import MigrationManager
 
 from app.classes.shared.command import MainPrompt