perf: batch export graph database queries
This commit is contained in:
+178
-87
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
from collections import defaultdict
|
||||
|
||||
from django.apps import apps
|
||||
from django.contrib.contenttypes.fields import GenericForeignKey
|
||||
@@ -55,6 +55,12 @@ def _descendants(obj):
|
||||
return [obj]
|
||||
|
||||
|
||||
def _chunks(values, size):
|
||||
values = list(values)
|
||||
for offset in range(0, len(values), size):
|
||||
yield values[offset : offset + size]
|
||||
|
||||
|
||||
def seed_scope(scope_type: str, scope_id: int):
|
||||
from dcim.models import Location, Region, Site
|
||||
from tenancy.models import Tenant, TenantGroup
|
||||
@@ -80,19 +86,45 @@ def seed_scope(scope_type: str, scope_id: int):
|
||||
|
||||
|
||||
class ObjectGraph:
|
||||
"""Collect scoped objects first and their forward dependencies second."""
|
||||
"""Collect scoped objects and dependencies using batched relation queries."""
|
||||
|
||||
def __init__(self, max_objects: int):
|
||||
def __init__(self, max_objects: int, query_batch_size: int = 500):
|
||||
self.max_objects = max_objects
|
||||
self.query_batch_size = max(1, query_batch_size)
|
||||
self.members: dict[str, models.Model] = {}
|
||||
self.dependencies: dict[str, models.Model] = {}
|
||||
self._custom_field_cache = {}
|
||||
self._models = exportable_models()
|
||||
self._reverse_fields, self._generic_fields, self._peer_fields = self._build_relation_indexes()
|
||||
|
||||
@property
|
||||
def objects(self) -> dict[str, models.Model]:
|
||||
return {**self.members, **self.dependencies}
|
||||
|
||||
def _build_relation_indexes(self):
|
||||
reverse_fields = defaultdict(lambda: defaultdict(list))
|
||||
generic_fields = defaultdict(list)
|
||||
peer_fields = defaultdict(list)
|
||||
for candidate_model in self._models:
|
||||
for relation_field in candidate_model._meta.concrete_fields:
|
||||
if not isinstance(relation_field, (models.ForeignKey, models.OneToOneField)):
|
||||
continue
|
||||
related_model = relation_field.remote_field.model
|
||||
if (
|
||||
relation_field.remote_field.on_delete in (models.CASCADE, models.PROTECT)
|
||||
or relation_field.name in SCOPE_LINK_FIELDS
|
||||
):
|
||||
reverse_fields[related_model][candidate_model].append(relation_field)
|
||||
if related_model._meta.label_lower in PEER_CONTAINER_MODELS:
|
||||
peer_fields[candidate_model].append(relation_field)
|
||||
for generic_field in candidate_model._meta.private_fields:
|
||||
if isinstance(generic_field, GenericForeignKey):
|
||||
generic_fields[candidate_model].append(generic_field)
|
||||
return reverse_fields, generic_fields, peer_fields
|
||||
|
||||
def _check_limit(self):
|
||||
if len(self.objects) > self.max_objects:
|
||||
object_count = len(self.members) + len(self.dependencies)
|
||||
if object_count > self.max_objects:
|
||||
raise GraphLimitError(
|
||||
f"Der Export würde mehr als {self.max_objects} Objekte enthalten. "
|
||||
"Bitte den Bereich verkleinern oder max_objects erhöhen."
|
||||
@@ -115,92 +147,123 @@ class ObjectGraph:
|
||||
self._check_limit()
|
||||
return True
|
||||
|
||||
def collect(self, seeds):
|
||||
queue = deque()
|
||||
def _group_by_model(self, objects):
|
||||
grouped = defaultdict(list)
|
||||
for obj in objects:
|
||||
grouped[type(obj)].append(obj)
|
||||
return grouped
|
||||
|
||||
def _expand_members(self, seeds):
|
||||
pending = {}
|
||||
for obj in seeds:
|
||||
if self.add_member(obj):
|
||||
queue.append(obj)
|
||||
pending[object_key(obj)] = obj
|
||||
|
||||
models_to_scan = exportable_models()
|
||||
while True:
|
||||
while queue:
|
||||
parent = queue.popleft()
|
||||
parent_model = type(parent)
|
||||
while pending:
|
||||
next_pending = {}
|
||||
for parent_model, parents in self._group_by_model(pending.values()).items():
|
||||
parent_ct = ContentType.objects.get_for_model(parent_model)
|
||||
for candidate_model in models_to_scan:
|
||||
query = models.Q()
|
||||
for field in candidate_model._meta.concrete_fields:
|
||||
if not isinstance(field, (models.ForeignKey, models.OneToOneField)):
|
||||
continue
|
||||
if field.remote_field.model is not parent_model:
|
||||
continue
|
||||
if (
|
||||
field.remote_field.on_delete not in (models.CASCADE, models.PROTECT)
|
||||
and field.name not in SCOPE_LINK_FIELDS
|
||||
):
|
||||
continue
|
||||
query |= models.Q(**{field.attname: parent.pk})
|
||||
for field in candidate_model._meta.private_fields:
|
||||
if isinstance(field, GenericForeignKey):
|
||||
query |= models.Q(**{field.ct_field: parent_ct, field.fk_field: parent.pk})
|
||||
if not query:
|
||||
continue
|
||||
for child in candidate_model.objects.filter(query).distinct().iterator():
|
||||
if self.add_member(child):
|
||||
queue.append(child)
|
||||
direct_candidates = self._reverse_fields.get(parent_model, {})
|
||||
peer_fields = self._peer_fields.get(parent_model, [])
|
||||
for parent_batch in _chunks(parents, self.query_batch_size):
|
||||
parent_ids = [obj.pk for obj in parent_batch]
|
||||
|
||||
promoted = False
|
||||
for obj in list(self.members.values()):
|
||||
for field in obj._meta.concrete_fields:
|
||||
if not isinstance(field, (models.ForeignKey, models.OneToOneField)):
|
||||
continue
|
||||
related = getattr(obj, field.name, None)
|
||||
if related is None or related._meta.label_lower not in PEER_CONTAINER_MODELS:
|
||||
continue
|
||||
if self.add_member(related):
|
||||
queue.append(related)
|
||||
promoted = True
|
||||
if not promoted:
|
||||
break
|
||||
if peer_fields:
|
||||
peer_names = [field.name for field in peer_fields]
|
||||
queryset = parent_model._default_manager.filter(pk__in=parent_ids).select_related(*peer_names)
|
||||
for parent in queryset.iterator(chunk_size=self.query_batch_size):
|
||||
for relation_field in peer_fields:
|
||||
related = getattr(parent, relation_field.name, None)
|
||||
if related is not None and self.add_member(related):
|
||||
next_pending[object_key(related)] = related
|
||||
|
||||
dependency_queue = deque(self.members.values())
|
||||
scanned = set()
|
||||
while dependency_queue:
|
||||
obj = dependency_queue.popleft()
|
||||
key = object_key(obj)
|
||||
if key in scanned:
|
||||
continue
|
||||
scanned.add(key)
|
||||
related_objects = []
|
||||
for field in obj._meta.concrete_fields:
|
||||
if isinstance(field, (models.ForeignKey, models.OneToOneField)):
|
||||
related = getattr(obj, field.name, None)
|
||||
if related is not None:
|
||||
related_objects.append(related)
|
||||
for field in obj._meta.private_fields:
|
||||
if isinstance(field, GenericForeignKey):
|
||||
related = getattr(obj, field.name, None)
|
||||
if related is not None:
|
||||
related_objects.append(related)
|
||||
for field in obj._meta.many_to_many:
|
||||
try:
|
||||
related_objects.extend(getattr(obj, field.name).all())
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
if hasattr(obj, "custom_field_data"):
|
||||
candidate_queries = defaultdict(models.Q)
|
||||
for candidate_model, relation_fields in direct_candidates.items():
|
||||
for relation_field in relation_fields:
|
||||
candidate_queries[candidate_model] |= models.Q(**{f"{relation_field.attname}__in": parent_ids})
|
||||
for candidate_model, generic_fields in self._generic_fields.items():
|
||||
for generic_field in generic_fields:
|
||||
candidate_queries[candidate_model] |= models.Q(
|
||||
**{
|
||||
generic_field.ct_field: parent_ct,
|
||||
f"{generic_field.fk_field}__in": parent_ids,
|
||||
}
|
||||
)
|
||||
|
||||
for candidate_model, query in candidate_queries.items():
|
||||
queryset = candidate_model._default_manager.filter(query).distinct()
|
||||
for child in queryset.iterator(chunk_size=self.query_batch_size):
|
||||
if self.add_member(child):
|
||||
next_pending[object_key(child)] = child
|
||||
pending = next_pending
|
||||
|
||||
def custom_fields_for_model(self, model):
|
||||
if model not in self._custom_field_cache:
|
||||
if hasattr(model, "custom_field_data"):
|
||||
from extras.models import CustomField
|
||||
|
||||
custom_fields = list(CustomField.objects.get_for_model(type(obj)))
|
||||
related_objects.extend(custom_fields)
|
||||
for custom_field in custom_fields:
|
||||
if custom_field.type not in ("object", "multiobject") or not custom_field.related_object_type:
|
||||
continue
|
||||
raw_value = obj.custom_field_data.get(custom_field.name)
|
||||
if raw_value in (None, "", []):
|
||||
continue
|
||||
target_model = custom_field.related_object_type.model_class()
|
||||
target_ids = raw_value if custom_field.type == "multiobject" else [raw_value]
|
||||
related_objects.extend(target_model._default_manager.filter(pk__in=target_ids))
|
||||
self._custom_field_cache[model] = list(CustomField.objects.get_for_model(model))
|
||||
else:
|
||||
self._custom_field_cache[model] = []
|
||||
return self._custom_field_cache[model]
|
||||
|
||||
def _load_objects(self, model, objects):
|
||||
foreign_keys = [
|
||||
field.name
|
||||
for field in model._meta.concrete_fields
|
||||
if isinstance(field, (models.ForeignKey, models.OneToOneField))
|
||||
]
|
||||
prefetch_fields = [field.name for field in model._meta.many_to_many]
|
||||
prefetch_fields.extend(
|
||||
field.name for field in model._meta.private_fields if isinstance(field, GenericForeignKey)
|
||||
)
|
||||
for object_batch in _chunks(objects, self.query_batch_size):
|
||||
object_ids = [obj.pk for obj in object_batch]
|
||||
queryset = model._default_manager.filter(pk__in=object_ids).select_related(*foreign_keys).order_by("pk")
|
||||
if prefetch_fields:
|
||||
queryset = queryset.prefetch_related(*prefetch_fields)
|
||||
yield from queryset.iterator(chunk_size=self.query_batch_size)
|
||||
|
||||
def iter_loaded_objects(self):
|
||||
grouped = self._group_by_model(self.objects.values())
|
||||
for model in sorted(grouped, key=lambda item: item._meta.label_lower):
|
||||
yield from self._load_objects(model, grouped[model])
|
||||
|
||||
def _related_objects(self, obj, custom_fields):
|
||||
related_objects = []
|
||||
for relation_field in obj._meta.concrete_fields:
|
||||
if isinstance(relation_field, (models.ForeignKey, models.OneToOneField)):
|
||||
related = getattr(obj, relation_field.name, None)
|
||||
if related is not None:
|
||||
related_objects.append(related)
|
||||
for generic_field in obj._meta.private_fields:
|
||||
if isinstance(generic_field, GenericForeignKey):
|
||||
related = getattr(obj, generic_field.name, None)
|
||||
if related is not None:
|
||||
related_objects.append(related)
|
||||
for many_to_many_field in obj._meta.many_to_many:
|
||||
try:
|
||||
related_objects.extend(getattr(obj, many_to_many_field.name).all())
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
related_objects.extend(custom_fields)
|
||||
return related_objects
|
||||
|
||||
def _custom_value_dependencies(self, objects, custom_fields):
|
||||
target_ids_by_model = defaultdict(set)
|
||||
for custom_field in custom_fields:
|
||||
if custom_field.type not in ("object", "multiobject") or not custom_field.related_object_type:
|
||||
continue
|
||||
target_model = custom_field.related_object_type.model_class()
|
||||
for obj in objects:
|
||||
raw_value = obj.custom_field_data.get(custom_field.name)
|
||||
if raw_value in (None, "", []):
|
||||
continue
|
||||
values = raw_value if custom_field.type == "multiobject" else [raw_value]
|
||||
target_ids_by_model[target_model].update(values)
|
||||
|
||||
for obj in objects:
|
||||
if (
|
||||
obj._meta.label_lower == "extras.customfield"
|
||||
and obj.type in ("object", "multiobject")
|
||||
@@ -208,9 +271,37 @@ class ObjectGraph:
|
||||
and obj.default not in (None, "", [])
|
||||
):
|
||||
target_model = obj.related_object_type.model_class()
|
||||
target_ids = obj.default if obj.type == "multiobject" else [obj.default]
|
||||
related_objects.extend(target_model._default_manager.filter(pk__in=target_ids))
|
||||
for related in related_objects:
|
||||
if self.add_dependency(related):
|
||||
dependency_queue.append(related)
|
||||
values = obj.default if obj.type == "multiobject" else [obj.default]
|
||||
target_ids_by_model[target_model].update(values)
|
||||
|
||||
for target_model, target_ids in target_ids_by_model.items():
|
||||
for target_id_batch in _chunks(target_ids, self.query_batch_size):
|
||||
yield from target_model._default_manager.filter(pk__in=target_id_batch).iterator(
|
||||
chunk_size=self.query_batch_size
|
||||
)
|
||||
|
||||
def _collect_dependencies(self):
|
||||
pending = dict(self.members)
|
||||
scanned = set()
|
||||
while pending:
|
||||
next_pending = {}
|
||||
for model, objects in self._group_by_model(pending.values()).items():
|
||||
custom_fields = self.custom_fields_for_model(model)
|
||||
loaded_objects = list(self._load_objects(model, objects))
|
||||
for related in self._custom_value_dependencies(loaded_objects, custom_fields):
|
||||
if self.add_dependency(related):
|
||||
next_pending[object_key(related)] = related
|
||||
for obj in loaded_objects:
|
||||
key = object_key(obj)
|
||||
if key in scanned:
|
||||
continue
|
||||
scanned.add(key)
|
||||
for related in self._related_objects(obj, custom_fields):
|
||||
if self.add_dependency(related):
|
||||
next_pending[object_key(related)] = related
|
||||
pending = next_pending
|
||||
|
||||
def collect(self, seeds):
|
||||
self._expand_members(seeds)
|
||||
self._collect_dependencies()
|
||||
return self
|
||||
|
||||
Reference in New Issue
Block a user