From 98c9cddcbfc827fa4c6c10607cc69044e4894d34 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Wed, 20 Dec 2023 17:45:18 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20es=E5=91=BD=E4=BB=A4=E8=AE=B0=E5=BD=95?=
 =?UTF-8?q?=E5=8F=AF=E4=BB=A5=E7=9C=8B=E5=88=B0=E5=85=B6=E4=BB=96=E8=B5=84?=
 =?UTF-8?q?=E4=BA=A7=E6=89=A7=E8=A1=8C=E7=9A=84=E5=91=BD=E4=BB=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/plugins/es.py            | 21 ++++++++++++++++-----
 apps/terminal/api/session/command.py |  4 ----
 apps/terminal/filters.py             | 13 +++++++++++--
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/apps/common/plugins/es.py b/apps/common/plugins/es.py
index 297f2e32a..8e2126f85 100644
--- a/apps/common/plugins/es.py
+++ b/apps/common/plugins/es.py
@@ -228,6 +228,18 @@ class ES(object):
             datetime_range['lte'] = datetime__lte
         return 'datetime', datetime_range
 
+    @staticmethod
+    def handle_exact_fields(exact):
+        _filter = []
+        for k, v in exact.items():
+            query = 'term'
+            if isinstance(v, list):
+                query = 'terms'
+            _filter.append({
+                query: {k: v}
+            })
+        return _filter
+
     def get_query_body(self, **kwargs):
         new_kwargs = {}
         for k, v in kwargs.items():
@@ -235,6 +247,8 @@ class ES(object):
                 v = str(v)
             if k == 'pk':
                 k = 'id'
+            if k.endswith('__in'):
+                k = k.replace('__in', '')
             new_kwargs[k] = v
         kwargs = new_kwargs
 
@@ -287,11 +301,8 @@ class ES(object):
                         {'match': {k: v}} for k, v in match.items()
                     ],
                     'should': should,
-                    'filter': [
-                                  {
-                                      'term': {k: v}
-                                  } for k, v in exact.items()
-                              ] + [
+                    'filter': self.handle_exact_fields(exact) +
+                              [
                                   {
                                       'range': {
                                           time_field_name: time_range
diff --git a/apps/terminal/api/session/command.py b/apps/terminal/api/session/command.py
index 387aebc3e..baec8f0fa 100644
--- a/apps/terminal/api/session/command.py
+++ b/apps/terminal/api/session/command.py
@@ -167,7 +167,6 @@ class CommandViewSet(JMSBulkModelViewSet):
 
     def get_queryset(self):
         command_storage_id = self.request.query_params.get('command_storage_id')
-        asset_id = self.request.query_params.get('asset_id')
         if not command_storage_id:
             return Command.objects.none()
 
@@ -176,9 +175,6 @@ class CommandViewSet(JMSBulkModelViewSet):
             raise StorageInvalid
         else:
             qs = storage.get_command_queryset()
-        if asset_id:
-            session_ids = Session.objects.filter(asset_id=asset_id).values_list('id', flat=True)
-            qs = qs.filter(session__in=list(session_ids))
         return qs
 
     def create(self, request, *args, **kwargs):
diff --git a/apps/terminal/filters.py b/apps/terminal/filters.py
index 40c048b28..0839c770c 100644
--- a/apps/terminal/filters.py
+++ b/apps/terminal/filters.py
@@ -2,7 +2,7 @@ from django.db.models import QuerySet
 from django_filters import rest_framework as filters
 
 from orgs.utils import filter_org_queryset
-from terminal.models import Command, CommandStorage
+from terminal.models import Command, CommandStorage, Session
 
 
 class CommandFilter(filters.FilterSet):
@@ -13,7 +13,7 @@ class CommandFilter(filters.FilterSet):
     user = filters.CharFilter(lookup_expr='startswith')
     input = filters.CharFilter(lookup_expr='icontains')
     asset = filters.CharFilter(field_name='asset', lookup_expr='icontains')
-    asset_id = filters.UUIDFilter(method='do_nothing')
+    asset_id = filters.UUIDFilter(method='filter_by_asset_id')
 
     class Meta:
         model = Command
@@ -48,6 +48,15 @@ class CommandFilter(filters.FilterSet):
         qs = qs.filter(**filters)
         return qs
 
+    def filter_by_asset_id(self, queryset, name, value):
+        asset_id = self.form.cleaned_data.get('asset_id')
+        filters = {}
+        if asset_id:
+            session_ids = Session.objects.filter(asset_id=asset_id).values_list('id', flat=True)
+            filters['session__in'] = list(session_ids)
+        queryset = queryset.filter(**filters)
+        return queryset
+
 
 class CommandFilterForStorageTree(CommandFilter):
     asset = filters.CharFilter(method='do_nothing')