diff --git a/netbox/ipam/migrations/0083_ipaddress_iprange_prefix_parent_data.py b/netbox/ipam/migrations/0083_ipaddress_iprange_prefix_parent_data.py index 321c51f38..45cf48c04 100644 --- a/netbox/ipam/migrations/0083_ipaddress_iprange_prefix_parent_data.py +++ b/netbox/ipam/migrations/0083_ipaddress_iprange_prefix_parent_data.py @@ -19,25 +19,28 @@ def draw_progress(count, total, length=20): sys.stdout.flush() -def set_ipaddress_prefix(apps, schema_editor): +def set_prefix(apps, schema_editor, model, attr='address', parent_model='Prefix'): start = time.time() - IPAddress = apps.get_model('ipam', 'IPAddress') - Prefix = apps.get_model('ipam', 'Prefix') + ChildModel = apps.get_model('ipam', model) + ParentModel = apps.get_model('ipam', parent_model) - addresses = IPAddress.objects.all() - i = 0 + addresses = ChildModel.objects.all() total = addresses.count() if total == 0: return + print('\r\n') + i = 0 draw_progress(i, total, 50) for address in addresses: i += 1 - prefixes = Prefix.objects.filter( - vrf=address.vrf, - prefix__net_contains_or_equals=str(address.address.ip), - prefix__net_mask_length__lte=address.address.prefixlen, + address_attr = getattr(address, attr) + prefixes = ParentModel.objects.filter( + prefix__net_contains_or_equals=str(address_attr.ip), + prefix__net_mask_length__lte=address_attr.prefixlen, ) + if hasattr(ParentModel, 'vrf'): + prefixes = prefixes.filter(vrf=address.vrf) address.prefix = prefixes.last() address.save() draw_progress(i, total, 50) @@ -46,37 +49,17 @@ def set_ipaddress_prefix(apps, schema_editor): print(f"\r\nElapsed Time: {end - start:.2f}s") +def set_ipaddress_prefix(apps, schema_editor): + set_prefix(apps, schema_editor, 'IPAddress') + + def unset_ipaddress_prefix(apps, schema_editor): IPAddress = apps.get_model('ipam', 'IPAddress') IPAddress.objects.update(prefix=None) def set_iprange_prefix(apps, schema_editor): - IPRange = apps.get_model('ipam', 'IPRange') - Prefix = apps.get_model('ipam', 'Prefix') - start = time.time() - - addresses = IPRange.objects.all() - i = 0 - total = addresses.count() - if total == 0: - return - - print('\r\n') - draw_progress(i, total, 50) - for address in addresses: - i += 1 - prefixes = Prefix.objects.filter( - vrf=address.vrf, - prefix__net_contains_or_equals=str(address.start_address.ip), - prefix__net_mask_length__lte=address.start_address.prefixlen, - ) - address.prefix = prefixes.last() - address.save() - draw_progress(i, total, 50) - - end = time.time() - print(f"\r\nElapsed Time: {end - start:.2f}s") + set_prefix(apps, schema_editor, 'IPRange', 'start_address') def unset_iprange_prefix(apps, schema_editor): @@ -85,30 +68,7 @@ def unset_iprange_prefix(apps, schema_editor): def set_prefix_aggregate(apps, schema_editor): - Prefix = apps.get_model('ipam', 'Prefix') - Aggregate = apps.get_model('ipam', 'Aggregate') - - start = time.time() - addresses = Prefix.objects.all() - i = 0 - total = addresses.count() - if total == 0: - return - - print('\r\n') - draw_progress(i, total, 50) - for address in addresses: - i += 1 - aggregates = Aggregate.objects.filter( - prefix__net_contains_or_equals=str(address.prefix.ip), - prefix__net_mask_length__lte=address.prefix.prefixlen, - ) - address.aggregate = aggregates.last() - address.save() - draw_progress(i, total, 50) - - end = time.time() - print(f"\r\nElapsed Time: {end - start:.2f}s") + set_prefix(apps, schema_editor, 'Prefix', 'prefix', 'Aggregate') def unset_prefix_aggregate(apps, schema_editor):