Merge pull request #436 from Bidaya0/optimise/vul_summary_query_union

union vul_summary count as one query to reduce IO time

iast/base/agent.py

@@ -7,6 +7,7 @@
 import json
 import re
 import time
+from django.db.models import Count
 
 from dongtai.models.agent import IastAgent
 from dongtai.models.project import IastProject
@@ -17,7 +18,7 @@
 from iast.base.project_version import get_project_version
 from dongtai.models.strategy import IastStrategyModel
 from dongtai.models.project_version import IastProjectVersion
-
+from functools import reduce
 
 def get_agents_with_project(project_name, users):
     """
@@ -88,23 +89,35 @@ def get_project_vul_count(users, queryset, auth_agents, project_id=None):
         current_version=1,
         user__in=users).values_list('id', 'project_id').all()
     versions_map = {version[1]: version[0] for version in versions}
+    qss = []
+    project_pair = {}
     for project in project_queryset:
         project_id = project['id']
         version_id = versions_map.get(project_id, 0)
         agent_queryset = auth_agents.filter(project_version_id=version_id,
                                             bind_project_id=project_id)
-        count = queryset.filter(agent__in=agent_queryset).values('id').count()
-        if count is False:
-            result.append({
-                "project_name": project['name'],
-                "count": 0,
-                "id": project_id
-            })
-        else:
+        qs = queryset.filter(agent__in=agent_queryset).extra({
+            "project_id":
+            project_id,
+        }).annotate(count=Count('id')).values('project_id',
+                                              'count')
+        project_pair[project_id] = project['name']
+        qss.append(qs)
+    if len(qss) > 1:
+        unionqs = qss[0].union(*qss[1::])
+    elif len(qss) == 1:
+        unionqs = qss[0]
+    else:
+        unionqs = None
+    if unionqs:
+        for project_result in unionqs:
             result.append({
-                "project_name": project['name'],
-                "count": count,
-                "id": project_id
+                "project_name":
+                project_pair[project_result['project_id']],
+                "count":
+                project_result['count'],
+                "project_id":
+                project_result['project_id'],
             })
     result = sorted(result, key=lambda item: item['count'], reverse=True)[:5]
     return result

iast/views/vul_summary.py

@@ -298,7 +298,7 @@ def get(self, request):
         vul_type = request.query_params.get('type')
         if vul_type:
             hook_types = HookType.objects.filter(name=vul_type).all()
-            strategys = IastStrategyModel.objects.filter(vul_name=vul_type).all() 
+            strategys = IastStrategyModel.objects.filter(vul_name=vul_type).all()
             q = Q(hook_type__in=hook_types,strategy_id=0) | Q(strategy__in=strategys)
             queryset = queryset.filter(q)
 
@@ -310,8 +310,10 @@ def get(self, request):
         queryset = queryset.filter(q)
 
         level_summary = queryset.values('level').order_by('level').annotate(total=Count('level'))
-        type_summary = queryset.values('hook_type_id','strategy_id').order_by(
-            'hook_type_id').annotate(total=Count('hook_type_id'))
+        type_summary = queryset.values(
+            'hook_type_id', 'strategy_id', 'hook_type__name',
+            'strategy__vul_name').order_by('hook_type_id').annotate(
+                total=Count('hook_type_id'))
 
         end['data']['language'] = self.get_languages(queryset.values('agent_id'))
 
@@ -350,10 +352,15 @@ def get_level_name(id):
 
 
 def get_hook_type_name(obj):
-    hook_type = HookType.objects.filter(pk=obj['hook_type_id']).first()
-    hook_type_name = hook_type.name if hook_type else None
-    strategy = IastStrategyModel.objects.filter(pk=obj['strategy_id']).first()
-    strategy_name = strategy.vul_name if strategy else None
+    #hook_type = HookType.objects.filter(pk=obj['hook_type_id']).first()
+    #hook_type_name = hook_type.name if hook_type else None
+    #strategy = IastStrategyModel.objects.filter(pk=obj['strategy_id']).first()
+    #strategy_name = strategy.vul_name if strategy else None
+    #type_ = list(
+    #    filter(lambda x: x is not None, [strategy_name, hook_type_name]))
     type_ = list(
-        filter(lambda x: x is not None, [strategy_name, hook_type_name]))
+        filter(lambda x: x is not None, [
+            obj.get('strategy__vul_name', None),
+            obj.get('hook_type__name', None)
+        ]))
     return type_[0] if type_ else ''