jumpserver/apps/assets/migrations/0093_auto_20220403_1627.py

# Generated by Django 3.1.14 on 2022-04-02 08:27

from django.utils import timezone
from django.db import migrations, models


def migrate_hardware(apps, *args):
    host_model = apps.get_model('assets', 'Host')
    asset_model = apps.get_model('assets', 'Asset')
    hardware_model = apps.get_model('assets', 'DeviceInfo')

    created = 0
    batch_size = 1000

    excludes = ['id', 'host', 'date_updated']
    fields = [f.name for f in hardware_model._meta.fields]
    fields = [name for name in fields if name not in excludes]

    while True:
        start = created
        end = created + batch_size
        hosts = host_model.objects.all()[start:end]
        asset_ids = [h.asset_ptr_id for h in hosts]
        assets = asset_model.objects.filter(id__in=asset_ids)
        asset_mapper = {a.id: a for a in assets}

        if not hosts:
            break

        hardware_infos = []
        for host in hosts:
            hardware = hardware_model()
            asset = asset_mapper[host.asset_ptr_id]
            hardware.host = host
            hardware.date_updated = timezone.now()
            for name in fields:
                setattr(hardware, name, getattr(asset, name))
            hardware_infos.append(hardware)

        hardware_model.objects.bulk_create(hardware_infos, ignore_conflicts=True)
        created += len(hardware_infos)


def migrate_to_host(apps, schema_editor):
    asset_model = apps.get_model("assets", "Asset")
    host_model = apps.get_model("assets", 'Host')
    db_alias = schema_editor.connection.alias

    created = 0
    batch_size = 1000

    while True:
        start = created
        end = created + batch_size
        assets = asset_model.objects.using(db_alias).all()[start:end]
        if not assets:
            break

        hosts = [host_model(asset_ptr=asset) for asset in assets]
        host_model.objects.using(db_alias).bulk_create(hosts, ignore_conflicts=True)
        created += len(hosts)


class Migration(migrations.Migration):

    dependencies = [
        ('assets', '0092_add_host'),
    ]

    operations = [
        migrations.RunPython(migrate_to_host),
        migrations.RunPython(migrate_hardware),
    ]