perf:Stored command records in ES support accurate searching.

This commit is contained in:
gerry 2025-04-11 15:42:55 +08:00 committed by feng626
parent ee97e45cc3
commit 3d6d2af268
2 changed files with 52 additions and 16 deletions

View File

@ -14,6 +14,7 @@ from uuid import UUID
from django.utils.translation import gettext_lazy as _ from django.utils.translation import gettext_lazy as _
from django.db.models import QuerySet as DJQuerySet from django.db.models import QuerySet as DJQuerySet
from django.db.models import Q
from elasticsearch7 import Elasticsearch from elasticsearch7 import Elasticsearch
from elasticsearch7.helpers import bulk from elasticsearch7.helpers import bulk
from elasticsearch7.exceptions import RequestError, SSLError from elasticsearch7.exceptions import RequestError, SSLError
@ -78,6 +79,10 @@ class ESClientV7(ESClientBase):
def get_sort(cls, field, direction): def get_sort(cls, field, direction):
return f'{field}:{direction}' return f'{field}:{direction}'
@classmethod
def get_sorts(cls, sorts: list):
return ','.join(sorts)
class ESClientV6(ESClientV7): class ESClientV6(ESClientV7):
@ -99,6 +104,10 @@ class ESClientV8(ESClientBase):
def get_sort(cls, field, direction): def get_sort(cls, field, direction):
return {field: {'order': direction}} return {field: {'order': direction}}
@classmethod
def get_sorts(cls, sorts: list):
return sorts
def get_es_client_version(**kwargs): def get_es_client_version(**kwargs):
try: try:
@ -190,8 +199,7 @@ class ES(object):
mappings['aliases'] = { mappings['aliases'] = {
self.query_index: {} self.query_index: {}
} }
if self.es.indices.exists(index=self.index):
return
try: try:
self.es.indices.create(index=self.index, body=mappings) self.es.indices.create(index=self.index, body=mappings)
except (RequestError, BadRequestError) as e: except (RequestError, BadRequestError) as e:
@ -245,6 +253,7 @@ class ES(object):
} }
if sort is not None: if sort is not None:
search_params['sort'] = sort search_params['sort'] = sort
logger.info('search_params: {}'.format(search_params))
data = self.es.search(**search_params) data = self.es.search(**search_params)
source_data = [] source_data = []
@ -319,10 +328,12 @@ class ES(object):
kwargs = new_kwargs kwargs = new_kwargs
index_in_field = 'id__in' index_in_field = 'id__in'
keyword_fields = self.keyword_fields
exact_fields = self.exact_fields exact_fields = self.exact_fields
match_fields = self.match_fields match_fields = self.match_fields
match = {} match = {}
search = []
exact = {} exact = {}
index = {} index = {}
@ -330,11 +341,19 @@ class ES(object):
index['values'] = kwargs[index_in_field] index['values'] = kwargs[index_in_field]
for k, v in kwargs.items(): for k, v in kwargs.items():
if k in exact_fields: if k in keyword_fields:
exact[k] = v exact[k] = v
elif k in exact_fields:
exact['{}.keyword'.format(k)] = v
elif k in match_fields: elif k in match_fields:
match[k] = v match[k] = v
args = kwargs.get('search')
for item in args:
for k, v in item.items():
if k in match_fields:
search.append(item)
# 处理时间 # 处理时间
time_field_name, time_range = self.handler_time_field(kwargs) time_field_name, time_range = self.handler_time_field(kwargs)
@ -363,10 +382,12 @@ class ES(object):
body = { body = {
'query': { 'query': {
'bool': { 'bool': {
'must': [ 'must': [],
'should': should + [
{'match': {k: v}} for k, v in match.items() {'match': {k: v}} for k, v in match.items()
] + [
{'match': item} for item in search
], ],
'should': should,
'filter': self.handle_exact_fields(exact) + 'filter': self.handle_exact_fields(exact) +
[ [
{ {
@ -403,6 +424,17 @@ class QuerySet(DJQuerySet):
_method_calls = {k: list(v) for k, v in groupby(self._method_calls, lambda x: x[0])} _method_calls = {k: list(v) for k, v in groupby(self._method_calls, lambda x: x[0])}
return _method_calls return _method_calls
def _grouped_search_args(self, query):
conditions = {}
for q in query:
for c in q.children:
if isinstance(c, Q):
child = self._grouped_search_args(c)
[conditions.setdefault(k, []).extend(v) for k, v in child.items()]
else:
conditions.setdefault(c[0], []).append(c[1])
return conditions
@lazyproperty @lazyproperty
def _filter_kwargs(self): def _filter_kwargs(self):
_method_calls = self._grouped_method_calls _method_calls = self._grouped_method_calls
@ -410,14 +442,14 @@ class QuerySet(DJQuerySet):
if not filter_calls: if not filter_calls:
return {} return {}
names, multi_args, multi_kwargs = zip(*filter_calls) names, multi_args, multi_kwargs = zip(*filter_calls)
args = {
key: value # input 输入
for arg in multi_args if arg multi_args = tuple(reduce(lambda x, y: x + y, (sub for sub in multi_args if sub),()))
for key, value in arg[0].children args = self._grouped_search_args(multi_args)
} striped_args = [{k.replace('__icontains', ''): v} for k, values in args.items() for v in values]
kwargs = reduce(lambda x, y: {**x, **y}, multi_kwargs, {}) kwargs = reduce(lambda x, y: {**x, **y}, multi_kwargs, {})
kwargs.update(args) striped_kwargs = {'search': striped_args}
striped_kwargs = {}
for k, v in kwargs.items(): for k, v in kwargs.items():
k = k.replace('__exact', '') k = k.replace('__exact', '')
k = k.replace('__startswith', '') k = k.replace('__startswith', '')
@ -428,6 +460,7 @@ class QuerySet(DJQuerySet):
@lazyproperty @lazyproperty
def _sort(self): def _sort(self):
order_by = self._grouped_method_calls.get('order_by') order_by = self._grouped_method_calls.get('order_by')
_sorts = [self._storage.client.get_sort('_score', 'desc')]
if order_by: if order_by:
for call in reversed(order_by): for call in reversed(order_by):
fields = call[1] fields = call[1]
@ -440,7 +473,10 @@ class QuerySet(DJQuerySet):
direction = 'asc' direction = 'asc'
field = field.lstrip('-+') field = field.lstrip('-+')
sort = self._storage.client.get_sort(field, direction) sort = self._storage.client.get_sort(field, direction)
return sort _sorts.append(sort)
break
sorts = self._storage.client.get_sorts(_sorts)
return sorts
def __execute(self): def __execute(self):
_filter_kwargs = self._filter_kwargs _filter_kwargs = self._filter_kwargs
@ -514,4 +550,4 @@ class QuerySet(DJQuerySet):
return iter(self.__execute()) return iter(self.__execute())
def __len__(self): def __len__(self):
return self.count() return self.count()

View File

@ -27,8 +27,8 @@ class CommandStore(ES):
"type": "long" "type": "long"
} }
} }
exact_fields = {} exact_fields = {'input', 'risk_level', 'user', 'asset', 'account'}
match_fields = {'input', 'risk_level', 'user', 'asset', 'system_user'} match_fields = {'input'}
keyword_fields = {'session', 'org_id'} keyword_fields = {'session', 'org_id'}
super().__init__(config, properties, keyword_fields, exact_fields, match_fields) super().__init__(config, properties, keyword_fields, exact_fields, match_fields)