|
| 1 | +/* |
| 2 | + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one |
| 3 | + * or more contributor license agreements. Licensed under the "Elastic License |
| 4 | + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side |
| 5 | + * Public License v 1"; you may not use this file except in compliance with, at |
| 6 | + * your election, the "Elastic License 2.0", the "GNU Affero General Public |
| 7 | + * License v3.0 only", or the "Server Side Public License, v 1". |
| 8 | + */ |
| 9 | + |
| 10 | +package org.elasticsearch.action.search; |
| 11 | + |
| 12 | +import org.elasticsearch.common.Strings; |
| 13 | +import org.elasticsearch.index.query.BoolQueryBuilder; |
| 14 | +import org.elasticsearch.index.query.BoostingQueryBuilder; |
| 15 | +import org.elasticsearch.index.query.ConstantScoreQueryBuilder; |
| 16 | +import org.elasticsearch.index.query.NestedQueryBuilder; |
| 17 | +import org.elasticsearch.index.query.QueryBuilder; |
| 18 | +import org.elasticsearch.index.query.RangeQueryBuilder; |
| 19 | +import org.elasticsearch.search.SearchService; |
| 20 | +import org.elasticsearch.search.builder.SearchSourceBuilder; |
| 21 | +import org.elasticsearch.search.sort.FieldSortBuilder; |
| 22 | +import org.elasticsearch.search.sort.ScoreSortBuilder; |
| 23 | +import org.elasticsearch.search.sort.SortBuilder; |
| 24 | +import org.elasticsearch.search.vectors.KnnVectorQueryBuilder; |
| 25 | + |
| 26 | +import java.util.ArrayList; |
| 27 | +import java.util.Arrays; |
| 28 | +import java.util.HashMap; |
| 29 | +import java.util.List; |
| 30 | +import java.util.Map; |
| 31 | + |
| 32 | +/** |
| 33 | + * Used to introspect a search request and extract metadata from it around the features it uses. |
| 34 | + */ |
| 35 | +public class SearchRequestIntrospector { |
| 36 | + |
| 37 | + /** |
| 38 | + * Introspects the provided search request and extracts metadata from it about some of its characteristics. |
| 39 | + * |
| 40 | + */ |
| 41 | + public static QueryMetadata introspectSearchRequest(SearchRequest searchRequest) { |
| 42 | + |
| 43 | + String target = introspectIndices(searchRequest.indices()); |
| 44 | + |
| 45 | + String pitOrScroll = null; |
| 46 | + if (searchRequest.scroll() != null) { |
| 47 | + pitOrScroll = SCROLL; |
| 48 | + } |
| 49 | + |
| 50 | + SearchSourceBuilder searchSourceBuilder = searchRequest.source(); |
| 51 | + if (searchSourceBuilder == null) { |
| 52 | + return new QueryMetadata(target, ScoreSortBuilder.NAME, HITS_ONLY, false, Strings.EMPTY_ARRAY, pitOrScroll); |
| 53 | + } |
| 54 | + |
| 55 | + if (searchSourceBuilder.pointInTimeBuilder() != null) { |
| 56 | + pitOrScroll = PIT; |
| 57 | + } |
| 58 | + |
| 59 | + final String primarySort; |
| 60 | + if (searchSourceBuilder.sorts() == null || searchSourceBuilder.sorts().isEmpty()) { |
| 61 | + primarySort = ScoreSortBuilder.NAME; |
| 62 | + } else { |
| 63 | + primarySort = introspectPrimarySort(searchSourceBuilder.sorts().getFirst()); |
| 64 | + } |
| 65 | + |
| 66 | + final String queryType = introspectQueryType(searchSourceBuilder); |
| 67 | + |
| 68 | + QueryMetadataBuilder queryMetadataBuilder = new QueryMetadataBuilder(); |
| 69 | + if (searchSourceBuilder.query() != null) { |
| 70 | + introspectQueryBuilder(searchSourceBuilder.query(), queryMetadataBuilder); |
| 71 | + } |
| 72 | + |
| 73 | + final boolean hasKnn = searchSourceBuilder.knnSearch().isEmpty() == false || queryMetadataBuilder.knnQuery; |
| 74 | + |
| 75 | + return new QueryMetadata( |
| 76 | + target, |
| 77 | + primarySort, |
| 78 | + queryType, |
| 79 | + hasKnn, |
| 80 | + queryMetadataBuilder.rangeFields.toArray(new String[0]), |
| 81 | + pitOrScroll |
| 82 | + ); |
| 83 | + } |
| 84 | + |
| 85 | + private static final class QueryMetadataBuilder { |
| 86 | + private boolean knnQuery = false; |
| 87 | + private final List<String> rangeFields = new ArrayList<>(); |
| 88 | + } |
| 89 | + |
| 90 | + public record QueryMetadata( |
| 91 | + String target, |
| 92 | + String primarySort, |
| 93 | + String queryType, |
| 94 | + boolean knn, |
| 95 | + String[] rangeFields, |
| 96 | + String pitOrScroll |
| 97 | + ) { |
| 98 | + |
| 99 | + public Map<String, Object> toAttributesMap() { |
| 100 | + Map<String, Object> attributes = new HashMap<>(); |
| 101 | + attributes.put(TARGET_ATTRIBUTE, target); |
| 102 | + attributes.put(SORT_ATTRIBUTE, primarySort); |
| 103 | + if (pitOrScroll == null) { |
| 104 | + attributes.put(QUERY_TYPE_ATTRIBUTE, queryType); |
| 105 | + } else { |
| 106 | + attributes.put(QUERY_TYPE_ATTRIBUTE, Arrays.asList(queryType, pitOrScroll)); |
| 107 | + } |
| 108 | + |
| 109 | + attributes.put(KNN_ATTRIBUTE, knn); |
| 110 | + attributes.put(RANGES_ATTRIBUTE, rangeFields); |
| 111 | + return attributes; |
| 112 | + } |
| 113 | + } |
| 114 | + |
| 115 | + private static final String TARGET_ATTRIBUTE = "target"; |
| 116 | + private static final String SORT_ATTRIBUTE = "sort"; |
| 117 | + private static final String QUERY_TYPE_ATTRIBUTE = "query_type"; |
| 118 | + private static final String KNN_ATTRIBUTE = "knn"; |
| 119 | + private static final String RANGES_ATTRIBUTE = "ranges"; |
| 120 | + |
| 121 | + private static final String TARGET_KIBANA = ".kibana"; |
| 122 | + private static final String TARGET_ML = ".ml"; |
| 123 | + private static final String TARGET_FLEET = ".fleet"; |
| 124 | + private static final String TARGET_SLO = ".slo"; |
| 125 | + private static final String TARGET_ALERTS = ".alerts"; |
| 126 | + private static final String TARGET_ELASTIC = ".elastic"; |
| 127 | + private static final String TARGET_DS = ".ds-"; |
| 128 | + private static final String TARGET_OTHERS = ".others"; |
| 129 | + private static final String TARGET_USER = "user"; |
| 130 | + |
| 131 | + static String introspectIndices(String[] indices) { |
| 132 | + // Note that indices are expected to be resolved, meaning wildcards are not handled on purpose |
| 133 | + // If indices resolve to data streams, the name of the data stream is returned as opposed to its backing indices |
| 134 | + if (indices.length == 1) { |
| 135 | + String index = indices[0]; |
| 136 | + if (index.startsWith(".")) { |
| 137 | + if (index.startsWith(TARGET_KIBANA)) { |
| 138 | + return TARGET_KIBANA; |
| 139 | + } |
| 140 | + if (index.startsWith(TARGET_ML)) { |
| 141 | + return TARGET_ML; |
| 142 | + } |
| 143 | + if (index.startsWith(TARGET_FLEET)) { |
| 144 | + return TARGET_FLEET; |
| 145 | + } |
| 146 | + if (index.startsWith(TARGET_SLO)) { |
| 147 | + return TARGET_SLO; |
| 148 | + } |
| 149 | + if (index.startsWith(TARGET_ALERTS)) { |
| 150 | + return TARGET_ALERTS; |
| 151 | + } |
| 152 | + if (index.startsWith(TARGET_ELASTIC)) { |
| 153 | + return TARGET_ELASTIC; |
| 154 | + } |
| 155 | + // this happens only when data streams backing indices are searched explicitly |
| 156 | + if (index.startsWith(TARGET_DS)) { |
| 157 | + return TARGET_DS; |
| 158 | + } |
| 159 | + return TARGET_OTHERS; |
| 160 | + } |
| 161 | + } |
| 162 | + return TARGET_USER; |
| 163 | + } |
| 164 | + |
| 165 | + private static final String TIMESTAMP = "@timestamp"; |
| 166 | + private static final String EVENT_INGESTED = "event.ingested"; |
| 167 | + private static final String _DOC = "_doc"; |
| 168 | + private static final String FIELD = "field"; |
| 169 | + |
| 170 | + static String introspectPrimarySort(SortBuilder<?> primarySortBuilder) { |
| 171 | + if (primarySortBuilder instanceof FieldSortBuilder fieldSort) { |
| 172 | + return switch (fieldSort.getFieldName()) { |
| 173 | + case TIMESTAMP -> TIMESTAMP; |
| 174 | + case EVENT_INGESTED -> EVENT_INGESTED; |
| 175 | + case _DOC -> _DOC; |
| 176 | + default -> FIELD; |
| 177 | + }; |
| 178 | + } |
| 179 | + return primarySortBuilder.getWriteableName(); |
| 180 | + } |
| 181 | + |
| 182 | + private static final String HITS_AND_AGGS = "hits_and_aggs"; |
| 183 | + private static final String HITS_ONLY = "hits_only"; |
| 184 | + private static final String AGGS_ONLY = "aggs_only"; |
| 185 | + private static final String COUNT_ONLY = "count_only"; |
| 186 | + private static final String PIT = "pit"; |
| 187 | + private static final String SCROLL = "scroll"; |
| 188 | + |
| 189 | + public static final Map<String, Object> SEARCH_SCROLL_ATTRIBUTES = Map.of(QUERY_TYPE_ATTRIBUTE, SCROLL); |
| 190 | + |
| 191 | + static String introspectQueryType(SearchSourceBuilder searchSourceBuilder) { |
| 192 | + int size = searchSourceBuilder.size(); |
| 193 | + if (size == -1) { |
| 194 | + size = SearchService.DEFAULT_SIZE; |
| 195 | + } |
| 196 | + if (searchSourceBuilder.aggregations() != null && size > 0) { |
| 197 | + return HITS_AND_AGGS; |
| 198 | + } |
| 199 | + if (searchSourceBuilder.aggregations() != null) { |
| 200 | + return AGGS_ONLY; |
| 201 | + } |
| 202 | + if (size > 0) { |
| 203 | + return HITS_ONLY; |
| 204 | + } |
| 205 | + return COUNT_ONLY; |
| 206 | + } |
| 207 | + |
| 208 | + static void introspectQueryBuilder(QueryBuilder queryBuilder, QueryMetadataBuilder queryMetadataBuilder) { |
| 209 | + switch (queryBuilder) { |
| 210 | + case BoolQueryBuilder bool: |
| 211 | + for (QueryBuilder must : bool.must()) { |
| 212 | + introspectQueryBuilder(must, queryMetadataBuilder); |
| 213 | + } |
| 214 | + for (QueryBuilder filter : bool.filter()) { |
| 215 | + introspectQueryBuilder(filter, queryMetadataBuilder); |
| 216 | + } |
| 217 | + if (bool.must().isEmpty() && bool.filter().isEmpty() && bool.mustNot().isEmpty() && bool.should().size() == 1) { |
| 218 | + introspectQueryBuilder(bool.should().getFirst(), queryMetadataBuilder); |
| 219 | + } |
| 220 | + // Note that should clauses are ignored unless there's only one that becomes mandatory |
| 221 | + // must_not clauses are also ignored for now |
| 222 | + break; |
| 223 | + case ConstantScoreQueryBuilder constantScore: |
| 224 | + introspectQueryBuilder(constantScore.innerQuery(), queryMetadataBuilder); |
| 225 | + break; |
| 226 | + case BoostingQueryBuilder boosting: |
| 227 | + introspectQueryBuilder(boosting.positiveQuery(), queryMetadataBuilder); |
| 228 | + break; |
| 229 | + case NestedQueryBuilder nested: |
| 230 | + introspectQueryBuilder(nested.query(), queryMetadataBuilder); |
| 231 | + break; |
| 232 | + case RangeQueryBuilder range: |
| 233 | + switch (range.fieldName()) { |
| 234 | + case TIMESTAMP -> queryMetadataBuilder.rangeFields.add(TIMESTAMP); |
| 235 | + case EVENT_INGESTED -> queryMetadataBuilder.rangeFields.add(EVENT_INGESTED); |
| 236 | + default -> queryMetadataBuilder.rangeFields.add(FIELD); |
| 237 | + } |
| 238 | + break; |
| 239 | + case KnnVectorQueryBuilder knn: |
| 240 | + queryMetadataBuilder.knnQuery = true; |
| 241 | + break; |
| 242 | + default: |
| 243 | + } |
| 244 | + } |
| 245 | +} |
0 commit comments