perf: batch export graph database queries

netbox_export/__init__.py

@@ -7,7 +7,7 @@ class NetBoxExportConfig(PluginConfig):
     name = "netbox_export"
     verbose_name = "NetBox-Export"
     description = "Portable ZIP export and import for tenants and locations"
-    version = "0.2.0"
+    version = "0.3.0"
     author = "NetBox Export contributors"
     base_url = "netbox-export"
     min_version = "4.6.0"
@@ -16,6 +16,7 @@ class NetBoxExportConfig(PluginConfig):
     default_settings: ClassVar[dict] = {
         "max_objects": 50000,
         "max_archive_size_mb": 250,
+        "query_batch_size": 500,
         "archive_signing_key": "",
     }
 

netbox_export/services/archive.py

@@ -8,6 +8,7 @@ import posixpath
 import zipfile
 from collections.abc import Iterable
 from dataclasses import dataclass
+from pathlib import PurePosixPath
 from typing import BinaryIO
 
 from django.core.serializers.json import DjangoJSONEncoder
@@ -18,6 +19,19 @@ FORMAT_NAME = "netbox-export"
 FORMAT_VERSION = 1
 MANIFEST_NAME = "manifest.json"
 OBJECTS_NAME = "objects.ndjson"
+PRECOMPRESSED_SUFFIXES = {
+    ".7z",
+    ".avi",
+    ".gz",
+    ".jpeg",
+    ".jpg",
+    ".mp3",
+    ".mp4",
+    ".pdf",
+    ".png",
+    ".webp",
+    ".zip",
+}
 
 
 def _json_bytes(value) -> bytes:
