Performance/libraries

backend/core/helpers.py

@@ -188,98 +188,73 @@ def _get_all_requirement_nodes_id_in_requirement_node(
 
 
 def get_sorted_requirement_nodes(
-    requirement_nodes: list,
-    requirements_assessed: list | None,
+    requirement_nodes: list, requirements_assessed: list | None
 ) -> dict:
     """
-    Recursive function to build framework groups tree
-    requirement_nodes: the list of all requirement_nodes
-    requirements_assessed: the list of all requirements_assessed
-    Returns a dictionary containing key=name and value={"description": description, "style": "leaf|node"}}
+    Optimized function to build a framework groups tree.
     """
-    requirement_assessment_from_requirement_id = (
+    # Preprocess requirements_assessed for O(1) lookup.
+    ra_dict = (
         {str(ra.requirement.id): ra for ra in requirements_assessed}
         if requirements_assessed
         else {}
     )
 
-    def get_sorted_requirement_nodes_rec(
-        requirement_nodes: list,
-        requirements_assessed: list,
-        start: list,
-    ) -> dict:
+    # Preprocess requirement_nodes into a dict by parent_urn for O(1) child lookup.
+    children_dict = {}
+    for node in requirement_nodes:
+        if node.parent_urn not in children_dict:
+            children_dict[node.parent_urn] = []
+        children_dict[node.parent_urn].append(node)
+
+    def get_sorted_requirement_nodes_rec(parent_urn):
         """
-        Recursive function to build framework groups tree, within get_sorted_requirements_and_groups
-        start: the initial list
+        Recursive helper function to build the framework groups tree.
         """
         result = {}
-        for node in start:
-            children = [
-                requirement_node
-                for requirement_node in requirement_nodes
-                if requirement_node.parent_urn == node.urn
-            ]
-            result[str(node.id)] = {
+        for node in children_dict.get(parent_urn, []):
+            node_info = {
                 "urn": node.urn,
                 "parent_urn": node.parent_urn,
                 "name": node.display_short(),
                 "node_content": node.display_long(),
-                "style": "node",
+                "style": "node" if node.urn in children_dict else "leaf",
                 "assessable": node.assessable,
                 "description": node.description,
-                "children": get_sorted_requirement_nodes_rec(
-                    requirement_nodes, requirements_assessed, children
-                ),
+                "children": get_sorted_requirement_nodes_rec(node.urn),
             }
-            for req in [
-                requirement_node
-                for requirement_node in requirement_nodes
-                if requirement_node.parent_urn == node.urn
-            ]:
-                if requirements_assessed:
-                    req_as = requirement_assessment_from_requirement_id[str(req.id)]
-                    result[str(node.id)]["children"][str(req.id)].update(
-                        {
-                            "urn": req.urn,
-                            "name": req.display_short(),
-                            "description": req.description,
-                            "ra_id": str(req_as.id),
-                            "leaf_content": req.display_long(),
-                            "status": req_as.status,
-                            "status_display": req_as.get_status_display(),
-                            "status_i18n": camel_case(req_as.status),
-                            "style": "leaf",
-                            "threats": ThreatReadSerializer(
-                                req.threats.all(), many=True
-                            ).data,
-                            "security_functions": SecurityFunctionReadSerializer(
-                                req.security_functions.all(), many=True
-                            ).data,
-                        }
-                    )
-                else:
-                    result[str(node.id)]["children"][str(req.id)].update(
+
+            # Add requirement assessment info if available
+            if ra_dict:
+                ra = ra_dict.get(str(node.id))
+                if ra:
+                    node_info.update(
                         {
-                            "urn": req.urn,
-                            "name": req.display_short(),
-                            "leaf_content": req.display_long(),
-                            "description": req.description,
-                            "style": "leaf",
+                            "ra_id": str(ra.id),
+                            "leaf_content": node_info.get("node_content", ""),
+                            "status": ra.status,
+                            "status_display": ra.get_status_display(),
+                            "status_i18n": camel_case(ra.status),
                             "threats": ThreatReadSerializer(
-                                req.threats.all(), many=True
+                                ra.requirement.threats.all(), many=True
                             ).data,
                             "security_functions": SecurityFunctionReadSerializer(
-                                req.security_functions.all(), many=True
+                                ra.requirement.security_functions.all(), many=True
                             ).data,
                         }
                     )
+                    node_info[
+                        "style"
+                    ] = "leaf"  # Update style to leaf if it has an assessment
+
+            result[str(node.id)] = node_info
+
         return result
 
+    # Initialize the recursive building from root nodes (those without a parent_urn).
     tree = get_sorted_requirement_nodes_rec(
-        requirement_nodes,
-        requirements_assessed,
-        [rg for rg in requirement_nodes if not rg.parent_urn],
-    )
+        None
+    )  # Assuming root nodes have `parent_urn` as None or similar.
 
     return tree
 

backend/core/models.py

@@ -1,4 +1,5 @@
 from django.apps import apps
+from django.forms.models import model_to_dict
 from django.contrib.auth import get_user_model
 from django.db import models
 from django.utils.translation import gettext_lazy as _
@@ -100,6 +101,23 @@ class Library(ReferentialObjectMixin):
         "self", blank=True, verbose_name=_("Dependencies"), symmetrical=False
     )
 
+    @property
+    def _objects(self):
+        res = {}
+        if self.frameworks.count() > 0:
+            res["framework"] = model_to_dict(self.frameworks.first())
+            res["framework"].update(self.frameworks.first().library_entry)
+        if self.threats.count() > 0:
+            res["threats"] = [model_to_dict(threat) for threat in self.threats.all()]
+        if self.security_functions.count() > 0:
+            res["security_functions"] = [
+                model_to_dict(security_function)
+                for security_function in self.security_functions.all()
+            ]
+        if self.risk_matrices.count() > 0:
+            res["risk_matrix"] = model_to_dict(self.risk_matrices.first())
+        return res
+
     @property
     def reference_count(self) -> int:
         """
@@ -304,6 +322,49 @@ def is_deletable(self) -> bool:
             return False
         return True
 
+    @property
+    def library_entry(self):
+        res = {}
+        requirement_nodes = self.get_requirement_nodes()
+        if requirement_nodes:
+            res["requirement_nodes"] = requirement_nodes
+
+        requirement_levels = self.get_requirement_levels()
+        if requirement_levels:
+            res["requirement_levels"] = requirement_levels
+
+        return res
+
+    def get_requirement_nodes(self):
+        # Prefetch related objects if they exist to reduce database queries.
+        # Adjust prefetch_related paths according to your model relationships.
+        nodes_queryset = self.requirement_nodes.prefetch_related(
+            "threats", "security_functions"
+        )
+        if nodes_queryset.exists():
+            return [self.process_node(node) for node in nodes_queryset]
+        return []
+
+    def process_node(self, node):
+        # Convert the node to dict and process threats and security functions.
+        node_dict = model_to_dict(node)
+        if node.threats.exists():
+            node_dict["threats"] = [
+                model_to_dict(threat) for threat in node.threats.all()
+            ]
+        if node.security_functions.exists():
+            node_dict["security_functions"] = [
+                model_to_dict(security_function)
+                for security_function in node.security_functions.all()
+            ]
+        return node_dict
+
+    def get_requirement_levels(self):
+        levels_queryset = self.requirement_levels.all()
+        if levels_queryset.exists():
+            return [model_to_dict(level) for level in levels_queryset]
+        return []
+
 
 class RequirementLevel(ReferentialObjectMixin):
     framework = models.ForeignKey(
@@ -312,6 +373,7 @@ class RequirementLevel(ReferentialObjectMixin):
         null=True,
         blank=True,
         verbose_name=_("Framework"),
+        related_name="requirement_levels",
     )
     level = models.IntegerField(null=False, blank=False, verbose_name=_("Level"))
 
@@ -339,6 +401,7 @@ class RequirementNode(ReferentialObjectMixin):
         null=True,
         blank=True,
         verbose_name=_("Framework"),
+        related_name="requirement_nodes",
     )
     parent_urn = models.CharField(
         max_length=100, null=True, blank=True, verbose_name=_("Parent URN")

backend/library/serializers.py

@@ -1,16 +1,20 @@
 from rest_framework import serializers
+from core.models import Library
+
+from core.serializers import (
+    BaseModelSerializer,
+    FrameworkReadSerializer,
+    RiskMatrixReadSerializer,
+    SecurityFunctionReadSerializer,
+    ThreatReadSerializer,
+)
 
 
 class LibraryObjectSerializer(serializers.Serializer):
-    type = serializers.ChoiceField(
-        choices=[
-            "risk_matrix",
-            "security_function",
-            "threat",
-            "framework",
-        ]
-    )
-    fields = serializers.DictField(child=serializers.CharField())
+    framework = FrameworkReadSerializer()
+    risk_matrix = RiskMatrixReadSerializer()
+    threats = ThreatReadSerializer(many=True)
+    security_functions = SecurityFunctionReadSerializer(many=True)
 
 
 class LibrarySerializer(serializers.Serializer):
@@ -22,6 +26,12 @@ class LibrarySerializer(serializers.Serializer):
     copyright = serializers.CharField()
 
 
+class LibraryModelSerializer(BaseModelSerializer):
+    class Meta:
+        model = Library
+        fields = "__all__"
+
+
 class LibraryUploadSerializer(serializers.Serializer):
     file = serializers.FileField(required=True)