perf: 优化命令记录慢的问题

apps/terminal/api/command.py

@@ -8,11 +8,14 @@ from rest_framework import viewsets
 from rest_framework import generics
 from rest_framework.fields import DateTimeField
 from rest_framework.response import Response
-from rest_framework import status
+from rest_framework.decorators import action
 from django.template import loader
 
+from terminal.models import CommandStorage
+from terminal.filters import CommandFilter
 from orgs.utils import current_org
 from common.permissions import IsOrgAdminOrAppUser, IsOrgAuditor, IsAppUser
+from common.const.http import GET
 from common.utils import get_logger
 from terminal.utils import send_command_alert_mail
 from terminal.serializers import InsecureCommandAlertSerializer
@@ -89,7 +92,7 @@ class CommandQueryMixin:
         return date_from_st, date_to_st
 
 
-class CommandViewSet(CommandQueryMixin, viewsets.ModelViewSet):
+class CommandViewSet(viewsets.ModelViewSet):
     """接受app发送来的command log, 格式如下
     {
         "user": "admin",
@@ -103,7 +106,16 @@ class CommandViewSet(CommandQueryMixin, viewsets.ModelViewSet):
 
     """
     command_store = get_command_storage()
+    permission_classes = [IsOrgAdminOrAppUser | IsOrgAuditor]
     serializer_class = SessionCommandSerializer
+    filterset_class = CommandFilter
+    ordering_fields = ('timestamp', )
+
+    def get_queryset(self):
+        command_storage_id = self.request.query_params.get('command_storage_id')
+        storage = CommandStorage.objects.get(id=command_storage_id)
+        qs = storage.get_command_queryset()
+        return qs
 
     def create(self, request, *args, **kwargs):
         serializer = self.serializer_class(data=request.data, many=True)

apps/terminal/api/storage.py

@@ -3,9 +3,15 @@
 
 from rest_framework import viewsets, generics, status
 from rest_framework.response import Response
+from rest_framework.request import Request
+from rest_framework.decorators import action
 from django.utils.translation import ugettext_lazy as _
+from django_filters import utils
 
+from terminal import const
+from common.const.http import GET
 from common.permissions import IsSuperUser
+from terminal.filters import CommandStorageFilter, CommandFilter, CommandFilterForStorageTree
 from ..models import CommandStorage, ReplayStorage
 from ..serializers import CommandStorageSerializer, ReplayStorageSerializer
 
@@ -30,11 +36,52 @@ class BaseStorageViewSetMixin:
 
 
 class CommandStorageViewSet(BaseStorageViewSetMixin, viewsets.ModelViewSet):
-    filterset_fields = ('name', 'type',)
-    search_fields = filterset_fields
+    search_fields = ('name', 'type',)
     queryset = CommandStorage.objects.all()
     serializer_class = CommandStorageSerializer
     permission_classes = (IsSuperUser,)
+    filterset_class = CommandStorageFilter
+
+    @action(methods=[GET], detail=False, filterset_class=CommandFilterForStorageTree)
+    def tree(self, request: Request):
+        storage_qs = self.get_queryset().exclude(name='null')
+        storages_with_count = []
+        for storage in storage_qs:
+            command_qs = storage.get_command_queryset()
+            filterset = CommandFilter(
+                data=request.query_params, queryset=command_qs,
+                request=request
+            )
+            if not filterset.is_valid():
+                raise utils.translate_validation(filterset.errors)
+            command_qs = filterset.qs
+            if storage.type == const.CommandStorageTypeChoices.es:
+                command_count = command_qs.count(limit_to_max_result_window=False)
+            else:
+                command_count = command_qs.count()
+            storages_with_count.append((storage, command_count))
+
+        root = {
+            'id': 'root',
+            'name': _('Command storages'),
+            'title': _('Command storages'),
+            'pId': '',
+            'isParent': True,
+            'open': True,
+        }
+
+        nodes = [
+            {
+                'id': storage.id,
+                'name': f'{storage.name}({storage.type})({command_count})',
+                'title': f'{storage.name}({storage.type})',
+                'pId': 'root',
+                'isParent': False,
+                'open': False,
+            } for storage, command_count in storages_with_count
+        ]
+        nodes.append(root)
+        return Response(data=nodes)
 
 
 class ReplayStorageViewSet(BaseStorageViewSetMixin, viewsets.ModelViewSet):