From 50088ef393667419997dc12b85c27b75bb28a400 Mon Sep 17 00:00:00 2001
From: jeremystretch <jstretch@ns1.com>
Date: Wed, 19 Oct 2022 15:40:06 -0400
Subject: [PATCH] Extend reindex command to support specifying particular
 models

---
 netbox/extras/management/commands/reindex.py | 78 ++++++++++++++++----
 1 file changed, 65 insertions(+), 13 deletions(-)

diff --git a/netbox/extras/management/commands/reindex.py b/netbox/extras/management/commands/reindex.py
index d8fc69a51..3ae1e21ea 100644
--- a/netbox/extras/management/commands/reindex.py
+++ b/netbox/extras/management/commands/reindex.py
@@ -1,4 +1,5 @@
-from django.core.management.base import BaseCommand
+from django.contrib.contenttypes.models import ContentType
+from django.core.management.base import BaseCommand, CommandError
 
 from extras.models import CachedValue
 from extras.registry import registry
@@ -6,20 +7,71 @@ from netbox.search.backends import search_backend
 
 
 class Command(BaseCommand):
-    """Reindex cached search values"""
-    help = 'Reindex cached search values.'
+    help = 'Reindex objects for search'
 
-    def handle(self, *args, **kwargs):
+    def add_arguments(self, parser):
+        parser.add_argument(
+            'args',
+            metavar='app_label[.ModelName]',
+            nargs='*',
+            help='One or more apps or models to reindex',
+        )
 
-        self.stdout.write('Clearing cached values...', ending="\n")
-        CachedValue.objects.all().delete()
+    def _get_indexers(self, *model_names):
+        indexers = {}
 
-        for app_label, models in registry['search'].items():
-            for name, idx in models.items():
-                self.stdout.write(f'Reindexing {app_label}.{name}...', ending="\n")
-                model = idx.model
-                for instance in model.objects.all():
-                    search_backend.caching_handler(model, instance)
+        # No models specified; pull in all registered indexers
+        if not model_names:
+            for app_label, models in registry['search'].items():
+                for _, idx in models.items():
+                    indexers[idx.model] = idx
+
+        # Return only indexers for the specified models
+        else:
+            for label in model_names:
+                try:
+                    app_label, model_name = label.lower().split('.')
+                except ValueError:
+                    raise CommandError(
+                        f"Invalid model: {label}. Model names must be in the format <app_label>.<model_name>."
+                    )
+                try:
+                    idx = registry['search'][app_label][model_name]
+                    indexers[idx.model] = idx
+                except KeyError:
+                    raise CommandError(f"No indexer found for {label}")
+
+        return indexers
+
+    def handle(self, *model_labels, **kwargs):
+
+        # Determine which models to reindex
+        indexers = self._get_indexers(*model_labels)
+        if not indexers:
+            raise CommandError("No indexers found!")
+        self.stdout.write(f'Reindexing {len(indexers)} models.')
+
+        # Clear all cached values for the specified models
+        self.stdout.write('Clearing cached values... ', ending='')
+        if model_labels:
+            content_types = [
+                ContentType.objects.get_for_model(model) for model in indexers.keys()
+            ]
+            del_count, _ = CachedValue.objects.filter(object_type__in=content_types).delete()
+        else:
+            del_count, _ = CachedValue.objects.all().delete()
+        self.stdout.write(f'{del_count} deleted.')
+
+        # Index models
+        for model, idx in indexers.items():
+            app_label = model._meta.app_label
+            model_name = model._meta.model_name
+            self.stdout.write(f'Reindexing {app_label}.{model_name}... ', ending='')
+            i = 0
+            for i, instance in enumerate(model.objects.all()):
+                search_backend.caching_handler(model, instance)
+            if i:
+                self.stdout.write(f'{i} created.')
 
         cache_size = CachedValue.objects.count()
-        self.stdout.write(f'Done. Generated {cache_size} cached values', ending="\n")
+        self.stdout.write(f'Done. Finished with {cache_size} cached values')