From f560693748d1f35e79ad9d006a8a9b75ef5ae37b Mon Sep 17 00:00:00 2001 From: Jeremy Stretch Date: Wed, 7 Oct 2020 10:23:15 -0400 Subject: [PATCH] Rewrite trace_paths management command and call in upgrade.sh --- .../dcim/management/commands/retrace_paths.py | 68 ---------------- .../dcim/management/commands/trace_paths.py | 81 +++++++++++++++++++ upgrade.sh | 5 ++ 3 files changed, 86 insertions(+), 68 deletions(-) delete mode 100644 netbox/dcim/management/commands/retrace_paths.py create mode 100644 netbox/dcim/management/commands/trace_paths.py diff --git a/netbox/dcim/management/commands/retrace_paths.py b/netbox/dcim/management/commands/retrace_paths.py deleted file mode 100644 index d11a85417..000000000 --- a/netbox/dcim/management/commands/retrace_paths.py +++ /dev/null @@ -1,68 +0,0 @@ -from django.contrib.contenttypes.models import ContentType -from django.core.management.base import BaseCommand -from django.core.management.color import no_style -from django.db import connection -from django.db.models import Q - -from dcim.models import CablePath -from dcim.signals import create_cablepath - -ENDPOINT_MODELS = ( - 'circuits.CircuitTermination', - 'dcim.ConsolePort', - 'dcim.ConsoleServerPort', - 'dcim.Interface', - 'dcim.PowerFeed', - 'dcim.PowerOutlet', - 'dcim.PowerPort', -) - - -class Command(BaseCommand): - help = "Recalculate natural ordering values for the specified models" - - def add_arguments(self, parser): - parser.add_argument( - 'args', metavar='app_label.ModelName', nargs='*', - help='One or more specific models (each prefixed with its app_label) to retrace', - ) - - def _get_content_types(self, model_names): - q = Q() - for model_name in model_names: - app_label, model = model_name.split('.') - q |= Q(app_label__iexact=app_label, model__iexact=model) - return ContentType.objects.filter(q) - - def handle(self, *model_names, **options): - # Determine the models for which we're retracing all paths - origin_types = self._get_content_types(model_names or ENDPOINT_MODELS) - self.stdout.write(f"Retracing paths for models: {', '.join([str(ct) for ct in origin_types])}") - - # Delete all existing CablePath instances - self.stdout.write(f"Deleting existing cable paths...") - deleted_count, _ = CablePath.objects.filter(origin_type__in=origin_types).delete() - self.stdout.write((self.style.SUCCESS(f' Deleted {deleted_count} paths'))) - - # Reset the SQL sequence. Can do this only if deleting _all_ CablePaths. - if not CablePath.objects.count(): - self.stdout.write(f'Resetting database sequence for CablePath...') - sequence_sql = connection.ops.sequence_reset_sql(no_style(), [CablePath]) - with connection.cursor() as cursor: - for sql in sequence_sql: - cursor.execute(sql) - self.stdout.write(self.style.SUCCESS(' Success.')) - - # Retrace interfaces - for ct in origin_types: - model = ct.model_class() - origins = model.objects.filter(cable__isnull=False) - print(f'Retracing {origins.count()} cabled {model._meta.verbose_name_plural}...') - i = 0 - for i, obj in enumerate(origins, start=1): - create_cablepath(obj) - if not i % 1000: - self.stdout.write(f' {i}') - self.stdout.write(self.style.SUCCESS(f' Retraced {i} {model._meta.verbose_name_plural}')) - - self.stdout.write(self.style.SUCCESS('Finished.')) diff --git a/netbox/dcim/management/commands/trace_paths.py b/netbox/dcim/management/commands/trace_paths.py new file mode 100644 index 000000000..47636a943 --- /dev/null +++ b/netbox/dcim/management/commands/trace_paths.py @@ -0,0 +1,81 @@ +from django.core.management.base import BaseCommand +from django.core.management.color import no_style +from django.db import connection + +from circuits.models import CircuitTermination +from dcim.models import CablePath, ConsolePort, ConsoleServerPort, Interface, PowerFeed, PowerOutlet, PowerPort +from dcim.signals import create_cablepath + +ENDPOINT_MODELS = ( + CircuitTermination, + ConsolePort, + ConsoleServerPort, + Interface, + PowerFeed, + PowerOutlet, + PowerPort +) + + +class Command(BaseCommand): + help = "Generate any missing cable paths among all cable termination objects in NetBox" + + def add_arguments(self, parser): + parser.add_argument( + "--force", action='store_true', dest='force', + help="Force recalculation of all existing cable paths" + ) + parser.add_argument( + "--no-input", action='store_true', dest='no_input', + help="Do not prompt user for any input/confirmation" + ) + + def handle(self, *model_names, **options): + + # If --force was passed, first delete all existing CablePaths + if options['force']: + cable_paths = CablePath.objects.all() + paths_count = cable_paths.count() + + # Prompt the user to confirm recalculation of all paths + if paths_count and not options['no_input']: + self.stdout.write(self.style.ERROR("WARNING: Forcing recalculation of all cable paths.")) + self.stdout.write( + f"This will delete and recalculate all {paths_count} existing cable paths. Are you sure?" + ) + confirmation = input("Type yes to confirm: ") + if confirmation != 'yes': + self.stdout.write(self.style.SUCCESS("Aborting")) + return + + # Delete all existing CablePath instances + self.stdout.write(f"Deleting {paths_count} existing cable paths...") + deleted_count, _ = CablePath.objects.all().delete() + self.stdout.write((self.style.SUCCESS(f' Deleted {deleted_count} paths'))) + + # Reinitialize the model's PK sequence + self.stdout.write(f'Resetting database sequence for CablePath model') + sequence_sql = connection.ops.sequence_reset_sql(no_style(), [CablePath]) + with connection.cursor() as cursor: + for sql in sequence_sql: + cursor.execute(sql) + + # Retrace paths + for model in ENDPOINT_MODELS: + origins = model.objects.filter(cable__isnull=False) + if not options['force']: + origins = origins.filter(_path__isnull=True) + origins_count = origins.count() + if not origins_count: + print(f'Found no missing {model._meta.verbose_name} paths; skipping') + continue + print(f'Retracing {origins_count} cabled {model._meta.verbose_name_plural}...') + i = 0 + for i, obj in enumerate(origins, start=1): + create_cablepath(obj) + # TODO: Come up with a better progress indicator + if not i % 1000: + self.stdout.write(f' {i}') + self.stdout.write(self.style.SUCCESS(f' Retraced {i} {model._meta.verbose_name_plural}')) + + self.stdout.write(self.style.SUCCESS('Finished.')) diff --git a/upgrade.sh b/upgrade.sh index 66ba7b39f..468f189b3 100755 --- a/upgrade.sh +++ b/upgrade.sh @@ -55,6 +55,11 @@ COMMAND="python3 netbox/manage.py migrate" echo "Applying database migrations ($COMMAND)..." eval $COMMAND || exit 1 +# Trace any missing cable paths (not typically needed) +COMMAND="python3 netbox/manage.py trace_paths --no-input" +echo "Checking for missing cable paths ($COMMAND)..." +eval $COMMAND || exit 1 + # Collect static files COMMAND="python3 netbox/manage.py collectstatic --no-input" echo "Collecting static files ($COMMAND)..."