@@ -59,11 +73,21 @@ def build_archive(manifest: dict, records: Iterable[dict], assets: dict[str, byt
         manifest["signature"] = _signature(manifest, signing_key)
 
     output = io.BytesIO()
-    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
+    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=1) as archive:
         archive.writestr(MANIFEST_NAME, _json_bytes(manifest))
         archive.writestr(OBJECTS_NAME, object_data)
         for path, content in sorted(assets.items()):
-            archive.writestr(_safe_member_name(path), content)
+            compression = (
+                zipfile.ZIP_STORED
+                if PurePosixPath(path).suffix.lower() in PRECOMPRESSED_SUFFIXES
+                else zipfile.ZIP_DEFLATED
+            )
+            archive.writestr(
+                _safe_member_name(path),
+                content,
+                compress_type=compression,
+                compresslevel=None if compression == zipfile.ZIP_STORED else 1,
+            )
     return output.getvalue()
 
 

netbox_export/services/codec.py

@@ -87,16 +87,23 @@ def external_identity(obj) -> dict:
     return {"model": label, "lookup": {"pk": encode_scalar(obj.pk)}}
 
 
-def _reference_spec(obj, exported_keys: set[str]):
-    if object_key(obj) in exported_keys:
-        return {"ref": object_key(obj)}
-    return {"external": external_identity(obj)}
+def _reference_for_pk(model, pk, exported_keys: set[str]):
+    key = f"{model._meta.label_lower}:{pk}"
+    if key in exported_keys:
+        return {"ref": key}
+    try:
+        target = model._default_manager.get(pk=pk)
+    except model.DoesNotExist:
+        return None
+    return {"external": external_identity(target)}
 
 
-def _encode_custom_field_data(obj, value: dict, exported_keys: set[str]):
-    from extras.models import CustomField
+def _encode_custom_field_data(obj, value: dict, exported_keys: set[str], custom_fields=None):
+    if custom_fields is None:
+        from extras.models import CustomField
 
-    definitions = {field.name: field for field in CustomField.objects.get_for_model(type(obj))}
+        custom_fields = CustomField.objects.get_for_model(type(obj))
+    definitions = {field.name: field for field in custom_fields}
     encoded = {}
     for name, raw_value in value.items():
         custom_field = definitions.get(name)
@@ -110,21 +117,17 @@ def _encode_custom_field_data(obj, value: dict, exported_keys: set[str]):
             continue
         target_model = custom_field.related_object_type.model_class()
         if custom_field.type == "object":
-            try:
-                target = target_model._default_manager.get(pk=raw_value)
-            except target_model.DoesNotExist:
-                encoded[name] = None
-            else:
-                encoded[name] = {"$type": "object_ref", "value": _reference_spec(target, exported_keys)}
+            reference = _reference_for_pk(target_model, raw_value, exported_keys)
+            encoded[name] = {"$type": "object_ref", "value": reference} if reference else None
         else:
-            targets = {str(item.pk): item for item in target_model._default_manager.filter(pk__in=raw_value)}
+            references = [
+                reference
+                for pk in raw_value
+                if (reference := _reference_for_pk(target_model, pk, exported_keys)) is not None
+            ]
             encoded[name] = {
                 "$type": "multiobject_ref",
-                "value": [
-                    _reference_spec(targets[str(pk)], exported_keys)
-                    for pk in raw_value
-                    if str(pk) in targets
-                ],
+                "value": references,
             }
     return encoded
 
@@ -138,23 +141,20 @@ def _encode_custom_field_default(custom_field, value, exported_keys: set[str]):
         return encode_scalar(value)
     target_model = custom_field.related_object_type.model_class()
     if custom_field.type == "object":
-        try:
-            target = target_model._default_manager.get(pk=value)
-        except target_model.DoesNotExist:
-            return None
-        return {"$type": "object_ref", "value": _reference_spec(target, exported_keys)}
-    targets = {str(item.pk): item for item in target_model._default_manager.filter(pk__in=value)}
+        reference = _reference_for_pk(target_model, value, exported_keys)
+        return {"$type": "object_ref", "value": reference} if reference else None
+    references = [
+        reference
+        for pk in value
+        if (reference := _reference_for_pk(target_model, pk, exported_keys)) is not None
+    ]
     return {
         "$type": "multiobject_ref",
-        "value": [
-            _reference_spec(targets[str(pk)], exported_keys)
-            for pk in value
-            if str(pk) in targets
-        ],
+        "value": references,
     }
 
 
-def serialize_object(obj, exported_keys: set[str], assets: dict[str, bytes]) -> dict:
+def serialize_object(obj, exported_keys: set[str], assets: dict[str, bytes], custom_fields=None) -> dict:
     record_id = object_key(obj)
     gfk_fields = generic_foreign_keys(type(obj))
     gfk_storage = {name for field in gfk_fields for name in (field.ct_field, field.fk_field)}
@@ -191,7 +191,7 @@ def serialize_object(obj, exported_keys: set[str], assets: dict[str, bytes]) ->
             continue
         value = field.value_from_object(obj)
         if field.name == "custom_field_data" and isinstance(value, dict):
-            scalars[field.name] = _encode_custom_field_data(obj, value, exported_keys)
+            scalars[field.name] = _encode_custom_field_data(obj, value, exported_keys, custom_fields)
         elif obj._meta.label_lower == "extras.customfield" and field.name == "default":
             scalars[field.name] = _encode_custom_field_default(obj, value, exported_keys)
         else:

netbox_export/services/exporter.py

@@ -12,21 +12,34 @@ from .codec import serialize_object
 from .graph import ObjectGraph, seed_scope
 
 
-def export_scope(scope_type: str, scope_id: int, *, max_objects: int, signing_key: str = ""):
+def export_scope(
+    scope_type: str,
+    scope_id: int,
+    *,
+    max_objects: int,
+    query_batch_size: int = 500,
+    signing_key: str = "",
+):
     root, seeds = seed_scope(scope_type, scope_id)
-    graph = ObjectGraph(max_objects=max_objects).collect(seeds)
+    graph = ObjectGraph(max_objects=max_objects, query_batch_size=query_batch_size).collect(seeds)
     objects = graph.objects
+    exported_keys = set(objects)
     assets = {}
     records = [
-        serialize_object(obj, set(objects), assets)
-        for _, obj in sorted(objects.items())
+        serialize_object(
+            obj,
+            exported_keys,
+            assets,
+            custom_fields=graph.custom_fields_for_model(type(obj)),
+        )
+        for obj in graph.iter_loaded_objects()
     ]
     counts = Counter(record["model"] for record in records)
     manifest = {
         "created_at": datetime.now(UTC).isoformat(),
         "source_instance": str(InstanceIdentity.local_id()),
         "source_netbox_version": getattr(getattr(settings, "RELEASE", None), "version", "4.6"),
-        "plugin_version": "0.2.0",
+        "plugin_version": "0.3.0",
         "scope": {
             "type": scope_type,
             "source_pk": str(scope_id),

netbox_export/services/graph.py

@@ -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

netbox_export/views.py

@@ -48,6 +48,7 @@ class DashboardView(UserPassesTestMixin, View):
                 form.cleaned_data["scope_type"],
                 scope.pk,
                 max_objects=int(get_plugin_config("netbox_export", "max_objects", 50000)),
+                query_batch_size=int(get_plugin_config("netbox_export", "query_batch_size", 500)),
                 signing_key=get_plugin_config("netbox_export", "archive_signing_key", ""),
             )
         except ExportImportError as exc: