Rewrite trace_paths management command and call in upgrade.sh

netbox/dcim/management/commands/retrace_paths.py

@@ -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.'))

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.'))

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)..."