Closes #22595: Introduce BULK_UPDATE_CHUNK_SIZE config parameter to limit max number of rows per bulk update (#22728)
This commit is contained in:
Jeremy Stretch
GitHub
parent
a24fbb06ce
commit
b62c384daf
|
|
@ -12,6 +12,20 @@ BASE_PATH = 'netbox/'
|
|||
|
||||
---
|
||||
|
||||
## BULK_UPDATE_CHUNK_SIZE
|
||||
|
||||
Default: `5000`
|
||||
|
||||
The maximum number of rows to affect in a single SQL `UPDATE` statement when NetBox performs a bulk update across many objects (for example, when recalculating cached counters or backfilling custom field data). On very large tables, an unbounded update spanning millions of rows can exceed the database's configured statement timeout; splitting the work into batches of at most this many rows bounds each statement while keeping the overall operation atomic.
|
||||
|
||||
Must be a positive integer, or `None` to disable chunking and issue each bulk update as a single unbounded statement.
|
||||
|
||||
```python
|
||||
BULK_UPDATE_CHUNK_SIZE = 5000
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## DATABASE_ROUTERS
|
||||
|
||||
Default: `[]` (empty list)
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
from django.core.cache import cache
|
||||
from django.db import models
|
||||
from django.db import models, router, transaction
|
||||
from django.urls import reverse
|
||||
from django.utils.translation import gettext
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
|
||||
from utilities.querysets import RestrictedQuerySet
|
||||
from utilities.querysets import RestrictedQuerySet, chunked_update
|
||||
|
||||
__all__ = (
|
||||
'ConfigRevision',
|
||||
|
|
@ -78,9 +78,15 @@ class ConfigRevision(models.Model):
|
|||
cache.set('config_version', self.pk, None)
|
||||
|
||||
if update_db:
|
||||
# Set all instances of ConfigRevision to false and set this instance to true
|
||||
ConfigRevision.objects.all().update(active=False)
|
||||
ConfigRevision.objects.filter(pk=self.pk).update(active=True)
|
||||
# Set all instances of ConfigRevision to false and set this instance to true. Wrap both
|
||||
# statements in a transaction so the "exactly one active revision" invariant is preserved
|
||||
# even when the deactivation is chunked into multiple statements. Resolve the write alias
|
||||
# once and pin the transaction and both querysets to it, so the transaction genuinely
|
||||
# covers the (potentially router-directed) writes performed by chunked_update().
|
||||
using = router.db_for_write(ConfigRevision)
|
||||
with transaction.atomic(using=using):
|
||||
chunked_update(ConfigRevision.objects.using(using).all(), active=False)
|
||||
ConfigRevision.objects.using(using).filter(pk=self.pk).update(active=True)
|
||||
|
||||
activate.alters_data = True
|
||||
|
||||
|
|
|
|||
|
|
@ -220,7 +220,9 @@ class DataSource(JobsMixin, PrimaryModel):
|
|||
continue
|
||||
|
||||
# Bulk update modified files
|
||||
updated_count = DataFile.objects.bulk_update(updated_files, ('last_updated', 'size', 'hash', 'data'))
|
||||
updated_count = DataFile.objects.bulk_update(
|
||||
updated_files, ('last_updated', 'size', 'hash', 'data'), batch_size=settings.BULK_UPDATE_CHUNK_SIZE
|
||||
)
|
||||
logger.debug(f"Updated {updated_count} files")
|
||||
|
||||
# Bulk delete deleted files
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from netbox.models import ChangeLoggedModel, PrimaryModel
|
|||
from utilities.conversion import to_meters
|
||||
from utilities.exceptions import AbortRequest
|
||||
from utilities.fields import ColorField, GenericArrayForeignKey
|
||||
from utilities.querysets import RestrictedQuerySet
|
||||
from utilities.querysets import RestrictedQuerySet, chunked_update
|
||||
from utilities.serialization import deserialize_object, serialize_object
|
||||
from wireless.models import WirelessLink
|
||||
|
||||
|
|
@ -791,7 +791,7 @@ class CablePath(models.Model):
|
|||
# Record a direct reference to this CablePath on its originating object(s)
|
||||
origin_model = self.origin_type.model_class()
|
||||
origin_ids = [decompile_path_node(node)[1] for node in self.path[0]]
|
||||
origin_model.objects.filter(pk__in=origin_ids).update(_path=self.pk)
|
||||
chunked_update(origin_model.objects.filter(pk__in=origin_ids), _path=self.pk)
|
||||
|
||||
def delete(self, *args, **kwargs):
|
||||
# Mirror save() - clear _path on origins to prevent stale references
|
||||
|
|
@ -799,7 +799,7 @@ class CablePath(models.Model):
|
|||
if self.path:
|
||||
origin_model = self.origin_type.model_class()
|
||||
origin_ids = [decompile_path_node(node)[1] for node in self.path[0]]
|
||||
origin_model.objects.filter(pk__in=origin_ids, _path=self.pk).update(_path=None)
|
||||
chunked_update(origin_model.objects.filter(pk__in=origin_ids, _path=self.pk), _path=None)
|
||||
|
||||
super().delete(*args, **kwargs)
|
||||
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from dcim.utils import (
|
|||
)
|
||||
from utilities.counters import update_counter
|
||||
from utilities.exceptions import AbortRequest
|
||||
from utilities.querysets import chunked_update
|
||||
|
||||
from .device_components import (
|
||||
ConsolePort,
|
||||
|
|
@ -825,10 +826,13 @@ class ModuleMovePlan:
|
|||
moved_front_port_pks = [obj.pk for obj in self.components[FrontPort]]
|
||||
moved_rear_port_pks = [obj.pk for obj in self.components[RearPort]]
|
||||
if moved_front_port_pks and moved_rear_port_pks:
|
||||
PortMapping.objects.filter(
|
||||
front_port_id__in=moved_front_port_pks,
|
||||
rear_port_id__in=moved_rear_port_pks,
|
||||
).update(device_id=self.new_device_id)
|
||||
chunked_update(
|
||||
PortMapping.objects.filter(
|
||||
front_port_id__in=moved_front_port_pks,
|
||||
rear_port_id__in=moved_rear_port_pks,
|
||||
),
|
||||
device_id=self.new_device_id,
|
||||
)
|
||||
|
||||
def _recompute_counters(self):
|
||||
# bulk updates bypass the signal-driven counters; apply exact deltas for both devices
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from collections.abc import Iterable, Mapping
|
|||
|
||||
import jsonschema
|
||||
import yaml
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import ValidationError
|
||||
from django.db import OperationalError, models, router, transaction
|
||||
from django.db.models.signals import post_save
|
||||
|
|
@ -577,7 +578,9 @@ class Module(TrackingModelMixin, PrimaryModel):
|
|||
instance.parent = self.module_bay
|
||||
update_fields = ['module', 'parent']
|
||||
|
||||
component_model.objects.bulk_update(update_instances, update_fields)
|
||||
component_model.objects.bulk_update(
|
||||
update_instances, update_fields, batch_size=settings.BULK_UPDATE_CHUNK_SIZE
|
||||
)
|
||||
for component in update_instances:
|
||||
post_save.send(
|
||||
sender=component_model,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from django.dispatch import receiver
|
|||
|
||||
from dcim.choices import CableEndChoices, LinkStatusChoices
|
||||
from netbox.search.backends import search_backend
|
||||
from utilities.querysets import chunked_update
|
||||
from virtualization.models import VMInterface
|
||||
|
||||
from .models import (
|
||||
|
|
@ -37,11 +38,11 @@ def handle_location_site_change(instance, created, **kwargs):
|
|||
(and to descendant Locations).
|
||||
"""
|
||||
if not created:
|
||||
instance.get_descendants().update(site=instance.site)
|
||||
chunked_update(instance.get_descendants(), site=instance.site)
|
||||
locations = instance.get_descendants(include_self=True).values_list('pk', flat=True)
|
||||
Rack.objects.filter(location__in=locations).update(site=instance.site)
|
||||
Device.objects.filter(location__in=locations).update(site=instance.site)
|
||||
PowerPanel.objects.filter(location__in=locations).update(site=instance.site)
|
||||
chunked_update(Rack.objects.filter(location__in=locations), site=instance.site)
|
||||
chunked_update(Device.objects.filter(location__in=locations), site=instance.site)
|
||||
chunked_update(PowerPanel.objects.filter(location__in=locations), site=instance.site)
|
||||
|
||||
|
||||
@receiver(post_save, sender=Rack)
|
||||
|
|
@ -50,7 +51,7 @@ def handle_rack_site_change(instance, created, **kwargs):
|
|||
Cascade a Rack's Site/Location assignment down to the Devices it contains.
|
||||
"""
|
||||
if not created:
|
||||
Device.objects.filter(rack=instance).update(site=instance.site, location=instance.location)
|
||||
chunked_update(Device.objects.filter(rack=instance), site=instance.site, location=instance.location)
|
||||
|
||||
|
||||
#
|
||||
|
|
@ -123,7 +124,7 @@ def update_connected_endpoints(instance, created, raw=False, **kwargs):
|
|||
# Update status of CablePaths if Cable status has been changed
|
||||
elif instance.status != instance._orig_status:
|
||||
if instance.status != LinkStatusChoices.STATUS_CONNECTED:
|
||||
CablePath.objects.filter(_nodes__contains=instance).update(is_active=False)
|
||||
chunked_update(CablePath.objects.filter(_nodes__contains=instance), is_active=False)
|
||||
else:
|
||||
rebuild_paths([instance])
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from django.db.models import F, Q
|
|||
from dcim.models import Device
|
||||
from extras.jobs import RenderConfigContextJob
|
||||
from extras.models.tags import TaggedItem
|
||||
from utilities.querysets import chunked_update
|
||||
from virtualization.models import VirtualMachine
|
||||
|
||||
|
||||
|
|
@ -30,7 +31,8 @@ def invalidate_config_context_for_objects(model_label, pks):
|
|||
return
|
||||
|
||||
Model = apps.get_model(model_label)
|
||||
updated = Model.objects.filter(pk__in=pks).update(
|
||||
updated = chunked_update(
|
||||
Model.objects.filter(pk__in=pks),
|
||||
_config_context_data=None,
|
||||
_config_context_generation=F('_config_context_generation') + 1,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,14 +6,6 @@ from extras.choices import LogLevelChoices
|
|||
# Custom fields
|
||||
CUSTOMFIELD_EMPTY_VALUES = (None, '', [])
|
||||
|
||||
# Maximum number of objects to update per query when provisioning, removing, or renaming custom
|
||||
# field data. Bounding the number of rows touched by each statement prevents very large tables from
|
||||
# exceeding the database statement timeout (JSONB updates rewrite each affected row). This value
|
||||
# sits at the throughput "knee": benchmarking jsonb_set() across a 1M-row table showed throughput
|
||||
# plateaus by ~5K rows/statement (raising it further yields no meaningful speedup), while keeping
|
||||
# each statement orders of magnitude below a typical statement timeout.
|
||||
CUSTOMFIELD_DATA_BATCH_SIZE = 5000
|
||||
|
||||
# ImageAttachment
|
||||
IMAGE_ATTACHMENT_IMAGE_FORMATS = {
|
||||
'avif': 'image/avif',
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from django.apps import apps
|
|||
from django.core.management.base import BaseCommand, CommandError
|
||||
|
||||
from utilities.fields import NaturalOrderingField
|
||||
from utilities.querysets import chunked_update
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
|
|
@ -93,7 +94,7 @@ class Command(BaseCommand):
|
|||
self.stdout.flush()
|
||||
|
||||
# Update each unique field value in bulk
|
||||
changed = model.objects.filter(name=value).update(**{field.name: naturalized_value})
|
||||
changed = chunked_update(model.objects.filter(name=value), **{field.name: naturalized_value})
|
||||
|
||||
if options['verbosity'] >= 2:
|
||||
self.stdout.write(f" ({changed})")
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import jsonschema
|
|||
from django import forms
|
||||
from django.conf import settings
|
||||
from django.core.validators import RegexValidator, ValidationError
|
||||
from django.db import models, transaction
|
||||
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
|
||||
|
|
@ -19,7 +19,6 @@ from jsonschema.exceptions import ValidationError as JSONValidationError
|
|||
|
||||
from core.models import ObjectType
|
||||
from extras.choices import *
|
||||
from extras.constants import CUSTOMFIELD_DATA_BATCH_SIZE
|
||||
from extras.data import CHOICE_SETS
|
||||
from extras.fields import ChoiceSetField
|
||||
from netbox.context import query_cache
|
||||
|
|
@ -44,7 +43,7 @@ from utilities.forms.fields import (
|
|||
from utilities.forms.utils import add_blank_choice
|
||||
from utilities.forms.widgets import APISelect, APISelectMultiple, DatePicker, DateTimePicker
|
||||
from utilities.jsonschema import validate_schema
|
||||
from utilities.querysets import RestrictedQuerySet
|
||||
from utilities.querysets import RestrictedQuerySet, chunked_update
|
||||
from utilities.templatetags.builtins.filters import render_markdown
|
||||
from utilities.validators import validate_regex
|
||||
|
||||
|
|
@ -329,32 +328,6 @@ class CustomField(CloningMixin, ExportTemplatesMixin, OwnerMixin, ChangeLoggedMo
|
|||
return self.choice_set.get_choice_color(value)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _update_object_data(model, **update_kwargs):
|
||||
"""
|
||||
Apply an UPDATE to the custom_field_data of every instance of the given model in batches,
|
||||
bounding the number of rows touched by each statement. A single unbounded UPDATE across
|
||||
millions of rows can exceed the database statement timeout, because JSONB updates rewrite
|
||||
each affected row in full. Batches are selected via keyset pagination on the primary key.
|
||||
|
||||
The batched updates are wrapped in a transaction so that the operation remains atomic, as
|
||||
it was when performed by a single UPDATE. This guards against partially-applied data (e.g.
|
||||
a renamed field landing on only some objects) should the loop be interrupted when not
|
||||
already running inside a request's transaction. Batching avoids the statement timeout
|
||||
regardless, as that limit applies per statement rather than per transaction.
|
||||
"""
|
||||
with transaction.atomic():
|
||||
last_pk = 0
|
||||
while True:
|
||||
pks = list(
|
||||
model.objects.filter(pk__gt=last_pk).order_by('pk')
|
||||
.values_list('pk', flat=True)[:CUSTOMFIELD_DATA_BATCH_SIZE]
|
||||
)
|
||||
if not pks:
|
||||
break
|
||||
model.objects.filter(pk__in=pks).update(**update_kwargs)
|
||||
last_pk = pks[-1]
|
||||
|
||||
def populate_initial_data(self, content_types):
|
||||
"""
|
||||
Populate initial custom field data upon either a) the creation of a new CustomField, or
|
||||
|
|
@ -367,8 +340,8 @@ class CustomField(CloningMixin, ExportTemplatesMixin, OwnerMixin, ChangeLoggedMo
|
|||
value = Value(self.default, models.JSONField())
|
||||
for ct in content_types:
|
||||
if model := ct.model_class():
|
||||
self._update_object_data(
|
||||
model,
|
||||
chunked_update(
|
||||
model.objects.all(),
|
||||
custom_field_data=Func(
|
||||
F('custom_field_data'),
|
||||
Value([self.name]),
|
||||
|
|
@ -384,8 +357,8 @@ class CustomField(CloningMixin, ExportTemplatesMixin, OwnerMixin, ChangeLoggedMo
|
|||
"""
|
||||
for ct in content_types:
|
||||
if model := ct.model_class():
|
||||
self._update_object_data(
|
||||
model,
|
||||
chunked_update(
|
||||
model.objects.all(),
|
||||
custom_field_data=F('custom_field_data') - self.name
|
||||
)
|
||||
|
||||
|
|
@ -396,8 +369,8 @@ class CustomField(CloningMixin, ExportTemplatesMixin, OwnerMixin, ChangeLoggedMo
|
|||
"""
|
||||
for ct in self.object_types.all():
|
||||
if model := ct.model_class():
|
||||
self._update_object_data(
|
||||
model,
|
||||
chunked_update(
|
||||
model.objects.all(),
|
||||
custom_field_data=Func(
|
||||
F('custom_field_data') - old_name,
|
||||
Value([new_name]),
|
||||
|
|
|
|||
|
|
@ -1,10 +1,9 @@
|
|||
import datetime
|
||||
import json
|
||||
from decimal import Decimal
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.core.exceptions import ValidationError
|
||||
from django.test import tag
|
||||
from django.test import override_settings, tag
|
||||
from django.urls import reverse
|
||||
from rest_framework import status
|
||||
|
||||
|
|
@ -674,7 +673,7 @@ class CustomFieldTestCase(TestCase):
|
|||
self.assertNotIn('field1', site.custom_field_data)
|
||||
self.assertEqual(site.custom_field_data['field2'], FIELD_DATA)
|
||||
|
||||
@patch('extras.models.customfields.CUSTOMFIELD_DATA_BATCH_SIZE', 2)
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=2)
|
||||
def test_batched_object_data_updates(self):
|
||||
"""
|
||||
Provisioning, renaming, and removing custom field data is applied in batches. Use a small
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from django.core.management.base import BaseCommand
|
|||
|
||||
from ipam.models import VRF, Prefix
|
||||
from ipam.utils import rebuild_prefixes
|
||||
from utilities.querysets import chunked_update
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
|
|
@ -11,7 +12,7 @@ class Command(BaseCommand):
|
|||
self.stdout.write(f'Rebuilding {Prefix.objects.count()} prefixes...')
|
||||
|
||||
# Reset existing counts
|
||||
Prefix.objects.update(_depth=0, _children=0)
|
||||
chunked_update(Prefix.objects.all(), _depth=0, _children=0)
|
||||
|
||||
# Rebuild the global table
|
||||
global_count = Prefix.objects.filter(vrf__isnull=True).count()
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ class RebuildPrefixesTestCase(TestCase):
|
|||
patch('ipam.management.commands.rebuild_prefixes.Prefix') as prefix_model,
|
||||
patch('ipam.management.commands.rebuild_prefixes.VRF') as vrf_model,
|
||||
patch('ipam.management.commands.rebuild_prefixes.rebuild_prefixes') as rebuild_prefixes,
|
||||
patch('ipam.management.commands.rebuild_prefixes.chunked_update') as chunked_update,
|
||||
):
|
||||
prefix_model.objects.count.return_value = 0
|
||||
prefix_model.objects.filter.return_value.count.return_value = 0
|
||||
|
|
@ -23,7 +24,7 @@ class RebuildPrefixesTestCase(TestCase):
|
|||
call_command('rebuild_prefixes', stdout=out)
|
||||
|
||||
rebuild_prefixes.assert_called_once_with(None)
|
||||
prefix_model.objects.update.assert_called_once_with(_depth=0, _children=0)
|
||||
chunked_update.assert_called_once_with(prefix_model.objects.all.return_value, _depth=0, _children=0)
|
||||
self.assertIn('Rebuilding 0 prefixes', out.getvalue())
|
||||
self.assertIn('Finished.', out.getvalue())
|
||||
|
||||
|
|
@ -68,6 +69,7 @@ class RebuildPrefixesTestCase(TestCase):
|
|||
patch('ipam.management.commands.rebuild_prefixes.Prefix') as prefix_model,
|
||||
patch('ipam.management.commands.rebuild_prefixes.VRF') as vrf_model,
|
||||
patch('ipam.management.commands.rebuild_prefixes.rebuild_prefixes') as rebuild_prefixes,
|
||||
patch('ipam.management.commands.rebuild_prefixes.chunked_update'),
|
||||
):
|
||||
prefix_model.objects.count.return_value = 3
|
||||
prefix_model.objects.filter.side_effect = [
|
||||
|
|
|
|||
|
|
@ -89,6 +89,11 @@ AUTH_PASSWORD_VALIDATORS = getattr(configuration, 'AUTH_PASSWORD_VALIDATORS', [
|
|||
},
|
||||
])
|
||||
BASE_PATH = trailing_slash(getattr(configuration, 'BASE_PATH', ''))
|
||||
BULK_UPDATE_CHUNK_SIZE = getattr(configuration, 'BULK_UPDATE_CHUNK_SIZE', 5000)
|
||||
if BULK_UPDATE_CHUNK_SIZE is not None and (type(BULK_UPDATE_CHUNK_SIZE) is not int or BULK_UPDATE_CHUNK_SIZE < 1):
|
||||
raise ImproperlyConfigured(
|
||||
f"BULK_UPDATE_CHUNK_SIZE must be a positive integer or None (found {BULK_UPDATE_CHUNK_SIZE!r})"
|
||||
)
|
||||
CHANGELOG_SKIP_EMPTY_CHANGES = getattr(configuration, 'CHANGELOG_SKIP_EMPTY_CHANGES', True)
|
||||
CENSUS_REPORTING_ENABLED = getattr(configuration, 'CENSUS_REPORTING_ENABLED', True)
|
||||
CORS_ORIGIN_ALLOW_ALL = getattr(configuration, 'CORS_ORIGIN_ALLOW_ALL', False)
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from django.db.models.signals import post_delete, post_save, pre_delete
|
|||
from netbox.registry import registry
|
||||
|
||||
from .fields import CounterCacheField
|
||||
from .querysets import chunked_update
|
||||
|
||||
|
||||
def get_counters_for_model(model):
|
||||
|
|
@ -37,7 +38,7 @@ def update_counts(model, field_name, related_query):
|
|||
subquery = Subquery(
|
||||
model.objects.filter(pk=OuterRef('pk')).annotate(_count=Count(related_query)).values('_count')
|
||||
)
|
||||
return model.objects.update(**{
|
||||
return chunked_update(model.objects.all(), **{
|
||||
field_name: subquery
|
||||
})
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
from django.db.models import Prefetch, QuerySet
|
||||
from django.conf import settings
|
||||
from django.db import router, transaction
|
||||
from django.db.models import Max, Prefetch, QuerySet
|
||||
|
||||
from users.constants import CONSTRAINT_TOKEN_USER
|
||||
from utilities.permissions import get_permission_for_model, permission_is_exempt, qs_filter_from_constraints
|
||||
|
|
@ -6,9 +8,70 @@ from utilities.permissions import get_permission_for_model, permission_is_exempt
|
|||
__all__ = (
|
||||
'RestrictedPrefetch',
|
||||
'RestrictedQuerySet',
|
||||
'chunked_update',
|
||||
)
|
||||
|
||||
|
||||
def chunked_update(queryset, chunk_size=None, **kwargs):
|
||||
"""
|
||||
Perform a bulk UPDATE on the given queryset, optionally splitting it into batches of at most
|
||||
`chunk_size` rows. Bounding the number of rows touched by each statement avoids exceeding the
|
||||
database's statement timeout when updating very large tables. Batches are selected via keyset
|
||||
pagination on the primary key and wrapped in a transaction so that the operation remains atomic,
|
||||
as it would be when performed by a single UPDATE. Returns the total number of rows updated
|
||||
(matching the return value of QuerySet.update()).
|
||||
|
||||
If `chunk_size` is None, it falls back to the BULK_UPDATE_CHUNK_SIZE configuration parameter
|
||||
(5000 by default). If that is also None, a single unbounded UPDATE is issued, identical to
|
||||
calling queryset.update(**kwargs) directly.
|
||||
|
||||
:param queryset: The QuerySet identifying the rows to update
|
||||
:param chunk_size: The maximum number of rows to update per statement (defaults to
|
||||
settings.BULK_UPDATE_CHUNK_SIZE)
|
||||
"""
|
||||
if chunk_size is None:
|
||||
chunk_size = settings.BULK_UPDATE_CHUNK_SIZE
|
||||
if chunk_size is not None and (type(chunk_size) is not int or chunk_size < 1):
|
||||
raise ValueError(f"chunk_size must be a positive integer or None (found {chunk_size!r})")
|
||||
if chunk_size is None:
|
||||
return queryset.update(**kwargs)
|
||||
|
||||
model = queryset.model
|
||||
|
||||
# Pin the entire operation to a single write database so that the PK lookups, the UPDATE
|
||||
# statements, and the enclosing transaction all use the same connection. This preserves an
|
||||
# explicit .using() on the queryset and otherwise honors the router's write destination.
|
||||
using = queryset._db or router.db_for_write(model)
|
||||
|
||||
count = 0
|
||||
last_pk = 0
|
||||
# Upper bound on the PKs to process. Established lazily (see below) only once a second batch is
|
||||
# known to be needed, so the common single-batch case incurs no extra aggregate query.
|
||||
max_pk = None
|
||||
with transaction.atomic(using=using):
|
||||
while True:
|
||||
batch = queryset.using(using).filter(pk__gt=last_pk).order_by('pk')
|
||||
if max_pk is not None:
|
||||
batch = batch.filter(pk__lte=max_pk)
|
||||
pks = list(batch.values_list('pk', flat=True)[:chunk_size])
|
||||
if not pks:
|
||||
break
|
||||
# Re-filter the original queryset by pk__in (rather than the model's default manager) so
|
||||
# that its own filters are preserved and rows no longer matching them are left untouched.
|
||||
count += queryset.using(using).filter(pk__in=pks).update(**kwargs)
|
||||
last_pk = pks[-1]
|
||||
# A batch shorter than chunk_size means the rows are exhausted; stop without issuing a
|
||||
# trailing (empty) lookup. This keeps a single-batch update to one SELECT and one UPDATE.
|
||||
if len(pks) < chunk_size:
|
||||
break
|
||||
# A full batch means more rows may remain. Capture the current maximum PK as an upper
|
||||
# bound (once) so that rows inserted while the operation runs cannot keep extending it.
|
||||
if max_pk is None:
|
||||
max_pk = queryset.using(using).aggregate(_max=Max('pk'))['_max']
|
||||
|
||||
return count
|
||||
|
||||
|
||||
class RestrictedPrefetch(Prefetch):
|
||||
"""
|
||||
Extend Django's Prefetch to accept a user and action to be passed to the
|
||||
|
|
|
|||
|
|
@ -0,0 +1,116 @@
|
|||
from django.db import connection
|
||||
from django.db.models import Count, F, IntegerField, OuterRef, Subquery
|
||||
from django.db.models.functions import Coalesce
|
||||
from django.test import TestCase, override_settings
|
||||
from django.test.utils import CaptureQueriesContext
|
||||
|
||||
from extras.models import Tag
|
||||
from utilities.querysets import chunked_update
|
||||
|
||||
|
||||
class ChunkedUpdateTestCase(TestCase):
|
||||
"""
|
||||
Tests for the chunked_update() helper, which performs a bulk UPDATE optionally split into
|
||||
batches bounded by the BULK_UPDATE_CHUNK_SIZE configuration parameter.
|
||||
"""
|
||||
@classmethod
|
||||
def setUpTestData(cls):
|
||||
Tag.objects.bulk_create([
|
||||
Tag(name=f'Tag {i}', slug=f'tag-{i}', weight=i)
|
||||
for i in range(1, 6) # Five tags, weights 1..5
|
||||
])
|
||||
|
||||
@staticmethod
|
||||
def _count_updates(queries):
|
||||
return len([q for q in queries if q['sql'].strip().upper().startswith('UPDATE')])
|
||||
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=None)
|
||||
def test_update_without_chunk_size(self):
|
||||
"""
|
||||
With BULK_UPDATE_CHUNK_SIZE set to None, a single unbounded UPDATE is issued.
|
||||
"""
|
||||
with CaptureQueriesContext(connection) as queries:
|
||||
count = chunked_update(Tag.objects.all(), weight=100)
|
||||
|
||||
self.assertEqual(count, 5)
|
||||
self.assertEqual(self._count_updates(queries.captured_queries), 1)
|
||||
self.assertEqual(Tag.objects.filter(weight=100).count(), 5)
|
||||
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=2)
|
||||
def test_update_with_chunk_size(self):
|
||||
"""
|
||||
With BULK_UPDATE_CHUNK_SIZE set, the update is split into batches; every row is updated
|
||||
exactly once and the total count is returned.
|
||||
"""
|
||||
with CaptureQueriesContext(connection) as queries:
|
||||
count = chunked_update(Tag.objects.all(), weight=100)
|
||||
|
||||
self.assertEqual(count, 5)
|
||||
# Five rows in batches of two → three UPDATE statements
|
||||
self.assertEqual(self._count_updates(queries.captured_queries), 3)
|
||||
self.assertEqual(Tag.objects.filter(weight=100).count(), 5)
|
||||
|
||||
def test_explicit_chunk_size_argument(self):
|
||||
"""
|
||||
An explicit chunk_size argument takes precedence over the configuration parameter.
|
||||
"""
|
||||
with CaptureQueriesContext(connection) as queries:
|
||||
count = chunked_update(Tag.objects.all(), chunk_size=2, weight=100)
|
||||
|
||||
self.assertEqual(count, 5)
|
||||
self.assertEqual(self._count_updates(queries.captured_queries), 3)
|
||||
self.assertEqual(Tag.objects.filter(weight=100).count(), 5)
|
||||
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=2)
|
||||
def test_f_expression_applied_once_per_row(self):
|
||||
"""
|
||||
An F() expression referencing the row's own column is applied exactly once per row, even
|
||||
when the update is chunked (chunks are disjoint by primary key).
|
||||
"""
|
||||
original = {tag.pk: tag.weight for tag in Tag.objects.all()}
|
||||
|
||||
count = chunked_update(Tag.objects.all(), weight=F('weight') + 1)
|
||||
|
||||
self.assertEqual(count, 5)
|
||||
for tag in Tag.objects.all():
|
||||
self.assertEqual(tag.weight, original[tag.pk] + 1)
|
||||
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=2)
|
||||
def test_correlated_subquery(self):
|
||||
"""
|
||||
A correlated subquery (OuterRef) resolves against each chunk's queryset, mirroring the
|
||||
counter-rebuild pattern in utilities.counters.update_counts().
|
||||
"""
|
||||
# Set each tag's weight to the number of tags sharing its slug (always 1), proving the
|
||||
# OuterRef binds correctly per-row across chunks.
|
||||
subquery = Subquery(
|
||||
Tag.objects.filter(slug=OuterRef('slug')).values('slug')
|
||||
.annotate(c=Count('pk')).values('c'),
|
||||
output_field=IntegerField()
|
||||
)
|
||||
count = chunked_update(Tag.objects.all(), chunk_size=2, weight=Coalesce(subquery, 0))
|
||||
|
||||
self.assertEqual(count, 5)
|
||||
self.assertEqual(Tag.objects.filter(weight=1).count(), 5)
|
||||
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=2)
|
||||
def test_filtered_queryset(self):
|
||||
"""
|
||||
Only rows matching the queryset's filter are updated when chunking.
|
||||
"""
|
||||
target_pks = list(Tag.objects.filter(weight__lte=3).values_list('pk', flat=True))
|
||||
|
||||
count = chunked_update(Tag.objects.filter(weight__lte=3), weight=0)
|
||||
|
||||
self.assertEqual(count, len(target_pks))
|
||||
self.assertEqual(Tag.objects.filter(weight=0).count(), len(target_pks))
|
||||
# Rows outside the filter are untouched (weights 4 and 5 remain)
|
||||
self.assertEqual(Tag.objects.filter(weight__gt=3).count(), 2)
|
||||
|
||||
@override_settings(BULK_UPDATE_CHUNK_SIZE=2)
|
||||
def test_empty_queryset(self):
|
||||
"""
|
||||
Updating an empty queryset is a no-op that returns zero.
|
||||
"""
|
||||
count = chunked_update(Tag.objects.filter(name='nonexistent'), weight=0)
|
||||
self.assertEqual(count, 0)
|
||||
|
|
@ -2,6 +2,8 @@ from django.db.models import Sum
|
|||
from django.db.models.signals import post_delete, post_save
|
||||
from django.dispatch import receiver
|
||||
|
||||
from utilities.querysets import chunked_update
|
||||
|
||||
from .models import Cluster, VirtualDisk, VirtualMachine
|
||||
|
||||
|
||||
|
|
@ -22,4 +24,4 @@ def update_virtualmachine_site(instance, **kwargs):
|
|||
Update the assigned site for all VMs to match that of the Cluster (if any).
|
||||
"""
|
||||
if instance._site:
|
||||
VirtualMachine.objects.filter(cluster=instance).update(site=instance._site)
|
||||
chunked_update(VirtualMachine.objects.filter(cluster=instance), site=instance._site)
|
||||
|
|
|
|||
Loading…
Reference in New Issue