Rewrite trace_paths management command and call in upgrade.sh

2024-05-10 07:54:54 +00:00 · 2020-10-07 10:23:15 -04:00
parent 6db3c65bcc
commit f560693748
3 changed files with 86 additions and 68 deletions
--- a/netbox/dcim/management/commands/trace_paths.py
+++ 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.'))