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

77 lines
2.3 KiB
Python

# 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 = []
hosts_updated = []
for host in hosts:
hardware = hardware_model()
asset = asset_mapper[host.asset_ptr_id]
hardware.date_updated = timezone.now()
for name in fields:
setattr(hardware, name, getattr(asset, name))
hardware_infos.append(hardware)
host.device_info_id = hardware.id
hosts_updated.append(host)
hardware_model.objects.bulk_create(hardware_infos, ignore_conflicts=True)
host_model.objects.bulk_update(hosts_updated, ['device_info_id'])
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),
]