Explorar o código

#21330 optimize the assignment of tags when saving an object (#21595)

* #21330 optimize object tag creation

* ruff fixes

* optimize

* review changes

* fix

* Update netbox/extras/managers.py

Co-authored-by: Jeremy Stretch <jstretch@netboxlabs.com>

---------

Co-authored-by: Jeremy Stretch <jstretch@netboxlabs.com>
Arthur Hanson hai 12 horas
pai
achega
b5bd8905ca

+ 67 - 0
netbox/extras/managers.py

@@ -0,0 +1,67 @@
+from django.db import router
+from django.db.models import signals
+from taggit.managers import _TaggableManager
+from taggit.utils import require_instance_manager
+
+__all__ = (
+    'NetBoxTaggableManager',
+)
+
+
+class NetBoxTaggableManager(_TaggableManager):
+    """
+    Extends taggit's _TaggableManager to replace the per-tag get_or_create loop in add() with a
+    single bulk_create() call, reducing SQL queries from O(N) to O(1) when assigning tags.
+    """
+
+    @require_instance_manager
+    def add(self, *tags, through_defaults=None, tag_kwargs=None, **kwargs):
+        self._remove_prefetched_objects()
+        if tag_kwargs is None:
+            tag_kwargs = {}
+        db = router.db_for_write(self.through, instance=self.instance)
+
+        tag_objs = self._to_tag_model_instances(tags, tag_kwargs)
+        new_ids = {t.pk for t in tag_objs}
+
+        # Determine which tags are not already assigned to this object
+        lookup = self._lookup_kwargs()
+        vals = set(
+            self.through._default_manager.using(db)
+            .values_list("tag_id", flat=True)
+            .filter(**lookup, tag_id__in=new_ids)
+        )
+        new_ids -= vals
+
+        if not new_ids:
+            return
+
+        signals.m2m_changed.send(
+            sender=self.through,
+            action="pre_add",
+            instance=self.instance,
+            reverse=False,
+            model=self.through.tag_model(),
+            pk_set=new_ids,
+            using=db,
+        )
+
+        # Use a single bulk INSERT instead of one get_or_create per tag.
+        self.through._default_manager.using(db).bulk_create(
+            [
+                self.through(tag=tag, **lookup, **(through_defaults or {}))
+                for tag in tag_objs
+                if tag.pk in new_ids
+            ],
+            ignore_conflicts=True,
+        )
+
+        signals.m2m_changed.send(
+            sender=self.through,
+            action="post_add",
+            instance=self.instance,
+            reverse=False,
+            model=self.through.tag_model(),
+            pk_set=new_ids,
+            using=db,
+        )

+ 3 - 0
netbox/netbox/api/serializers/features.py

@@ -53,8 +53,11 @@ class TaggableModelSerializer(serializers.Serializer):
 
     def _save_tags(self, instance, tags):
         if tags:
+            # Cache tags on instance so serialize_object() can reuse them without a DB query
+            instance._tags = tags
             instance.tags.set([t.name for t in tags])
         else:
+            instance._tags = []
             instance.tags.clear()
 
         return instance

+ 3 - 1
netbox/netbox/models/features.py

@@ -15,6 +15,7 @@ from core.choices import JobStatusChoices, ObjectChangeActionChoices
 from core.models import ObjectType
 from extras.choices import *
 from extras.constants import CUSTOMFIELD_EMPTY_VALUES
+from extras.managers import NetBoxTaggableManager
 from extras.utils import is_taggable
 from netbox.config import get_config
 from netbox.constants import CORE_APPS
@@ -487,11 +488,12 @@ class JournalingMixin(models.Model):
 class TagsMixin(models.Model):
     """
     Enables support for tag assignment. Assigned tags can be managed via the `tags` attribute,
-    which is a `TaggableManager` instance.
+    which is a `NetBoxTaggableManager` instance.
     """
     tags = TaggableManager(
         through='extras.TaggedItem',
         ordering=('weight', 'name'),
+        manager=NetBoxTaggableManager,
     )
 
     class Meta: