1
0
mirror of https://github.com/netbox-community/netbox.git synced 2024-05-10 07:54:54 +00:00
netbox-community-netbox/netbox/dcim/migrations/0158_populate_cable_terminations.py
2022-07-15 10:07:35 -04:00

88 lines
2.9 KiB
Python

import sys
from django.db import migrations
def cache_related_objects(termination):
"""
Replicate caching logic from CableTermination.cache_related_objects()
"""
attrs = {}
# Device components
if getattr(termination, 'device', None):
attrs['_device'] = termination.device
attrs['_rack'] = termination.device.rack
attrs['_location'] = termination.device.location
attrs['_site'] = termination.device.site
# Power feeds
elif getattr(termination, 'rack', None):
attrs['_rack'] = termination.rack
attrs['_location'] = termination.rack.location
attrs['_site'] = termination.rack.site
# Circuit terminations
elif getattr(termination, 'site', None):
attrs['_site'] = termination.site
return attrs
def populate_cable_terminations(apps, schema_editor):
"""
Replicate terminations from the Cable model into CableTermination instances.
"""
ContentType = apps.get_model('contenttypes', 'ContentType')
Cable = apps.get_model('dcim', 'Cable')
CableTermination = apps.get_model('dcim', 'CableTermination')
# Retrieve the necessary data from Cable objects
cables = Cable.objects.values(
'id', 'termination_a_type', 'termination_a_id', 'termination_b_type', 'termination_b_id'
)
# Queue CableTerminations to be created
cable_terminations = []
cable_count = cables.count()
for i, cable in enumerate(cables, start=1):
for cable_end in ('a', 'b'):
# We must manually instantiate the termination object, because GFK fields are not
# supported within migrations.
termination_ct = ContentType.objects.get(pk=cable[f'termination_{cable_end}_type'])
termination_model = apps.get_model(termination_ct.app_label, termination_ct.model)
termination = termination_model.objects.get(pk=cable[f'termination_{cable_end}_id'])
cable_terminations.append(CableTermination(
cable_id=cable['id'],
cable_end=cable_end.upper(),
termination_type_id=cable[f'termination_{cable_end}_type'],
termination_id=cable[f'termination_{cable_end}_id'],
**cache_related_objects(termination)
))
# Output progress occasionally
if 'test' not in sys.argv and not i % 100:
progress = float(i) * 100 / cable_count
if i == 100:
print('')
sys.stdout.write(f"\r Updated {i}/{cable_count} cables ({progress:.2f}%)")
sys.stdout.flush()
# Bulk create the termination objects
CableTermination.objects.bulk_create(cable_terminations, batch_size=100)
class Migration(migrations.Migration):
dependencies = [
('dcim', '0157_new_cabling_models'),
]
operations = [
migrations.RunPython(
code=populate_cable_terminations,
reverse_code=migrations.RunPython.noop
),
]