vorta/src/vorta/store/migrations.py

from datetime import datetime
import peewee as pw
from playhouse.migrate import SqliteMigrator, migrate
from .models import (DB, RepoModel, BackupProfileModel, SourceFileModel,
                     SettingsModel, ArchiveModel, WifiSettingModel, EventLogModel)


def run_migrations(current_schema, db_connection):
    """
    Apply new schema versions to database.

    See http://docs.peewee-orm.com/en/latest/peewee/playhouse.html#schema-migrations
    """
    migrator = SqliteMigrator(db_connection)

    if current_schema.version < 4:  # version 3 to 4
        _apply_schema_update(
            current_schema, 4,
            migrator.add_column(ArchiveModel._meta.table_name, 'duration', pw.FloatField(null=True)),
            migrator.add_column(ArchiveModel._meta.table_name, 'size', pw.IntegerField(null=True))
        )
    if current_schema.version < 5:
        _apply_schema_update(
            current_schema, 5,
            migrator.drop_not_null(WifiSettingModel._meta.table_name, 'last_connected'),
        )

    if current_schema.version < 6:
        _apply_schema_update(
            current_schema, 6,
            migrator.add_column(EventLogModel._meta.table_name, 'repo_url', pw.CharField(null=True))
        )

    if current_schema.version < 7:
        _apply_schema_update(
            current_schema, 7,
            migrator.rename_column(SourceFileModel._meta.table_name, 'config_id', 'profile_id'),
            migrator.drop_column(EventLogModel._meta.table_name, 'profile_id'),
            migrator.add_column(EventLogModel._meta.table_name, 'profile', pw.CharField(null=True))
        )

    if current_schema.version < 8:
        _apply_schema_update(
            current_schema, 8,
            migrator.add_column(BackupProfileModel._meta.table_name,
                                'prune_keep_within', pw.CharField(null=True)))

    if current_schema.version < 9:
        _apply_schema_update(
            current_schema, 9,
            migrator.add_column(BackupProfileModel._meta.table_name, 'new_archive_name',
                                pw.CharField(default="{hostname}-{profile_slug}-{now:%Y-%m-%dT%H:%M:%S}")),
            migrator.add_column(BackupProfileModel._meta.table_name, 'prune_prefix',
                                pw.CharField(default="{hostname}-{profile_slug}-")),
        )

    if current_schema.version < 10:
        _apply_schema_update(
            current_schema, 10,
            migrator.add_column(BackupProfileModel._meta.table_name, 'pre_backup_cmd',
                                pw.CharField(default='')),
            migrator.add_column(BackupProfileModel._meta.table_name, 'post_backup_cmd',
                                pw.CharField(default='')),
        )

    if current_schema.version < 11:
        _apply_schema_update(current_schema, 11)
        for profile in BackupProfileModel:
            if profile.compression == 'zstd':
                profile.compression = 'zstd,3'
            if profile.compression == 'lzma,6':
                profile.compression = 'auto,lzma,6'
            profile.save()

    if current_schema.version < 12:
        _apply_schema_update(
            current_schema, 12,
            migrator.add_column(RepoModel._meta.table_name,
                                'extra_borg_arguments', pw.CharField(default='')))

    if current_schema.version < 13:
        # Migrate ArchiveModel data to new table to remove unique constraint from snapshot_id column.
        tables = DB.get_tables()
        if ArchiveModel.select().count() == 0 and 'snapshotmodel' in tables:
            cursor = DB.execute_sql('select * from snapshotmodel;')
            fields = [ArchiveModel.id, ArchiveModel.snapshot_id, ArchiveModel.name, ArchiveModel.repo,
                      ArchiveModel.time, ArchiveModel.duration, ArchiveModel.size]
            data = [row for row in cursor.fetchall()]
            with DB.atomic():
                size = 1000
                for i in range(0, len(data), size):
                    ArchiveModel.insert_many(data[i: i + size], fields=fields).execute()

        _apply_schema_update(current_schema, 13)

    if current_schema.version < 14:
        _apply_schema_update(
            current_schema, 14,
            migrator.add_column(SettingsModel._meta.table_name,
                                'str_value', pw.CharField(default='')))

    if current_schema.version < 15:
        _apply_schema_update(
            current_schema, 15,
            migrator.add_column(BackupProfileModel._meta.table_name,
                                'dont_run_on_metered_networks', pw.BooleanField(default=True))
        )

    if current_schema.version < 16:
        _apply_schema_update(
            current_schema, 16,
            migrator.add_column(SourceFileModel._meta.table_name,
                                'dir_size', pw.BigIntegerField(default=-1)),
            migrator.add_column(SourceFileModel._meta.table_name,
                                'dir_files_count', pw.BigIntegerField(default=-1)),
            migrator.add_column(SourceFileModel._meta.table_name,
                                'path_isdir', pw.BooleanField(default=False))
        )

    if current_schema.version < 17:
        _apply_schema_update(
            current_schema, 17,
            migrator.add_column(RepoModel._meta.table_name,
                                'create_backup_cmd', pw.CharField(default=''))
        )

    if current_schema.version < 18:
        _apply_schema_update(
            current_schema, 18,
            migrator.add_column(BackupProfileModel._meta.table_name,
                                'schedule_interval_unit', pw.CharField(default='hours')),
            migrator.add_column(BackupProfileModel._meta.table_name,
                                'schedule_interval_count', pw.IntegerField(default=3)),
            migrator.add_column(BackupProfileModel._meta.table_name,
                                'schedule_make_up_missed', pw.BooleanField(default=False)),
            migrator.add_column(EventLogModel._meta.table_name,
                                'end_time', pw.DateTimeField(default=datetime.now))
        )


def _apply_schema_update(current_schema, version_after, *operations):
    with DB.atomic():
        migrate(*operations)
        current_schema.version = version_after
        current_schema.changed_at = datetime.now()
        current_schema.save()