crafty-4/app/migrations/20240216_rework_servers_uuid_part2.py

import datetime
import uuid
import peewee
import logging

from app.classes.shared.console import Console
from app.classes.shared.migration import Migrator, MigrateHistory
from app.classes.models.management import (
    AuditLog,
    Webhooks,
    Schedules,
    Backups,
)
from app.classes.models.server_permissions import RoleServers
from app.classes.models.servers import Servers

logger = logging.getLogger(__name__)


def migrate(migrator: Migrator, database, **kwargs):
    """
    Write your migrations here.
    """
    db = database

    this_migration = MigrateHistory.get_or_none(MigrateHistory.name == "20240216_rework_servers_uuid_part2")
    if this_migration is not None:
        Console.debug("Update database already done, skipping this part")
        return

    try:
        logger.info("Migrating Data from Int to UUID (Foreign Keys)")
        Console.info("Migrating Data from Int to UUID (Foreign Keys)")
        # Changes on Audit Log Table
        for audit_log in AuditLog.select():
            old_server_id = audit_log.server_id_id
            if old_server_id == "0" or old_server_id is None:
                server_uuid = None
            else:
                try:
                    server = Servers.get_by_id(old_server_id)
                    server_uuid = server.server_uuid
                except:
                    server_uuid = old_server_id
            AuditLog.update(server_id=server_uuid).where(
                AuditLog.audit_id == audit_log.audit_id
            ).execute()

        # Changes on Webhooks Log Table
        for webhook in Webhooks.select():
            old_server_id = webhook.server_id_id
            try:
                server = Servers.get_by_id(old_server_id)
                server_uuid = server.server_uuid
            except:
                server_uuid = old_server_id
            Webhooks.update(server_id=server_uuid).where(
                Webhooks.id == webhook.id
            ).execute()

        # Changes on Schedules Log Table
        for schedule in Schedules.select():
            old_server_id = schedule.server_id_id
            try:
                server = Servers.get_by_id(old_server_id)
                server_uuid = server.server_uuid
            except:
                server_uuid = old_server_id
            Schedules.update(server_id=server_uuid).where(
                Schedules.schedule_id == schedule.schedule_id
            ).execute()

        # Changes on Backups Log Table
        for backup in Backups.select():
            old_server_id = backup.server_id_id
            try:
                server = Servers.get_by_id(old_server_id)
                server_uuid = server.server_uuid
            except:
                server_uuid = old_server_id
            Backups.update(server_id=server_uuid).where(
                Backups.server_id == old_server_id
            ).execute()

        # Changes on RoleServers Log Table
        for role_servers in RoleServers.select():
            old_server_id = role_servers.server_id_id
            try:
                server = Servers.get_by_id(old_server_id)
                server_uuid = server.server_uuid
            except:
                server_uuid = old_server_id
            RoleServers.update(server_id=server_uuid).where(
                RoleServers.role_id == role_servers.id
                and RoleServers.server_id == old_server_id
            ).execute()

        logger.info("Migrating Data from Int to UUID (Foreign Keys) : SUCCESS")
        Console.info("Migrating Data from Int to UUID (Foreign Keys) : SUCCESS")

    except Exception as ex:
        logger.error("Error while migrating Data from Int to UUID (Foreign Keys)")
        logger.error(ex)
        Console.error("Error while migrating Data from Int to UUID (Foreign Keys)")
        Console.error(ex)
        last_migration = MigrateHistory.get_by_id(MigrateHistory.select().count())
        last_migration.delete()
        return

    try:
        logger.info("Migrating Data from Int to UUID (Primary Keys)")
        Console.info("Migrating Data from Int to UUID (Primary Keys)")
        # Migrating servers from the old id type to the new one
        for server in Servers.select():
            Servers.update(server_id=server.server_uuid).where(
                Servers.server_id == server.server_id
            ).execute()

        logger.info("Migrating Data from Int to UUID (Primary Keys) : SUCCESS")
        Console.info("Migrating Data from Int to UUID (Primary Keys) : SUCCESS")

    except Exception as ex:
        logger.error("Error while migrating Data from Int to UUID (Primary Keys)")
        logger.error(ex)
        Console.error("Error while migrating Data from Int to UUID (Primary Keys)")
        Console.error(ex)
        last_migration = MigrateHistory.get_by_id(MigrateHistory.select().count())
        last_migration.delete()
        return

    return


def rollback(migrator: Migrator, database, **kwargs):
    """
    Write your rollback migrations here.
    """
    db = database

    # Changes on Server Table
    migrator.alter_column_type(
        "servers",
        "server_id",
        peewee.AutoField(),
    )

    # Changes on Audit Log Table
    migrator.alter_column_type(
        AuditLog,
        "server_id",
        peewee.IntegerField(default=None, index=True),
    )

    # Changes on Webhook Table
    migrator.alter_column_type(
        Webhooks,
        "server_id",
        peewee.IntegerField(null=True),
    )