From 4db034c39b39b3fa64e55efec8375c8b26a169aa Mon Sep 17 00:00:00 2001 From: Andrew Date: Thu, 27 Jan 2022 22:06:44 -0500 Subject: [PATCH 1/2] Fix revert. Somehow migrations file was destroyed --- app/classes/shared/migration.py | 450 ++++++++++++++++++++++++++++++++ main.py | 1 - 2 files changed, 450 insertions(+), 1 deletion(-) diff --git a/app/classes/shared/migration.py b/app/classes/shared/migration.py index e69de29b..8d25a0fe 100644 --- a/app/classes/shared/migration.py +++ b/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, '', '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)) \ No newline at end of file diff --git a/main.py b/main.py index 2c4ab085..e86c5153 100644 --- a/main.py +++ b/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 From 350ffcba6302f6cd1822151e959752dd5c960164 Mon Sep 17 00:00:00 2001 From: Andrew Date: Thu, 27 Jan 2022 22:13:01 -0500 Subject: [PATCH 2/2] Lint ignore migrations --- app/classes/shared/migration.py | 1 + 1 file changed, 1 insertion(+) diff --git a/app/classes/shared/migration.py b/app/classes/shared/migration.py index 8d25a0fe..f200f1c7 100644 --- a/app/classes/shared/migration.py +++ b/app/classes/shared/migration.py @@ -1,3 +1,4 @@ +# pylint: skip-file from datetime import datetime import logging import typing as t