Просмотр исходного кода

Closes #18980: Optimize update of object data when adding/removing custom fields (#18983)

* Employ native PostgreSQL functions for updating object JSON data when adding/removing custom fields

* Optimize rename_object_data()

* remove_stale_data() should validate model class
Jeremy Stretch 10 месяцев назад
Родитель
Сommit
af5a600583
1 измененных файлов с 31 добавлено и 16 удалено
  1. 31 16
      netbox/extras/models/customfields.py

+ 31 - 16
netbox/extras/models/customfields.py

@@ -9,6 +9,8 @@ from django.conf import settings
 from django.contrib.postgres.fields import ArrayField
 from django.contrib.postgres.fields import ArrayField
 from django.core.validators import RegexValidator, ValidationError
 from django.core.validators import RegexValidator, ValidationError
 from django.db import models
 from django.db import models
+from django.db.models import F, Func, Value
+from django.db.models.expressions import RawSQL
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.html import escape
 from django.utils.html import escape
 from django.utils.safestring import mark_safe
 from django.utils.safestring import mark_safe
@@ -281,12 +283,20 @@ class CustomField(CloningMixin, ExportTemplatesMixin, ChangeLoggedModel):
         Populate initial custom field data upon either a) the creation of a new CustomField, or
         Populate initial custom field data upon either a) the creation of a new CustomField, or
         b) the assignment of an existing CustomField to new object types.
         b) the assignment of an existing CustomField to new object types.
         """
         """
+        if self.default is None:
+            # We have to convert None to a JSON null for jsonb_set()
+            value = RawSQL("'null'::jsonb", [])
+        else:
+            value = Value(self.default, models.JSONField())
         for ct in content_types:
         for ct in content_types:
-            model = ct.model_class()
-            instances = model.objects.exclude(**{'custom_field_data__contains': self.name})
-            for instance in instances:
-                instance.custom_field_data[self.name] = self.default
-            model.objects.bulk_update(instances, ['custom_field_data'], batch_size=100)
+            ct.model_class().objects.update(
+                custom_field_data=Func(
+                    F('custom_field_data'),
+                    Value([self.name]),
+                    value,
+                    function='jsonb_set'
+                )
+            )
 
 
     def remove_stale_data(self, content_types):
     def remove_stale_data(self, content_types):
         """
         """
@@ -295,22 +305,27 @@ class CustomField(CloningMixin, ExportTemplatesMixin, ChangeLoggedModel):
         """
         """
         for ct in content_types:
         for ct in content_types:
             if model := ct.model_class():
             if model := ct.model_class():
-                instances = model.objects.filter(custom_field_data__has_key=self.name)
-                for instance in instances:
-                    del instance.custom_field_data[self.name]
-                model.objects.bulk_update(instances, ['custom_field_data'], batch_size=100)
+                model.objects.update(
+                    custom_field_data=F('custom_field_data') - self.name
+                )
 
 
     def rename_object_data(self, old_name, new_name):
     def rename_object_data(self, old_name, new_name):
         """
         """
-        Called when a CustomField has been renamed. Updates all assigned object data.
+        Called when a CustomField has been renamed. Removes the original key and inserts the new
+        one, copying the value of the old key.
         """
         """
         for ct in self.object_types.all():
         for ct in self.object_types.all():
-            model = ct.model_class()
-            params = {f'custom_field_data__{old_name}__isnull': False}
-            instances = model.objects.filter(**params)
-            for instance in instances:
-                instance.custom_field_data[new_name] = instance.custom_field_data.pop(old_name)
-            model.objects.bulk_update(instances, ['custom_field_data'], batch_size=100)
+            ct.model_class().objects.update(
+                custom_field_data=Func(
+                    F('custom_field_data') - old_name,
+                    Value([new_name]),
+                    Func(
+                        F('custom_field_data'),
+                        function='jsonb_extract_path_text',
+                        template=f"to_jsonb(%(expressions)s -> '{old_name}')"
+                    ),
+                    function='jsonb_set')
+            )
 
 
     def clean(self):
     def clean(self):
         super().clean()
         super().clean()