217 lines
8.1 KiB
Python
217 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
|
|
from django.apps import apps
|
|
from django.contrib.contenttypes.fields import GenericForeignKey
|
|
from django.contrib.contenttypes.models import ContentType
|
|
from django.db import models
|
|
|
|
from .codec import object_key
|
|
from .exceptions import GraphLimitError
|
|
|
|
EXCLUDED_APP_LABELS = {
|
|
"account",
|
|
"admin",
|
|
"auth",
|
|
"contenttypes",
|
|
"sessions",
|
|
"users",
|
|
"netbox_export",
|
|
}
|
|
EXCLUDED_MODELS = {
|
|
"core.job",
|
|
"core.objectchange",
|
|
"extras.eventrule",
|
|
"extras.journalentry",
|
|
"extras.notification",
|
|
"extras.notificationgroup",
|
|
"extras.savedfilter",
|
|
"extras.subscription",
|
|
}
|
|
SCOPE_LINK_FIELDS = {"tenant", "site", "location", "region"}
|
|
PEER_CONTAINER_MODELS = {"dcim.cable", "circuits.circuit", "circuits.virtualcircuit"}
|
|
|
|
|
|
def is_exportable_model(model) -> bool:
|
|
opts = model._meta
|
|
return bool(
|
|
opts.managed
|
|
and not opts.abstract
|
|
and not opts.proxy
|
|
and not opts.auto_created
|
|
and opts.app_label not in EXCLUDED_APP_LABELS
|
|
and opts.label_lower not in EXCLUDED_MODELS
|
|
)
|
|
|
|
|
|
def exportable_models():
|
|
return tuple(model for model in apps.get_models() if is_exportable_model(model))
|
|
|
|
|
|
def _descendants(obj):
|
|
if hasattr(obj, "get_descendants"):
|
|
return list(obj.get_descendants(include_self=True))
|
|
return [obj]
|
|
|
|
|
|
def seed_scope(scope_type: str, scope_id: int):
|
|
from dcim.models import Location, Region, Site
|
|
from tenancy.models import Tenant, TenantGroup
|
|
|
|
models_by_scope = {
|
|
"tenant_group": TenantGroup,
|
|
"tenant": Tenant,
|
|
"region": Region,
|
|
"site": Site,
|
|
"location": Location,
|
|
}
|
|
model = models_by_scope[scope_type]
|
|
root = model.objects.get(pk=scope_id)
|
|
seeds = _descendants(root)
|
|
|
|
if scope_type == "tenant_group":
|
|
group_ids = [obj.pk for obj in seeds]
|
|
seeds.extend(Tenant.objects.filter(group_id__in=group_ids))
|
|
elif scope_type == "region":
|
|
region_ids = [obj.pk for obj in seeds]
|
|
seeds.extend(Site.objects.filter(region_id__in=region_ids))
|
|
return root, seeds
|
|
|
|
|
|
class ObjectGraph:
|
|
"""Collect scoped objects first and their forward dependencies second."""
|
|
|
|
def __init__(self, max_objects: int):
|
|
self.max_objects = max_objects
|
|
self.members: dict[str, models.Model] = {}
|
|
self.dependencies: dict[str, models.Model] = {}
|
|
|
|
@property
|
|
def objects(self) -> dict[str, models.Model]:
|
|
return {**self.members, **self.dependencies}
|
|
|
|
def _check_limit(self):
|
|
if len(self.objects) > 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."
|
|
)
|
|
|
|
def add_member(self, obj) -> bool:
|
|
key = object_key(obj)
|
|
if key in self.members:
|
|
return False
|
|
self.dependencies.pop(key, None)
|
|
self.members[key] = obj
|
|
self._check_limit()
|
|
return True
|
|
|
|
def add_dependency(self, obj) -> bool:
|
|
key = object_key(obj)
|
|
if key in self.members or key in self.dependencies or not is_exportable_model(type(obj)):
|
|
return False
|
|
self.dependencies[key] = obj
|
|
self._check_limit()
|
|
return True
|
|
|
|
def collect(self, seeds):
|
|
queue = deque()
|
|
for obj in seeds:
|
|
if self.add_member(obj):
|
|
queue.append(obj)
|
|
|
|
models_to_scan = exportable_models()
|
|
while True:
|
|
while queue:
|
|
parent = queue.popleft()
|
|
parent_model = type(parent)
|
|
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)
|
|
|
|
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
|
|
|
|
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"):
|
|
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))
|
|
if (
|
|
obj._meta.label_lower == "extras.customfield"
|
|
and obj.type in ("object", "multiobject")
|
|
and obj.related_object_type
|
|
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)
|
|
return self
|