working views

validity/forms/general.py

@@ -1,12 +1,11 @@
 from core.forms.mixins import SyncedDataMixin
-from core.models import DataSource
 from dcim.models import DeviceType, Location, Manufacturer, Platform, Site
-from django.forms import CharField, PasswordInput, Textarea, ValidationError
+from django.forms import CharField, Textarea, ValidationError
 from django.utils.translation import gettext_lazy as _
 from extras.models import Tag
 from netbox.forms import NetBoxModelForm
 from tenancy.models import Tenant
-from utilities.forms.fields import DynamicModelChoiceField, DynamicModelMultipleChoiceField
+from utilities.forms.fields import DynamicModelMultipleChoiceField
 
 from validity import models
 

validity/managers.py

@@ -1,5 +1,5 @@
 from itertools import chain
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TypeVar
 
 from django.contrib.postgres.aggregates import ArrayAgg
 from django.db.models import (
@@ -10,40 +10,38 @@
     ExpressionWrapper,
     F,
     FloatField,
-    OuterRef,
     Prefetch,
     Q,
     QuerySet,
     Value,
     When,
 )
 from django.db.models.fields.json import KeyTextTransform
-from django.db.models.functions import Cast, JSONObject
+from django.db.models.functions import Cast
 from netbox.models import RestrictedQuerySet
 
 from validity import settings
 from validity.choices import DeviceGroupByChoices, SeverityChoices
 from validity.utils.orm import QuerySetMap, RegexpReplace
 
 
-if TYPE_CHECKING:
-    from validity.models.base import BaseModel
-
-
 class CustomPrefetchMixin(QuerySet):
     """
     Allows to prefetch objects without direct relations
+    Many-objects are prefetched as generators
     """
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.custom_prefetches = {}
 
-    def custom_prefetch(self, field: str, prefetch_qs: QuerySet):
+    def custom_prefetch(self, field: str, prefetch_qs: QuerySet, many: bool = False):
         pk_field = field + "_id"
         pk_values = self.values_list(pk_field, flat=True)
+        if many:
+            pk_values = chain.from_iterable(pk_values)
         prefetched_objects = prefetch_qs.filter(pk__in=pk_values)
-        self.custom_prefetches[field] = QuerySetMap(prefetched_objects)
+        self.custom_prefetches[field] = (many, QuerySetMap(prefetched_objects))
         return self
 
     def _clone(self, *args, **kwargs):
@@ -56,8 +54,13 @@ def _fetch_all(self):
         for item in self._result_cache:
             if not isinstance(item, self.model):
                 continue
-            for prefetched_field, qs_dict in self.custom_prefetches.items():
-                setattr(item, prefetched_field, qs_dict[item.pk])
+            for prefetched_field, (many, qs_dict) in self.custom_prefetches.items():
+                prefetch_pk_values = getattr(item, prefetched_field + "_id")
+                if many:
+                    prefetch_values = (qs_dict[pk] for pk in prefetch_pk_values)
+                else:
+                    prefetch_values = qs_dict[prefetch_pk_values]
+                setattr(item, prefetched_field, prefetch_values)
 
 
 class ComplianceTestQS(RestrictedQuerySet):
@@ -102,19 +105,6 @@ def delete_old(self):
         return self.filter(report=None).last_more_than(settings.store_last_results).delete()
 
 
