From dcf52c28a6bd7ed169c71fd8c006e56db7b73f13 Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Thu, 10 Oct 2024 01:23:11 -0700 Subject: [PATCH] Add field type to query shape and extensive unit tests (#140) * Add query type to shape and extensive unit tests Signed-off-by: Siddhant Deshmukh * Log query shape if trace log enabled Signed-off-by: Siddhant Deshmukh * Refactor code, handle multifield with test, remove properties cache Signed-off-by: Siddhant Deshmukh * Use only first index and optimization to get properties once per query Signed-off-by: Siddhant Deshmukh * Further refactoring Signed-off-by: Siddhant Deshmukh --------- Signed-off-by: Siddhant Deshmukh --- .../core/listener/QueryInsightsListener.java | 24 +- .../categorizer/QueryShapeGenerator.java | 296 ++++++-- .../categorizer/QueryShapeVisitor.java | 37 +- .../categorizer/SearchQueryCategorizer.java | 5 - .../categorizer/QueryShapeGeneratorTests.java | 686 +++++++++++++++++- .../categorizer/QueryShapeVisitorTests.java | 12 +- 6 files changed, 957 insertions(+), 103 deletions(-) diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java index 2b18692..aa141f3 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java +++ b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java @@ -58,6 +58,7 @@ public final class QueryInsightsListener extends SearchRequestOperationsListener private final ClusterService clusterService; private boolean groupingFieldNameEnabled; private boolean groupingFieldTypeEnabled; + private final QueryShapeGenerator queryShapeGenerator; /** * Constructor for QueryInsightsListener @@ -87,6 +88,7 @@ public QueryInsightsListener( super(initiallyEnabled); this.clusterService = clusterService; this.queryInsightsService = queryInsightsService; + this.queryShapeGenerator = new QueryShapeGenerator(clusterService); // Setting endpoints set up for top n queries, including enabling top n queries, window size, and top n size // Expected metricTypes are Latency, CPU, and Memory. @@ -270,9 +272,25 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final attributes.put(Attribute.PHASE_LATENCY_MAP, searchRequestContext.phaseTookMap()); attributes.put(Attribute.TASK_RESOURCE_USAGES, tasksResourceUsages); - if (queryInsightsService.isGroupingEnabled()) { - String hashcode = QueryShapeGenerator.getShapeHashCodeAsString(request.source(), groupingFieldNameEnabled); - attributes.put(Attribute.QUERY_HASHCODE, hashcode); + if (queryInsightsService.isGroupingEnabled() || log.isTraceEnabled()) { + // Generate the query shape only if grouping is enabled or trace logging is enabled + final String queryShape = queryShapeGenerator.buildShape( + request.source(), + groupingFieldNameEnabled, + groupingFieldTypeEnabled, + searchRequestContext.getSuccessfulSearchShardIndices() + ); + + // Print the query shape if tracer is enabled + if (log.isTraceEnabled()) { + log.trace("Query Shape:\n{}", queryShape); + } + + // Add hashcode attribute when grouping is enabled + if (queryInsightsService.isGroupingEnabled()) { + String hashcode = queryShapeGenerator.getShapeHashCodeAsString(queryShape); + attributes.put(Attribute.QUERY_HASHCODE, hashcode); + } } Map labels = new HashMap<>(); diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java index 4aaeff8..f16257b 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java @@ -8,13 +8,20 @@ package org.opensearch.plugin.insights.core.service.categorizer; +import java.io.IOException; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import org.apache.lucene.util.BytesRef; +import org.opensearch.cluster.metadata.MappingMetadata; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.hash.MurmurHash3; import org.opensearch.core.common.io.stream.NamedWriteable; +import org.opensearch.core.index.Index; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.WithFieldName; import org.opensearch.search.aggregations.AggregationBuilder; @@ -29,61 +36,159 @@ public class QueryShapeGenerator { static final String EMPTY_STRING = ""; static final String ONE_SPACE_INDENT = " "; + private final ClusterService clusterService; + private final String NO_FIELD_TYPE_VALUE = ""; + private final ConcurrentHashMap> fieldTypeMap; + + public QueryShapeGenerator(ClusterService clusterService) { + this.clusterService = clusterService; + this.fieldTypeMap = new ConcurrentHashMap<>(); + } /** - * Method to get query shape hash code given a source - * @param source search request source - * @param showFields whether to include field data in query shape - * @return Hash code of query shape as a MurmurHash3.Hash128 object (128-bit) + * Gets the hash code of the query shape given a search source. + * + * @param source search request source + * @param showFieldName whether to include field names in the query shape + * @param showFieldType whether to include field types in the query shape + * @param successfulSearchShardIndices the set of indices that were successfully searched + * @return Hash code of the query shape as a MurmurHash3.Hash128 object (128-bit) */ - public static MurmurHash3.Hash128 getShapeHashCode(SearchSourceBuilder source, Boolean showFields) { - final String shape = buildShape(source, showFields); + public MurmurHash3.Hash128 getShapeHashCode( + SearchSourceBuilder source, + Boolean showFieldName, + Boolean showFieldType, + Set successfulSearchShardIndices + ) { + final String shape = buildShape(source, showFieldName, showFieldType, successfulSearchShardIndices); final BytesRef shapeBytes = new BytesRef(shape); return MurmurHash3.hash128(shapeBytes.bytes, 0, shapeBytes.length, 0, new MurmurHash3.Hash128()); } - public static String getShapeHashCodeAsString(SearchSourceBuilder source, Boolean showFields) { - MurmurHash3.Hash128 hashcode = getShapeHashCode(source, showFields); + /** + * Gets the hash code of the query shape as a string. + * + * @param source search request source + * @param showFieldName whether to include field names in the query shape + * @param showFieldType whether to include field types in the query shape + * @param successfulSearchShardIndices the set of indices that were successfully searched + * @return Hash code of the query shape as a string + */ + public String getShapeHashCodeAsString( + SearchSourceBuilder source, + Boolean showFieldName, + Boolean showFieldType, + Set successfulSearchShardIndices + ) { + MurmurHash3.Hash128 hashcode = getShapeHashCode(source, showFieldName, showFieldType, successfulSearchShardIndices); String hashAsString = Long.toHexString(hashcode.h1) + Long.toHexString(hashcode.h2); return hashAsString; } /** - * Method to build search query shape given a source - * @param source search request source - * @param showFields whether to append field data - * @return Search query shape as String + * Gets the hash code of the query shape as a string. + * + * @param queryShape query shape as input + * @return Hash code of the query shape as a string */ - public static String buildShape(SearchSourceBuilder source, Boolean showFields) { + public String getShapeHashCodeAsString(String queryShape) { + final BytesRef shapeBytes = new BytesRef(queryShape); + MurmurHash3.Hash128 hashcode = MurmurHash3.hash128(shapeBytes.bytes, 0, shapeBytes.length, 0, new MurmurHash3.Hash128()); + return Long.toHexString(hashcode.h1) + Long.toHexString(hashcode.h2); + } + + /** + * Builds the search query shape given a source. + * + * @param source search request source + * @param showFieldName whether to append field names + * @param showFieldType whether to append field types + * @param successfulSearchShardIndices the set of indices that were successfully searched + * @return Search query shape as a String + */ + public String buildShape( + SearchSourceBuilder source, + Boolean showFieldName, + Boolean showFieldType, + Set successfulSearchShardIndices + ) { + Index firstIndex = null; + Map propertiesAsMap = null; + if (successfulSearchShardIndices != null) { + firstIndex = successfulSearchShardIndices.iterator().next(); + propertiesAsMap = getPropertiesMapForIndex(firstIndex); + } + StringBuilder shape = new StringBuilder(); - shape.append(buildQueryShape(source.query(), showFields)); - shape.append(buildAggregationShape(source.aggregations(), showFields)); - shape.append(buildSortShape(source.sorts(), showFields)); + shape.append(buildQueryShape(source.query(), showFieldName, showFieldType, propertiesAsMap, firstIndex)); + shape.append(buildAggregationShape(source.aggregations(), showFieldName, showFieldType, propertiesAsMap, firstIndex)); + shape.append(buildSortShape(source.sorts(), showFieldName, showFieldType, propertiesAsMap, firstIndex)); return shape.toString(); } + private Map getPropertiesMapForIndex(Index index) { + Map indexMapping; + try { + indexMapping = clusterService.state().metadata().findMappings(new String[] { index.getName() }, input -> str -> true); + } catch (IOException e) { + // If an error occurs while retrieving mappings, return an empty map + return Collections.emptyMap(); + } + + MappingMetadata mappingMetadata = indexMapping.get(index.getName()); + if (mappingMetadata == null) { + return Collections.emptyMap(); + } + + Map propertiesMap = (Map) mappingMetadata.getSourceAsMap().get("properties"); + if (propertiesMap == null) { + return Collections.emptyMap(); + } + return propertiesMap; + } + /** - * Method to build query-section shape - * @param queryBuilder search request query builder - * @param showFields whether to append field data - * @return Query-section shape as String + * Builds the query-section shape. + * + * @param queryBuilder search request query builder + * @param showFieldName whether to append field names + * @param showFieldType whether to append field types + * @param propertiesAsMap properties + * @param index index + * @return Query-section shape as a String */ - static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) { + String buildQueryShape( + QueryBuilder queryBuilder, + Boolean showFieldName, + Boolean showFieldType, + Map propertiesAsMap, + Index index + ) { if (queryBuilder == null) { return EMPTY_STRING; } - QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(this, propertiesAsMap, index, showFieldName, showFieldType); queryBuilder.visit(shapeVisitor); - return shapeVisitor.prettyPrintTree(EMPTY_STRING, showFields); + return shapeVisitor.prettyPrintTree(EMPTY_STRING, showFieldName, showFieldType); } /** - * Method to build aggregation shape - * @param aggregationsBuilder search request aggregation builder - * @param showFields whether to append field data - * @return Aggregation shape as String + * Builds the aggregation shape. + * + * @param aggregationsBuilder search request aggregation builder + * @param showFieldName whether to append field names + * @param showFieldType whether to append field types + * @param propertiesAsMap properties + * @param index index + * @return Aggregation shape as a String */ - static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) { + String buildAggregationShape( + AggregatorFactories.Builder aggregationsBuilder, + Boolean showFieldName, + Boolean showFieldType, + Map propertiesAsMap, + Index index + ) { if (aggregationsBuilder == null) { return EMPTY_STRING; } @@ -92,17 +197,23 @@ static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuil aggregationsBuilder.getPipelineAggregatorFactories(), new StringBuilder(), new StringBuilder(), - showFields + showFieldName, + showFieldType, + propertiesAsMap, + index ); return aggregationShape.toString(); } - static StringBuilder recursiveAggregationShapeBuilder( + StringBuilder recursiveAggregationShapeBuilder( Collection aggregationBuilders, Collection pipelineAggregations, StringBuilder outputBuilder, StringBuilder baseIndent, - Boolean showFields + Boolean showFieldName, + Boolean showFieldType, + Map propertiesAsMap, + Index index ) { //// Normal Aggregations //// if (aggregationBuilders.isEmpty() == false) { @@ -112,8 +223,8 @@ static StringBuilder recursiveAggregationShapeBuilder( for (AggregationBuilder aggBuilder : aggregationBuilders) { StringBuilder stringBuilder = new StringBuilder(); stringBuilder.append(baseIndent).append(ONE_SPACE_INDENT.repeat(2)).append(aggBuilder.getType()); - if (showFields) { - stringBuilder.append(buildFieldDataString(aggBuilder)); + if (showFieldName || showFieldType) { + stringBuilder.append(buildFieldDataString(aggBuilder, propertiesAsMap, index, showFieldName, showFieldType)); } stringBuilder.append("\n"); @@ -124,7 +235,10 @@ static StringBuilder recursiveAggregationShapeBuilder( aggBuilder.getPipelineAggregations(), stringBuilder, baseIndent.append(ONE_SPACE_INDENT.repeat(4)), - showFields + showFieldName, + showFieldType, + propertiesAsMap, + index ); baseIndent.delete(0, 4); } @@ -162,12 +276,22 @@ static StringBuilder recursiveAggregationShapeBuilder( } /** - * Method to build sort shape - * @param sortBuilderList search request sort builders list - * @param showFields whether to append field data - * @return Sort shape as String + * Builds the sort shape. + * + * @param sortBuilderList search request sort builders list + * @param showFieldName whether to append field names + * @param showFieldType whether to append field types + * @param propertiesAsMap properties + * @param index index + * @return Sort shape as a String */ - static String buildSortShape(List> sortBuilderList, Boolean showFields) { + String buildSortShape( + List> sortBuilderList, + Boolean showFieldName, + Boolean showFieldType, + Map propertiesAsMap, + Index index + ) { if (sortBuilderList == null || sortBuilderList.isEmpty()) { return EMPTY_STRING; } @@ -178,8 +302,8 @@ static String buildSortShape(List> sortBuilderList, Boolean showF for (SortBuilder sortBuilder : sortBuilderList) { StringBuilder stringBuilder = new StringBuilder(); stringBuilder.append(ONE_SPACE_INDENT.repeat(2)).append(sortBuilder.order()); - if (showFields) { - stringBuilder.append(buildFieldDataString(sortBuilder)); + if (showFieldName || showFieldType) { + stringBuilder.append(buildFieldDataString(sortBuilder, propertiesAsMap, index, showFieldName, showFieldType)); } shapeStrings.add(stringBuilder.toString()); } @@ -191,15 +315,97 @@ static String buildSortShape(List> sortBuilderList, Boolean showF } /** - * Method to build field data - * @return String: comma separated list with leading space in square brackets + * Builds a field data string from a builder. + * + * @param builder aggregation or sort builder + * @param propertiesAsMap properties + * @param index index + * @param showFieldName whether to include field names + * @param showFieldType whether to include field types + * @return Field data string * Ex: " [my_field, width:5]" */ - static String buildFieldDataString(NamedWriteable builder) { + String buildFieldDataString( + NamedWriteable builder, + Map propertiesAsMap, + Index index, + Boolean showFieldName, + Boolean showFieldType + ) { List fieldDataList = new ArrayList<>(); if (builder instanceof WithFieldName) { - fieldDataList.add(((WithFieldName) builder).fieldName()); + String fieldName = ((WithFieldName) builder).fieldName(); + if (showFieldName) { + fieldDataList.add(fieldName); + } + if (showFieldType) { + String fieldType = getFieldType(fieldName, propertiesAsMap, index); + if (fieldType != null && !fieldType.isEmpty()) { + fieldDataList.add(fieldType); + } + } } return " [" + String.join(", ", fieldDataList) + "]"; } + + String getFieldType(String fieldName, Map propertiesAsMap, Index index) { + if (propertiesAsMap == null || index == null) { + return null; + } + // Attempt to get field type from cache + String fieldType = getFieldTypeFromCache(fieldName, index); + + if (fieldType != null) { + return fieldType; + } + + // Retrieve field type from mapping and cache it if found + fieldType = getFieldTypeFromProperties(fieldName, propertiesAsMap); + + // Cache field type or NO_FIELD_TYPE_VALUE if not found + fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()) + .putIfAbsent(fieldName, fieldType != null ? fieldType : NO_FIELD_TYPE_VALUE); + + return fieldType; + } + + String getFieldTypeFromProperties(String fieldName, Map propertiesAsMap) { + if (propertiesAsMap == null) { + return null; + } + + String[] fieldParts = fieldName.split("\\."); + Map currentProperties = propertiesAsMap; + + for (int depth = 0; depth < fieldParts.length; depth++) { + Object currentMapping = currentProperties.get(fieldParts[depth]); + + if (currentMapping instanceof Map) { + Map currentMap = (Map) currentMapping; + + // Navigate into nested properties if available + if (currentMap.containsKey("properties")) { + currentProperties = (Map) currentMap.get("properties"); + } + // Handle multifields (e.g., title.raw) + else if (currentMap.containsKey("fields") && depth + 1 < fieldParts.length) { + currentProperties = (Map) currentMap.get("fields"); + } + // Return type if found + else if (currentMap.containsKey("type")) { + return (String) currentMap.get("type"); + } else { + return null; + } + } else { + return null; + } + } + + return null; + } + + String getFieldTypeFromCache(String fieldName, Index index) { + return fieldTypeMap.getOrDefault(index, new ConcurrentHashMap<>()).get(fieldName); + } } diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java index d4d0b5b..a046a52 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java @@ -9,7 +9,6 @@ package org.opensearch.plugin.insights.core.service.categorizer; import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.ONE_SPACE_INDENT; -import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.buildFieldDataString; import java.util.ArrayList; import java.util.EnumMap; @@ -18,6 +17,7 @@ import java.util.Map; import org.apache.lucene.search.BooleanClause; import org.opensearch.common.SetOnce; +import org.opensearch.core.index.Index; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilderVisitor; @@ -28,11 +28,16 @@ public final class QueryShapeVisitor implements QueryBuilderVisitor { private final SetOnce queryType = new SetOnce<>(); private final SetOnce fieldData = new SetOnce<>(); private final Map> childVisitors = new EnumMap<>(BooleanClause.Occur.class); + private final QueryShapeGenerator queryShapeGenerator; + private final Map propertiesAsMap; + private final Index index; + private final Boolean showFieldName; + private final Boolean showFieldType; @Override public void accept(QueryBuilder queryBuilder) { queryType.set(queryBuilder.getName()); - fieldData.set(buildFieldDataString(queryBuilder)); + fieldData.set(queryShapeGenerator.buildFieldDataString(queryBuilder, propertiesAsMap, index, showFieldName, showFieldType)); } @Override @@ -47,7 +52,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { @Override public void accept(QueryBuilder qb) { - currentChild = new QueryShapeVisitor(); + currentChild = new QueryShapeVisitor(queryShapeGenerator, propertiesAsMap, index, showFieldName, showFieldType); childVisitorList.add(currentChild); currentChild.accept(qb); } @@ -85,13 +90,15 @@ public String toJson() { /** * Pretty print the query builder tree - * @param indent indent size - * @param showFields whether to print field data + * + * @param indent indent size + * @param showFieldName whether to print field name + * @param showFieldType * @return Query builder tree as a pretty string */ - public String prettyPrintTree(String indent, Boolean showFields) { + public String prettyPrintTree(String indent, Boolean showFieldName, Boolean showFieldType) { StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get()); - if (showFields) { + if (showFieldName || showFieldType) { outputBuilder.append(fieldData.get()); } outputBuilder.append("\n"); @@ -101,7 +108,7 @@ public String prettyPrintTree(String indent, Boolean showFields) { .append(entry.getKey().name().toLowerCase(Locale.ROOT)) .append(":\n"); for (QueryShapeVisitor child : entry.getValue()) { - outputBuilder.append(child.prettyPrintTree(indent + ONE_SPACE_INDENT.repeat(4), showFields)); + outputBuilder.append(child.prettyPrintTree(indent + ONE_SPACE_INDENT.repeat(4), showFieldName, showFieldType)); } } return outputBuilder.toString(); @@ -110,5 +117,17 @@ public String prettyPrintTree(String indent, Boolean showFields) { /** * Default constructor */ - public QueryShapeVisitor() {} + public QueryShapeVisitor( + QueryShapeGenerator queryShapeGenerator, + Map propertiesAsMap, + Index index, + Boolean showFieldName, + Boolean showFieldType + ) { + this.queryShapeGenerator = queryShapeGenerator; + this.propertiesAsMap = propertiesAsMap; + this.index = index; + this.showFieldName = showFieldName; + this.showFieldType = showFieldType; + } } diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java index df3cb7b..43ceddd 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCategorizer.java @@ -87,11 +87,6 @@ public void categorize(SearchQueryRecord record) { incrementQueryTypeCounters(source.query(), measurements); incrementQueryAggregationCounters(source.aggregations(), measurements); incrementQuerySortCounters(source.sorts(), measurements); - - if (logger.isTraceEnabled()) { - String searchShape = QueryShapeGenerator.buildShape(source, true); - logger.trace(searchShape); - } } private void incrementQuerySortCounters(List> sorts, Map measurements) { diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java index 0226464..dc603a8 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGeneratorTests.java @@ -8,17 +8,623 @@ package org.opensearch.plugin.insights.core.service.categorizer; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.index.query.QueryBuilders.boolQuery; +import static org.opensearch.index.query.QueryBuilders.geoDistanceQuery; +import static org.opensearch.index.query.QueryBuilders.matchAllQuery; +import static org.opensearch.index.query.QueryBuilders.matchQuery; +import static org.opensearch.index.query.QueryBuilders.nestedQuery; +import static org.opensearch.index.query.QueryBuilders.rangeQuery; +import static org.opensearch.index.query.QueryBuilders.regexpQuery; +import static org.opensearch.index.query.QueryBuilders.termQuery; +import static org.opensearch.index.query.QueryBuilders.termsQuery; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import org.apache.lucene.search.join.ScoreMode; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.MappingMetadata; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.hash.MurmurHash3; +import org.opensearch.common.unit.DistanceUnit; +import org.opensearch.core.compress.CompressorRegistry; +import org.opensearch.core.index.Index; +import org.opensearch.index.query.ScriptQueryBuilder; import org.opensearch.plugin.insights.SearchSourceBuilderUtils; +import org.opensearch.script.Script; +import org.opensearch.script.ScriptType; +import org.opensearch.search.aggregations.AggregationBuilders; +import org.opensearch.search.aggregations.PipelineAggregatorBuilders; +import org.opensearch.search.aggregations.bucket.histogram.DateHistogramInterval; import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.search.sort.FieldSortBuilder; +import org.opensearch.search.sort.SortOrder; import org.opensearch.test.OpenSearchTestCase; public final class QueryShapeGeneratorTests extends OpenSearchTestCase { + final Set successfulSearchShardIndices = Set.of(new Index("index1", UUID.randomUUID().toString())); + final QueryShapeGenerator queryShapeGenerator; + + private ClusterService mockClusterService; + private ClusterState mockClusterState; + private Metadata mockMetaData; + + public QueryShapeGeneratorTests() { + CompressorRegistry.defaultCompressor(); + this.mockClusterService = mock(ClusterService.class); + this.mockClusterState = mock(ClusterState.class); + this.mockMetaData = mock(Metadata.class); + this.queryShapeGenerator = new QueryShapeGenerator(mockClusterService); + + when(mockClusterService.state()).thenReturn(mockClusterState); + when(mockClusterState.metadata()).thenReturn(mockMetaData); + } + + public void setUpMockMappings(String indexName, Map mappingProperties) throws IOException { + MappingMetadata mockMappingMetadata = mock(MappingMetadata.class); + + when(mockMappingMetadata.getSourceAsMap()).thenReturn(mappingProperties); + + final Map indexMappingMap = Map.of(indexName, mockMappingMetadata); + when(mockMetaData.findMappings(any(), any())).thenReturn(indexMappingMap); + } + + public void testBasicSearchWithFieldNameAndType() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "field1", + Map.of("type", "keyword"), + "field2", + Map.of("type", "text"), + "field3", + Map.of("type", "text"), + "field4", + Map.of("type", "long") + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder(); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [field1, keyword]\n" + + " filter:\n" + + " match [field2, text]\n" + + " range [field4, long]\n" + + " should:\n" + + " regexp [field3, text]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + // If field type is not found we leave the field type blank + public void testFieldTypeNotFound() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("field1", Map.of("type", "keyword")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query( + boolQuery().must(termQuery("field1", "value1")) + .filter(matchQuery("field2", "value2")) + .filter(rangeQuery("field4").gte(10).lte(20)) + .should(regexpQuery("field3", ".*pattern.*")) + ); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [field1, keyword]\n" + + " filter:\n" + + " match [field2]\n" + + " range [field4]\n" + + " should:\n" + + " regexp [field3]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testEmptyMappings() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of() // No fields defined + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query( + boolQuery().must(termQuery("field1", "value1")) + .filter(matchQuery("field2", "value2")) + .filter(rangeQuery("field4").gte(10).lte(20)) + .should(regexpQuery("field3", ".*pattern.*")) + ); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [field1]\n" + + " filter:\n" + + " match [field2]\n" + + " range [field4]\n" + + " should:\n" + + " regexp [field3]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + // Field type should be inferred from both the mappings + public void testMultipleIndexMappings() throws IOException { + setUpMockMappings("index2", Map.of("properties", Map.of("field1", Map.of("type", "keyword")))); + + setUpMockMappings("index1", Map.of("properties", Map.of("field2", Map.of("type", "text"), "field4", Map.of("type", "long")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder(); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [field1]\n" + + " filter:\n" + + " match [field2, text]\n" + + " range [field4, long]\n" + + " should:\n" + + " regexp [field3]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testDifferentFieldTypes() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "field1", + Map.of("type", "keyword"), + "field2", + Map.of("type", "text"), + "field3", + Map.of("type", "integer"), + "field4", + Map.of("type", "boolean") + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder(); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [field1, keyword]\n" + + " filter:\n" + + " match [field2, text]\n" + + " range [field4, boolean]\n" + + " should:\n" + + " regexp [field3, integer]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithNestedProperties() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "nestedField", + Map.of( + "type", + "nested", + "properties", + Map.of("subField1", Map.of("type", "keyword"), "subField2", Map.of("type", "text")) + ) + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(termQuery("nestedField.subField1", "value1")).filter(matchQuery("nestedField.subField2", "value2"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [nestedField.subField1, keyword]\n" + + " filter:\n" + + " match [nestedField.subField2, text]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithArrayType() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of("field1", Map.of("type", "keyword"), "field2", Map.of("type", "text"), "field3", Map.of("type", "keyword")) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query( + boolQuery().must(termQuery("field1", "value1")) + .filter(matchQuery("field2", "value2")) + .should(termsQuery("field3", "value3a", "value3b")) + ); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [field1, keyword]\n" + + " filter:\n" + + " match [field2, text]\n" + + " should:\n" + + " terms [field3, keyword]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithDateType() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("dateField", Map.of("type", "date")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().filter(rangeQuery("dateField").gte("2024-01-01").lte("2024-12-31"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + " filter:\n" + " range [dateField, date]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithGeoPointType() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("location", Map.of("type", "geo_point")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().filter(geoDistanceQuery("location").point(40.73, -74.1).distance(200, DistanceUnit.KILOMETERS))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + " filter:\n" + " geo_distance [location, geo_point]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithBinaryType() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("binaryField", Map.of("type", "binary")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(termQuery("binaryField", "base64EncodedString"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + " must:\n" + " term [binaryField, binary]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithMixedTypes() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "mixedField", + Map.of( + "type", + "object", + "properties", + Map.of("subField1", Map.of("type", "keyword"), "subField2", Map.of("type", "text")) + ) + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(termQuery("mixedField.subField1", "value1")).filter(matchQuery("mixedField.subField2", "value2"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [mixedField.subField1, keyword]\n" + + " filter:\n" + + " match [mixedField.subField2, text]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithInvalidQueries() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("field1", Map.of("type", "keyword")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query( + boolQuery().must(termQuery("invalidField", "value1")) // Invalid field + ); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + " must:\n" + " term [invalidField]\n"; // No type info, just the invalid field + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithDeeplyNestedStructure() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "level1", + Map.of( + "type", + "nested", + "properties", + Map.of("level2", Map.of("type", "nested", "properties", Map.of("level3", Map.of("type", "keyword")))) + ) + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(termQuery("level1.level2.level3", "value1"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + " must:\n" + " term [level1.level2.level3, keyword]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + // We are not parsing fields for scripts + public void testFieldWithScriptedQuery() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("scriptedField", Map.of("type", "long")))); + + Script script = new Script( + ScriptType.INLINE, + "p.params.threshold < doc['scriptedField'].value", + "mockscript", + Map.of("threshold", 100) + ); + + ScriptQueryBuilder scriptedQuery = new ScriptQueryBuilder(script); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder().query(scriptedQuery); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "script []\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testDynamicTemplateMappingWithTypeInference() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "dynamic_templates", + List.of( + Map.of("fields", Map.of("mapping", Map.of("type", "short"), "match_mapping_type", "string", "path_match", "status*")) + ), + "properties", + Map.of("status_code", Map.of("type", "short"), "other_field", Map.of("type", "text")) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(termQuery("status_code", "200")).filter(matchQuery("other_field", "value"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [status_code, short]\n" + + " filter:\n" + + " match [other_field, text]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFieldWithIpAddressType() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("ip_address", Map.of("type", "ip", "ignore_malformed", true)))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(termQuery("ip_address", "192.168.1.1")).filter(termQuery("ip_address", "invalid_ip"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " term [ip_address, ip]\n" + + " filter:\n" + + " term [ip_address, ip]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + // Nested query not working as expected + public void testNestedQueryType() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "patients", + Map.of( + "type", + "nested", + "properties", + Map.of("name", Map.of("type", "text"), "age", Map.of("type", "integer"), "smoker", Map.of("type", "boolean")) + ) + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query( + nestedQuery( + "patients", + boolQuery().should(termQuery("patients.smoker", true)).should(rangeQuery("patients.age").gte(75)), + ScoreMode.Avg + ) + ); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "nested []\n" + " must:\n" + " bool []\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testFlatObjectQueryType() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of( + "issue", + Map.of( + "type", + "flat_object", + "properties", + Map.of( + "number", + Map.of("type", "keyword"), + "labels", + Map.of( + "properties", + Map.of( + "version", + Map.of("type", "keyword"), + "category", + Map.of("properties", Map.of("type", Map.of("type", "keyword"), "level", Map.of("type", "keyword"))) + ) + ) + ) + ) + ) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(matchQuery("issue.labels.category.level", "bug")); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "match [issue.labels.category.level, keyword]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testQueryTypeWithSorting() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of("age", Map.of("type", "integer"), "name", Map.of("type", "keyword"), "score", Map.of("type", "float")) + ) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(matchQuery("name", "John")) + .sort(new FieldSortBuilder("age").order(SortOrder.ASC)) + .sort(new FieldSortBuilder("score").order(SortOrder.DESC)); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "match [name, keyword]\n" + "sort:\n" + " asc [age, integer]\n" + " desc [score, float]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + public void testQueryTypeWithAggregations() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("price", Map.of("type", "double"), "category", Map.of("type", "keyword")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(matchAllQuery()) + .aggregation(AggregationBuilders.terms("categories").field("category")) + .aggregation(AggregationBuilders.avg("average_price").field("price")); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "match_all []\n" + "aggregation:\n" + " avg [price, double]\n" + " terms [category, keyword]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + // No field name and type being parsed for pipeline aggregations + public void testQueryTypeWithPipelineAggregation() throws IOException { + setUpMockMappings("index1", Map.of("properties", Map.of("sales", Map.of("type", "double"), "timestamp", Map.of("type", "date")))); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(matchAllQuery()) + .aggregation( + AggregationBuilders.dateHistogram("sales_over_time") + .field("timestamp") + .calendarInterval(DateHistogramInterval.MONTH) + .subAggregation(AggregationBuilders.sum("total_sales").field("sales")) + ) + .aggregation(PipelineAggregatorBuilders.derivative("sales_derivative", "total_sales")); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "match_all []\n" + + "aggregation:\n" + + " date_histogram [timestamp, date]\n" + + " aggregation:\n" + + " sum [sales, double]\n" + + " pipeline aggregation:\n" + + " derivative\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } + + // Should cache empty value when we do not find a field type to avoid doing the search again + public void testFieldTypeCachingForNonExistentField() throws IOException { + setUpMockMappings( + "index1", + Map.of( + "properties", + Map.of("age", Map.of("type", "integer"), "name", Map.of("type", "keyword"), "score", Map.of("type", "float")) + ) + ); + + QueryShapeGenerator queryShapeGeneratorSpy = spy(queryShapeGenerator); + + String nonExistentField = "nonExistentField"; + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(matchQuery(nonExistentField, "value")); + + queryShapeGeneratorSpy.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), any(Index.class)); + + queryShapeGeneratorSpy.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + verify(queryShapeGeneratorSpy, atLeastOnce()).getFieldTypeFromCache(eq(nonExistentField), any(Index.class)); + } + + public void testMultifieldQueryCombined() throws IOException { + setUpMockMappings( + "index1", + Map.of("properties", Map.of("title", Map.of("type", "text", "fields", Map.of("raw", Map.of("type", "keyword"))))) + ); + + SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder() + .query(boolQuery().must(matchQuery("title", "eg")).should(termQuery("title.raw", "e_g"))); + + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, true, successfulSearchShardIndices); + + String expectedShowFieldsTrue = "bool []\n" + + " must:\n" + + " match [title, text]\n" + + " should:\n" + + " term [title.raw, keyword]\n"; + + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + } public void testComplexSearch() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createDefaultSearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, false, successfulSearchShardIndices); String expectedShowFieldsTrue = "bool []\n" + " must:\n" + " term [field1]\n" @@ -53,7 +659,7 @@ public void testComplexSearch() { + " asc [album]\n"; assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String shapeShowFieldsFalse = queryShapeGenerator.buildShape(sourceBuilder, false, false, successfulSearchShardIndices); String expectedShowFieldsFalse = "bool\n" + " must:\n" + " term\n" @@ -89,36 +695,10 @@ public void testComplexSearch() { assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse); } - public void testQueryShape() { - SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder(); - - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); - String expectedShowFieldsTrue = "bool []\n" - + " must:\n" - + " term [field1]\n" - + " filter:\n" - + " match [field2]\n" - + " range [field4]\n" - + " should:\n" - + " regexp [field3]\n"; - assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); - String expectedShowFieldsFalse = "bool\n" - + " must:\n" - + " term\n" - + " filter:\n" - + " match\n" - + " range\n" - + " should:\n" - + " regexp\n"; - assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse); - } - public void testAggregationShape() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createAggregationSearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, false, successfulSearchShardIndices); String expectedShowFieldsTrue = "aggregation:\n" + " significant_text []\n" + " terms [key]\n" @@ -140,7 +720,7 @@ public void testAggregationShape() { + " max_bucket\n"; assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String shapeShowFieldsFalse = queryShapeGenerator.buildShape(sourceBuilder, false, false, successfulSearchShardIndices); String expectedShowFieldsFalse = "aggregation:\n" + " significant_text\n" + " terms\n" @@ -166,11 +746,11 @@ public void testAggregationShape() { public void testSortShape() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createSortSearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, false, successfulSearchShardIndices); String expectedShowFieldsTrue = "sort:\n" + " desc [color]\n" + " desc [vendor]\n" + " asc [price]\n" + " asc [album]\n"; assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String shapeShowFieldsFalse = queryShapeGenerator.buildShape(sourceBuilder, false, false, successfulSearchShardIndices); String expectedShowFieldsFalse = "sort:\n" + " desc\n" + " desc\n" + " asc\n" + " asc\n"; assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse); } @@ -181,17 +761,43 @@ public void testHashCode() { SearchSourceBuilder querySourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder(); // showFields true - MurmurHash3.Hash128 defaultHashTrue = QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, true); - MurmurHash3.Hash128 queryHashTrue = QueryShapeGenerator.getShapeHashCode(querySourceBuilder, true); - assertEquals(defaultHashTrue, QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, true)); - assertEquals(queryHashTrue, QueryShapeGenerator.getShapeHashCode(querySourceBuilder, true)); + MurmurHash3.Hash128 defaultHashTrue = queryShapeGenerator.getShapeHashCode( + defaultSourceBuilder, + true, + false, + successfulSearchShardIndices + ); + MurmurHash3.Hash128 queryHashTrue = queryShapeGenerator.getShapeHashCode( + querySourceBuilder, + true, + false, + successfulSearchShardIndices + ); + assertEquals( + defaultHashTrue, + queryShapeGenerator.getShapeHashCode(defaultSourceBuilder, true, false, successfulSearchShardIndices) + ); + assertEquals(queryHashTrue, queryShapeGenerator.getShapeHashCode(querySourceBuilder, true, false, successfulSearchShardIndices)); assertNotEquals(defaultHashTrue, queryHashTrue); // showFields false - MurmurHash3.Hash128 defaultHashFalse = QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, false); - MurmurHash3.Hash128 queryHashFalse = QueryShapeGenerator.getShapeHashCode(querySourceBuilder, false); - assertEquals(defaultHashFalse, QueryShapeGenerator.getShapeHashCode(defaultSourceBuilder, false)); - assertEquals(queryHashFalse, QueryShapeGenerator.getShapeHashCode(querySourceBuilder, false)); + MurmurHash3.Hash128 defaultHashFalse = queryShapeGenerator.getShapeHashCode( + defaultSourceBuilder, + false, + false, + successfulSearchShardIndices + ); + MurmurHash3.Hash128 queryHashFalse = queryShapeGenerator.getShapeHashCode( + querySourceBuilder, + false, + false, + successfulSearchShardIndices + ); + assertEquals( + defaultHashFalse, + queryShapeGenerator.getShapeHashCode(defaultSourceBuilder, false, false, successfulSearchShardIndices) + ); + assertEquals(queryHashFalse, queryShapeGenerator.getShapeHashCode(querySourceBuilder, false, false, successfulSearchShardIndices)); assertNotEquals(defaultHashFalse, queryHashFalse); // Compare field data on vs off diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java index 298f4f8..5c9e4ba 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitorTests.java @@ -8,6 +8,10 @@ package org.opensearch.plugin.insights.core.service.categorizer; +import static org.mockito.Mockito.mock; + +import java.util.HashMap; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.ConstantScoreQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; @@ -27,7 +31,13 @@ public void testQueryShapeVisitor() { .mustNot(new RegexpQueryBuilder("color", "red.*")) ) .must(new TermsQueryBuilder("genre", "action", "drama", "romance")); - QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor( + new QueryShapeGenerator(mock(ClusterService.class)), + new HashMap<>(), + null, + false, + false + ); builder.visit(shapeVisitor); assertEquals( "{\"type\":\"bool\",\"must\"[{\"type\":\"term\"},{\"type\":\"terms\"}],\"filter\"[{\"type\":\"constant_score\",\"filter\"[{\"type\":\"range\"}]}],\"should\"[{\"type\":\"bool\",\"must\"[{\"type\":\"match\"}],\"must_not\"[{\"type\":\"regexp\"}]}]}",