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 33ea533..d3cf48d 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 @@ -15,6 +15,7 @@ import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getTopNSizeSetting; import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getTopNWindowSizeSetting; +import java.io.IOException; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -28,6 +29,7 @@ import org.opensearch.action.search.SearchRequestContext; import org.opensearch.action.search.SearchRequestOperationsListener; import org.opensearch.action.search.SearchTask; +import org.opensearch.cluster.metadata.MappingMetadata; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo; @@ -54,6 +56,7 @@ public final class QueryInsightsListener extends SearchRequestOperationsListener private final QueryInsightsService queryInsightsService; private final ClusterService clusterService; + private final QueryShapeGenerator queryShapeGenerator; /** * Constructor for QueryInsightsListener @@ -81,6 +84,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. @@ -241,7 +245,15 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final ); } - String hashcode = QueryShapeGenerator.getShapeHashCodeAsString(request.source(), false); + String hashcode = queryShapeGenerator.getShapeHashCodeAsString( + request.source(), + false, + searchRequestContext.getSuccessfulSearchShardIndices() + ); + + String shape = queryShapeGenerator.buildShape(request.source(), true, searchRequestContext.getSuccessfulSearchShardIndices()); + System.out.println("===== Query Shape ====="); + System.out.println(shape); Map attributes = new HashMap<>(); attributes.put(Attribute.SEARCH_TYPE, request.searchType().toString().toLowerCase(Locale.ROOT)); diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/MappingVisitor.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/MappingVisitor.java new file mode 100644 index 0000000..95a0502 --- /dev/null +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/MappingVisitor.java @@ -0,0 +1,76 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.plugin.insights.core.service.categorizer; + +import java.util.Map; + +/** + * Visitor pattern for obtaining index mappings + * + * @opensearch.internal + */ +final class MappingVisitor { + + private MappingVisitor() {} + + @SuppressWarnings("unchecked") + static String visitMapping(Map propertiesAsMap, String[] fieldName, int idx) { + if (propertiesAsMap.containsKey(fieldName[idx])) { + Map fieldMapping = (Map) propertiesAsMap.get(fieldName[idx]); + if (idx == fieldName.length - 1) { + return (String) fieldMapping.get("type"); + } else { + // Multi field case + if (fieldMapping.containsKey("fields")) { + return visitMapping((Map) fieldMapping.get("fields"), fieldName, idx + 1); + } else { + return null; + } + } + } + + // fieldName not found at current level + // call visitMapping() for any fields with subfields (contains "properties" key) + for (Object v : propertiesAsMap.values()) { + if (v instanceof Map) { + Map fieldMapping = (Map) v; + if (fieldMapping.containsKey("properties")) { + return visitMapping((Map) fieldMapping.get("properties"), fieldName, idx); + } + } + } + // no mapping found + return null; + } + +} 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..8116d90 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,6 +36,13 @@ public class QueryShapeGenerator { static final String EMPTY_STRING = ""; static final String ONE_SPACE_INDENT = " "; + private final ClusterService clusterService; + 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 @@ -36,14 +50,14 @@ public class QueryShapeGenerator { * @param showFields whether to include field data in query shape * @return Hash code of 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 showFields, Set successfulSearchShardIndices) { + final String shape = buildShape(source, showFields, 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); + public String getShapeHashCodeAsString(SearchSourceBuilder source, Boolean showFields, Set successfulSearchShardIndices) { + MurmurHash3.Hash128 hashcode = getShapeHashCode(source, showFields, successfulSearchShardIndices); String hashAsString = Long.toHexString(hashcode.h1) + Long.toHexString(hashcode.h2); return hashAsString; } @@ -54,11 +68,11 @@ public static String getShapeHashCodeAsString(SearchSourceBuilder source, Boolea * @param showFields whether to append field data * @return Search query shape as String */ - public static String buildShape(SearchSourceBuilder source, Boolean showFields) { + public String buildShape(SearchSourceBuilder source, Boolean showFields, Set successfulSearchShardIndices) { 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(), showFields, successfulSearchShardIndices)); + shape.append(buildAggregationShape(source.aggregations(), showFields, successfulSearchShardIndices)); + shape.append(buildSortShape(source.sorts(), showFields, successfulSearchShardIndices)); return shape.toString(); } @@ -68,11 +82,11 @@ public static String buildShape(SearchSourceBuilder source, Boolean showFields) * @param showFields whether to append field data * @return Query-section shape as String */ - static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) { + String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields, Set successfulSearchShardIndices) { if (queryBuilder == null) { return EMPTY_STRING; } - QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(this, successfulSearchShardIndices); queryBuilder.visit(shapeVisitor); return shapeVisitor.prettyPrintTree(EMPTY_STRING, showFields); } @@ -83,7 +97,11 @@ static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) { * @param showFields whether to append field data * @return Aggregation shape as String */ - static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) { + String buildAggregationShape( + AggregatorFactories.Builder aggregationsBuilder, + Boolean showFields, + Set successfulSearchShardIndices + ) { if (aggregationsBuilder == null) { return EMPTY_STRING; } @@ -92,17 +110,19 @@ static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuil aggregationsBuilder.getPipelineAggregatorFactories(), new StringBuilder(), new StringBuilder(), - showFields + showFields, + successfulSearchShardIndices ); return aggregationShape.toString(); } - static StringBuilder recursiveAggregationShapeBuilder( + StringBuilder recursiveAggregationShapeBuilder( Collection aggregationBuilders, Collection pipelineAggregations, StringBuilder outputBuilder, StringBuilder baseIndent, - Boolean showFields + Boolean showFields, + Set successfulSearchShardIndices ) { //// Normal Aggregations //// if (aggregationBuilders.isEmpty() == false) { @@ -113,7 +133,7 @@ static StringBuilder recursiveAggregationShapeBuilder( StringBuilder stringBuilder = new StringBuilder(); stringBuilder.append(baseIndent).append(ONE_SPACE_INDENT.repeat(2)).append(aggBuilder.getType()); if (showFields) { - stringBuilder.append(buildFieldDataString(aggBuilder)); + stringBuilder.append(buildFieldDataString(aggBuilder, successfulSearchShardIndices)); } stringBuilder.append("\n"); @@ -124,7 +144,8 @@ static StringBuilder recursiveAggregationShapeBuilder( aggBuilder.getPipelineAggregations(), stringBuilder, baseIndent.append(ONE_SPACE_INDENT.repeat(4)), - showFields + showFields, + successfulSearchShardIndices ); baseIndent.delete(0, 4); } @@ -167,7 +188,7 @@ static StringBuilder recursiveAggregationShapeBuilder( * @param showFields whether to append field data * @return Sort shape as String */ - static String buildSortShape(List> sortBuilderList, Boolean showFields) { + String buildSortShape(List> sortBuilderList, Boolean showFields, Set successfulSearchShardIndices) { if (sortBuilderList == null || sortBuilderList.isEmpty()) { return EMPTY_STRING; } @@ -179,7 +200,7 @@ static String buildSortShape(List> sortBuilderList, Boolean showF StringBuilder stringBuilder = new StringBuilder(); stringBuilder.append(ONE_SPACE_INDENT.repeat(2)).append(sortBuilder.order()); if (showFields) { - stringBuilder.append(buildFieldDataString(sortBuilder)); + stringBuilder.append(buildFieldDataString(sortBuilder, successfulSearchShardIndices)); } shapeStrings.add(stringBuilder.toString()); } @@ -195,11 +216,47 @@ static String buildSortShape(List> sortBuilderList, Boolean showF * @return String: comma separated list with leading space in square brackets * Ex: " [my_field, width:5]" */ - static String buildFieldDataString(NamedWriteable builder) { + String buildFieldDataString(NamedWriteable builder, Set successfulSearchShardIndices) { List fieldDataList = new ArrayList<>(); if (builder instanceof WithFieldName) { - fieldDataList.add(((WithFieldName) builder).fieldName()); + String fieldName = ((WithFieldName) builder).fieldName(); + fieldDataList.add(fieldName); + fieldDataList.add(getFieldType(fieldName, successfulSearchShardIndices)); } return " [" + String.join(", ", fieldDataList) + "]"; } + + String getFieldType(String fieldName, Set successfulSearchShardIndices) { + for (Index index : successfulSearchShardIndices) { + Map indexMap = fieldTypeMap.get(index); + if (indexMap != null) { + String fieldType = indexMap.get(fieldName); + if (fieldType != null) { + return fieldType; + } + } + } + Map allIndicesMap; + try { + allIndicesMap = clusterService.state() + .metadata() + .findMappings(successfulSearchShardIndices.stream().map(Index::getName).toArray(String[]::new), input -> (str -> true)); + } catch (IOException e) { + return null; + } + for (Index index : successfulSearchShardIndices) { + MappingMetadata mappingMetadata = allIndicesMap.get(index.getName()); + if (mappingMetadata != null) { + @SuppressWarnings("unchecked") + Map propertiesAsMap = (Map) mappingMetadata.getSourceAsMap().get("properties"); + String fieldType = MappingVisitor.visitMapping(propertiesAsMap, fieldName.split("\\."), 0); + if (fieldType != null) { + // add item to cache + fieldTypeMap.computeIfAbsent(index, k -> new ConcurrentHashMap<>()).put(fieldName, fieldType); + return fieldType; + } + } + } + return null; + } } 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..24aed09 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,15 +9,16 @@ 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; import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.Set; 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 +29,13 @@ 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 Set successfulSearchShardIndices; @Override public void accept(QueryBuilder queryBuilder) { queryType.set(queryBuilder.getName()); - fieldData.set(buildFieldDataString(queryBuilder)); + fieldData.set(queryShapeGenerator.buildFieldDataString(queryBuilder, successfulSearchShardIndices)); } @Override @@ -47,7 +50,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { @Override public void accept(QueryBuilder qb) { - currentChild = new QueryShapeVisitor(); + currentChild = new QueryShapeVisitor(queryShapeGenerator, successfulSearchShardIndices); childVisitorList.add(currentChild); currentChild.accept(qb); } @@ -110,5 +113,8 @@ public String prettyPrintTree(String indent, Boolean showFields) { /** * Default constructor */ - public QueryShapeVisitor() {} + public QueryShapeVisitor(QueryShapeGenerator queryShapeGenerator, Set successfulSearchShardIndices) { + this.queryShapeGenerator = queryShapeGenerator; + this.successfulSearchShardIndices = successfulSearchShardIndices; + } } 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..a3bf461 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,71 @@ package org.opensearch.plugin.insights.core.service.categorizer; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import java.util.Set; +import java.util.UUID; + +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.core.index.Index; +import org.opensearch.index.IndexService; +import org.opensearch.index.mapper.MapperService; import org.opensearch.plugin.insights.SearchSourceBuilderUtils; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.test.OpenSearchTestCase; public final class QueryShapeGeneratorTests extends OpenSearchTestCase { + final Set successfulSearchShardIndices = Set.of( + new Index("index1", UUID.randomUUID().toString()), + new Index("index2", UUID.randomUUID().toString()) + ); + final QueryShapeGenerator queryShapeGenerator; + + public QueryShapeGeneratorTests() { + final ClusterService mockClusterService = mock(ClusterService.class); + final ClusterState mockClusterState = mock(ClusterState.class); + final Metadata mockMetaData = mock(Metadata.class); +// final MappingMetadata mockMappingMetadata = mock(MappingMetadata.class); + final MappingMetadata mappingMetadata = new MappingMetadata(MapperService.SINGLE_MAPPING_NAME, Collections.emptyMap()); + final Map map = Map.of("index1", mappingMetadata); + when(mockClusterService.state()).thenReturn(mockClusterState); + when(mockClusterState.metadata()).thenReturn(mockMetaData); + Map finalMap = Map.of( + "properties", Map.of( + "age", Map.of("type", "long"), + "name", Map.of( + "type", "text", + "fields", Map.of( + "keyword_field_name", Map.of( + "type", "keyword", + "ignore_above", 256 + ) + ) + ) + ) + ); +// when(mockMappingMetadata.getSourceAsMap()).thenReturn(finalMap); + try { + when(mockMetaData.findMappings(any(), any())).thenReturn(map); + } catch (IOException e) { + throw new RuntimeException(e); + } + this.queryShapeGenerator = new QueryShapeGenerator(mockClusterService); + } public void testComplexSearch() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createDefaultSearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, successfulSearchShardIndices); String expectedShowFieldsTrue = "bool []\n" + " must:\n" + " term [field1]\n" @@ -53,7 +107,7 @@ public void testComplexSearch() { + " asc [album]\n"; assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String shapeShowFieldsFalse = queryShapeGenerator.buildShape(sourceBuilder, false, successfulSearchShardIndices); String expectedShowFieldsFalse = "bool\n" + " must:\n" + " term\n" @@ -92,7 +146,7 @@ public void testComplexSearch() { public void testQueryShape() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createQuerySearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, successfulSearchShardIndices); String expectedShowFieldsTrue = "bool []\n" + " must:\n" + " term [field1]\n" @@ -103,7 +157,7 @@ public void testQueryShape() { + " regexp [field3]\n"; assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String shapeShowFieldsFalse = queryShapeGenerator.buildShape(sourceBuilder, false, successfulSearchShardIndices); String expectedShowFieldsFalse = "bool\n" + " must:\n" + " term\n" @@ -118,7 +172,7 @@ public void testQueryShape() { public void testAggregationShape() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createAggregationSearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, successfulSearchShardIndices); String expectedShowFieldsTrue = "aggregation:\n" + " significant_text []\n" + " terms [key]\n" @@ -140,7 +194,7 @@ public void testAggregationShape() { + " max_bucket\n"; assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String shapeShowFieldsFalse = queryShapeGenerator.buildShape(sourceBuilder, false, successfulSearchShardIndices); String expectedShowFieldsFalse = "aggregation:\n" + " significant_text\n" + " terms\n" @@ -166,11 +220,11 @@ public void testAggregationShape() { public void testSortShape() { SearchSourceBuilder sourceBuilder = SearchSourceBuilderUtils.createSortSearchSourceBuilder(); - String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = queryShapeGenerator.buildShape(sourceBuilder, true, 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, successfulSearchShardIndices); String expectedShowFieldsFalse = "sort:\n" + " desc\n" + " desc\n" + " asc\n" + " asc\n"; assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse); } @@ -181,17 +235,25 @@ 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, + successfulSearchShardIndices + ); + MurmurHash3.Hash128 queryHashTrue = queryShapeGenerator.getShapeHashCode(querySourceBuilder, true, successfulSearchShardIndices); + assertEquals(defaultHashTrue, queryShapeGenerator.getShapeHashCode(defaultSourceBuilder, true, successfulSearchShardIndices)); + assertEquals(queryHashTrue, queryShapeGenerator.getShapeHashCode(querySourceBuilder, true, 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, + successfulSearchShardIndices + ); + MurmurHash3.Hash128 queryHashFalse = queryShapeGenerator.getShapeHashCode(querySourceBuilder, false, successfulSearchShardIndices); + assertEquals(defaultHashFalse, queryShapeGenerator.getShapeHashCode(defaultSourceBuilder, false, successfulSearchShardIndices)); + assertEquals(queryHashFalse, queryShapeGenerator.getShapeHashCode(querySourceBuilder, 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..830603a 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.Set; +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,7 @@ 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)), Set.of()); 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\"}]}]}",