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_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

    count = 0
    batch_size = 1000

    while True:
        assets = asset_model.objects.using(db_alias).all()[count:count+batch_size]
        if not assets:
            break
        count += len(assets)
        hosts = [host_model(asset_ptr=asset) for asset in assets]
        host_model.objects.using(db_alias).bulk_create(hosts, ignore_conflicts=True)


def migrate_hardware_info(apps, *args):
    asset_model = apps.get_model("assets", "Asset")

    count = 0
    batch_size = 1000
    hardware_fields = [
        'vendor', 'model', 'sn', 'cpu_model', 'cpu_count', 'cpu_cores',
        'cpu_vcpus', 'memory', 'disk_total', 'disk_info', 'os', 'os_arch',
        'os_version', 'hostname_raw', 'number'
    ]

    while True:
        assets = asset_model.objects.all()[count:count+batch_size]
        if not assets:
            break
        count += len(assets)

        updated = []
        for asset in assets:
            info = {field: getattr(asset, field) for field in hardware_fields if getattr(asset, field)}
            if not info:
                continue
            asset.info = info
            updated.append(asset)
        asset_model.objects.bulk_update(updated, ['info'])


class Migration(migrations.Migration):

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

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