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 802660e6..900f61ab 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,7 +8,27 @@ package org.opensearch.plugin.insights.core.service.categorizer; +import org.opensearch.index.query.AbstractGeometryQueryBuilder; +import org.opensearch.index.query.CommonTermsQueryBuilder; +import org.opensearch.index.query.ExistsQueryBuilder; +import org.opensearch.index.query.FieldMaskingSpanQueryBuilder; +import org.opensearch.index.query.FuzzyQueryBuilder; +import org.opensearch.index.query.GeoDistanceQueryBuilder; +import org.opensearch.index.query.GeoPolygonQueryBuilder; +import org.opensearch.index.query.MatchBoolPrefixQueryBuilder; +import org.opensearch.index.query.MatchPhrasePrefixQueryBuilder; +import org.opensearch.index.query.MatchPhraseQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.index.query.MultiTermQueryBuilder; +import org.opensearch.index.query.PrefixQueryBuilder; import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.index.query.RegexpQueryBuilder; +import org.opensearch.index.query.SpanNearQueryBuilder; +import org.opensearch.index.query.SpanTermQueryBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.TermsQueryBuilder; +import org.opensearch.index.query.WildcardQueryBuilder; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregatorFactories; import org.opensearch.search.aggregations.PipelineAggregationBuilder; @@ -21,25 +41,59 @@ import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.function.Function; public class QueryShapeGenerator { static final String TWO_SPACE_INDENT = " "; + static final Map, List>> QUERY_FIELD_DATA_MAP = Map.ofEntries( + Map.entry(AbstractGeometryQueryBuilder.class, List.of(obj -> ((AbstractGeometryQueryBuilder) obj).fieldName())), + Map.entry(CommonTermsQueryBuilder.class, List.of(obj -> ((CommonTermsQueryBuilder) obj).fieldName())), + Map.entry(org.opensearch.index.query.ExistsQueryBuilder.class, List.of(obj -> ((ExistsQueryBuilder) obj).fieldName())), + Map.entry( + org.opensearch.index.query.FieldMaskingSpanQueryBuilder.class, + List.of(obj -> ((FieldMaskingSpanQueryBuilder) obj).fieldName()) + ), + Map.entry(FuzzyQueryBuilder.class, List.of(obj -> ((FuzzyQueryBuilder) obj).fieldName())), + Map.entry( + org.opensearch.index.query.GeoBoundingBoxQueryBuilder.class, + List.of(obj -> ((org.opensearch.index.query.GeoBoundingBoxQueryBuilder) obj).fieldName()) + ), + Map.entry(org.opensearch.index.query.GeoDistanceQueryBuilder.class, List.of(obj -> ((GeoDistanceQueryBuilder) obj).fieldName())), + Map.entry(GeoPolygonQueryBuilder.class, List.of(obj -> ((GeoPolygonQueryBuilder) obj).fieldName())), + Map.entry(MatchBoolPrefixQueryBuilder.class, List.of(obj -> ((MatchBoolPrefixQueryBuilder) obj).fieldName())), + Map.entry(MatchQueryBuilder.class, List.of(obj -> ((MatchQueryBuilder) obj).fieldName())), + Map.entry(org.opensearch.index.query.MatchPhraseQueryBuilder.class, List.of(obj -> ((MatchPhraseQueryBuilder) obj).fieldName())), + Map.entry(MatchPhrasePrefixQueryBuilder.class, List.of(obj -> ((MatchPhrasePrefixQueryBuilder) obj).fieldName())), + Map.entry(MultiTermQueryBuilder.class, List.of(obj -> ((MultiTermQueryBuilder) obj).fieldName())), + Map.entry(PrefixQueryBuilder.class, List.of(obj -> ((PrefixQueryBuilder) obj).fieldName())), + Map.entry(RangeQueryBuilder.class, List.of(obj -> ((RangeQueryBuilder) obj).fieldName())), + Map.entry(RegexpQueryBuilder.class, List.of(obj -> ((RegexpQueryBuilder) obj).fieldName())), + Map.entry( + SpanNearQueryBuilder.SpanGapQueryBuilder.class, + List.of(obj -> ((SpanNearQueryBuilder.SpanGapQueryBuilder) obj).fieldName()) + ), + Map.entry(SpanTermQueryBuilder.class, List.of(obj -> ((SpanTermQueryBuilder) obj).fieldName())), + Map.entry(TermQueryBuilder.class, List.of(obj -> ((TermQueryBuilder) obj).fieldName())), + Map.entry(TermsQueryBuilder.class, List.of(obj -> ((TermsQueryBuilder) obj).fieldName())), + Map.entry(WildcardQueryBuilder.class, List.of(obj -> ((WildcardQueryBuilder) obj).fieldName())) + ); public static String buildShape(SearchSourceBuilder source, Boolean showFields) { StringBuilder shape = new StringBuilder(); - shape.append(buildQueryShape(source.query())); + shape.append(buildQueryShape(source.query(), showFields)); shape.append(buildAggregationShape(source.aggregations(), showFields)); shape.append(buildSortShape(source.sorts(), showFields)); return shape.toString(); } - static String buildQueryShape(QueryBuilder queryBuilder) { + static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) { if (queryBuilder == null) { return ""; } QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); queryBuilder.visit(shapeVisitor); - return shapeVisitor.prettyPrintTree(""); + return shapeVisitor.prettyPrintTree("", showFields); } static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) { 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 626e05dc..4e604978 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 @@ -18,17 +18,31 @@ import java.util.List; import java.util.Locale; import java.util.Map; +import java.util.function.Function; + +import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.QUERY_FIELD_DATA_MAP; +import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.TWO_SPACE_INDENT; /** * Class to traverse the QueryBuilder tree and capture the query shape */ 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); @Override - public void accept(QueryBuilder qb) { - queryType.set(qb.getName()); + public void accept(QueryBuilder queryBuilder) { + queryType.set(queryBuilder.getName()); + + List fieldDataList = new ArrayList<>(); + List> methods = QUERY_FIELD_DATA_MAP.get(queryBuilder.getClass()); + if (methods != null) { + for (Function lambda : methods) { + fieldDataList.add(lambda.apply(queryBuilder)); + } + } + fieldData.set(String.join(", ", fieldDataList)); } @Override @@ -84,12 +98,16 @@ public String toJson() { * @param indent indent size * @return Query builder tree as a pretty string */ - public String prettyPrintTree(String indent) { - StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get()).append("\n"); + public String prettyPrintTree(String indent, Boolean showFields) { + StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get()); + if (showFields) { + outputBuilder.append(" [").append(fieldData.get()).append("]"); + } + outputBuilder.append("\n"); for (Map.Entry> entry : childVisitors.entrySet()) { - outputBuilder.append(indent).append(" ").append(entry.getKey().name().toLowerCase(Locale.ROOT)).append(":\n"); + outputBuilder.append(indent).append(TWO_SPACE_INDENT).append(entry.getKey().name().toLowerCase(Locale.ROOT)).append(":\n"); for (QueryShapeVisitor child : entry.getValue()) { - outputBuilder.append(child.prettyPrintTree(indent + " ")); + outputBuilder.append(child.prettyPrintTree(indent + TWO_SPACE_INDENT.repeat(2), showFields)); } } return outputBuilder.toString(); diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeGeneratorTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeGeneratorTests.java index 56c9de46..b49a3b62 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeGeneratorTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeGeneratorTests.java @@ -11,6 +11,7 @@ import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.RangeQueryBuilder; import org.opensearch.index.query.RegexpQueryBuilder; import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator; @@ -24,15 +25,35 @@ public final class QueryShapeGeneratorTests extends OpenSearchTestCase { public void testQueryShape() { SearchSourceBuilder sourceBuilder = new SearchSourceBuilder(); sourceBuilder.size(0); - TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field", "value2"); - MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("tags", "php"); - RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field", "text"); - sourceBuilder.query(new BoolQueryBuilder().must(termQueryBuilder).filter(matchQueryBuilder).should(regexpQueryBuilder)); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field1", "value2"); + MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("field2", "php"); + RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field3", "text"); + RangeQueryBuilder rangeQueryBuilder = new RangeQueryBuilder("field4"); + sourceBuilder.query( + new BoolQueryBuilder().must(termQueryBuilder).filter(matchQueryBuilder).should(regexpQueryBuilder).filter(rangeQueryBuilder) + ); - String shape = QueryShapeGenerator.buildShape(sourceBuilder, true); + 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 expected = "bool\n" + " must:\n" + " term\n" + " filter:\n" + " match\n" + " should:\n" + " regexp\n"; - assertEquals(expected, shape); + 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() { @@ -46,16 +67,25 @@ public void testAggregationShape() { .subAggregation(new TermsAggregationBuilder("child-agg2").userValueTypeHint(ValueType.STRING).field("key.sub2")) ); - String shape = QueryShapeGenerator.buildShape(sourceBuilder, true); - - String expected = "aggregation:\n" + String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String expectedShowFieldsTrue = "aggregation:\n" + " terms [key]\n" + " aggregation:\n" + " terms [key.sub1]\n" + " terms [key.sub2]\n" + " terms [model]\n" + " terms [type]\n"; - assertEquals(expected, shape); + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); + + String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String expectedShowFieldsFalse = "aggregation:\n" + + " terms\n" + + " terms\n" + + " terms\n" + + " aggregation:\n" + + " terms\n" + + " terms\n"; + assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse); } public void testSortShape() { @@ -65,9 +95,12 @@ public void testSortShape() { sourceBuilder.sort("price", SortOrder.ASC); sourceBuilder.sort("album", SortOrder.ASC); - String shape = QueryShapeGenerator.buildShape(sourceBuilder, true); + String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true); + String expectedShowFieldsTrue = "sort:\n" + " asc [album]\n" + " asc [price]\n" + " desc [color]\n" + " desc [vendor]\n"; + assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue); - String expected = "sort:\n" + " asc [album]\n" + " asc [price]\n" + " desc [color]\n" + " desc [vendor]\n"; - assertEquals(expected, shape); + String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false); + String expectedShowFieldsFalse = "sort:\n" + " asc\n" + " asc\n" + " desc\n" + " desc\n"; + assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse); } }