-class JSONObjMixin:
-    def as_json(self):
-        return self.values(json=JSONObject(**{f: f for f in self.model.json_fields}))
-
-
-class ConfigSerializerQS(JSONObjMixin, RestrictedQuerySet):
-    pass
-
-
-class NameSetQS(JSONObjMixin, RestrictedQuerySet):
-    pass
-
-
 def percentage(field1: str, field2: str) -> Case:
     return Case(
         When(Q(**{f"{field2}__gt": 0}), then=Value(100.0) * F(field1) / F(field2)),
@@ -127,16 +117,20 @@ class VDataFileQS(RestrictedQuerySet):
     pass
 
 
-class VDataSourceQS(RestrictedQuerySet):
+class VDataSourceQS(CustomPrefetchMixin, RestrictedQuerySet):
     def annotate_config_path(self):
         return self.annotate(device_config_path=KeyTextTransform("device_config_path", "custom_field_data"))
 
     def prefetch_config_files(self):
         from validity.models import VDataFile
 
-        config_path = RegexpReplace(OuterRef("device_config_path"), Value("{{.+?}}"), Value(".+?"))
-        return self.annotate_config_path().prefetch_related(
-            Prefetch("datafiles", VDataFile.objects.filter(path__regex=config_path), to_attr="config_files")
+        config_path = RegexpReplace(F("device_config_path"), Value("{{.+?}}"), Value(".+?"))
+        path_filter = Q(datafiles__path__regex=config_path)
+        return (
+            self.annotate_config_path()
+            .annotate(config_files_id=ArrayAgg(F("datafiles__pk"), filter=path_filter))
+            .annotate(config_file_count=Count("datafiles__pk", filter=path_filter))
+            .custom_prefetch("config_files", VDataFile.objects.all(), many=True)
         )
 
 
@@ -177,10 +171,6 @@ def delete_old(self):
 _QS = TypeVar("_QS", bound=QuerySet)
 
 
-def annotate_json(qs: _QS, field: str, annotate_model: type["BaseModel"]) -> _QS:
-    return qs.annotate(**{f"json_{field}": annotate_model.objects.filter(pk=OuterRef(f"{field}_id")).as_json()})
-
-
 class VDeviceQS(CustomPrefetchMixin, RestrictedQuerySet):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)

validity/models/base.py

@@ -1,5 +1,4 @@
 import logging
-from collections.abc import Iterable
 
 from django.db import models
 from django.urls import reverse

validity/models/data.py

@@ -23,16 +23,11 @@ def configfiles_by_path(self) -> QuerySetMap:
         Returns "path: datafile" mapping for config files
         """
         assert hasattr(self, "config_files"), "You must call .prefetch_config_files() first"
-        return QuerySetMap(self.config_files.all(), attribute="path")
+        return QuerySetMap(self.config_files, attribute="path")
 
     class Meta:
         proxy = True
 
-    @cached_property
-    def config_file_count(self) -> int:
-        assert hasattr(self, "config_files"), "You must call .prefetch_config_files() first"
-        return self.config_files.all().count()
-
     @property
     def is_default(self):
         return self.cf.get("device_config_default", False)

validity/models/device.py

@@ -24,7 +24,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
     @property
     def config_path(self) -> str:
-        assert hasattr(self, "datasource"), "You must prefetch datasource first"
+        assert hasattr(self, "data_source"), "You must prefetch datasource first"
         template = Environment(loader=BaseLoader()).from_string(self.data_source.config_path_template)
         return template.render(device=self)
 

validity/models/nameset.py

@@ -4,7 +4,6 @@
 from django.db import models
 from django.utils.translation import gettext_lazy as _
 
-from validity.managers import NameSetQS
 from .base import BaseModel, DataSourceMixin
 from .test import ComplianceTest
 
@@ -18,8 +17,6 @@ class NameSet(DataSourceMixin, BaseModel):
     )
     definitions = models.TextField(help_text=_("Here you can write python functions or imports"), blank=True)
 
-    objects = NameSetQS.as_manager()
-
     clone_fields = ("description", "_global", "tests", "definitions", "data_source", "data_file")
     text_db_field_name = "definitions"
 

validity/models/serializer.py

@@ -4,7 +4,6 @@
 from django.utils.translation import gettext_lazy as _
 
 from validity.choices import ConfigExtractionChoices
-from validity.managers import ConfigSerializerQS
 from validity.netbox_changes import DEVICE_ROLE_RELATION
 from .base import BaseModel, DataSourceMixin
 
@@ -16,8 +15,6 @@ class ConfigSerializer(DataSourceMixin, BaseModel):
     )
     ttp_template = models.TextField(_("TTP Template"), blank=True)
 
-    objects = ConfigSerializerQS.as_manager()
-
     clone_fields = ("ttp_template", "extraction_method", "data_source", "data_file")
     text_db_field_name = "ttp_template"
 

validity/templatetags/validity.py

@@ -1,6 +1,5 @@
 from typing import Any
 
-from dcim.models import Device
 from django import template
 from django.db.models import Model
 from django.utils.safestring import mark_safe

validity/utils/orm.py

@@ -42,3 +42,7 @@ def __contains__(self, key):
     def get(self, key, default=None):
         self._evaluate()
         return self._map.get(key, default)
+
+    @property
+    def model(self):
+        return self._qs.model