diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java b/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java index a36c794a779ed..44585dc0cfd77 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java +++ b/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java @@ -29,7 +29,18 @@ import org.opensearch.search.aggregations.bucket.terms.SignificantTermsAggregationBuilder; import org.opensearch.search.aggregations.bucket.terms.SignificantTextAggregationBuilder; import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder; -import org.opensearch.search.aggregations.metrics.*; +import org.opensearch.search.aggregations.metrics.AvgAggregationBuilder; +import org.opensearch.search.aggregations.metrics.CardinalityAggregationBuilder; +import org.opensearch.search.aggregations.metrics.ExtendedStatsAggregationBuilder; +import org.opensearch.search.aggregations.metrics.GeoCentroidAggregationBuilder; +import org.opensearch.search.aggregations.metrics.MaxAggregationBuilder; +import org.opensearch.search.aggregations.metrics.MedianAbsoluteDeviationAggregationBuilder; +import org.opensearch.search.aggregations.metrics.MinAggregationBuilder; +import org.opensearch.search.aggregations.metrics.PercentileRanksAggregationBuilder; +import org.opensearch.search.aggregations.metrics.PercentilesAggregationBuilder; +import org.opensearch.search.aggregations.metrics.ScriptedMetricAggregationBuilder; +import org.opensearch.search.aggregations.metrics.SumAggregationBuilder; +import org.opensearch.search.aggregations.metrics.TopHitsAggregationBuilder; import org.opensearch.telemetry.metrics.tags.Tags; import java.util.Collection; diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java index 134575c599581..bb3184aea9034 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java +++ b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java @@ -9,8 +9,35 @@ package org.opensearch.action.search; import org.apache.lucene.search.BooleanClause; -import org.opensearch.index.query.*; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.BoostingQueryBuilder; +import org.opensearch.index.query.ConstantScoreQueryBuilder; +import org.opensearch.index.query.DisMaxQueryBuilder; +import org.opensearch.index.query.DistanceFeatureQueryBuilder; +import org.opensearch.index.query.ExistsQueryBuilder; +import org.opensearch.index.query.FieldMaskingSpanQueryBuilder; +import org.opensearch.index.query.FuzzyQueryBuilder; +import org.opensearch.index.query.GeoBoundingBoxQueryBuilder; +import org.opensearch.index.query.GeoDistanceQueryBuilder; +import org.opensearch.index.query.GeoPolygonQueryBuilder; +import org.opensearch.index.query.GeoShapeQueryBuilder; +import org.opensearch.index.query.IntervalQueryBuilder; +import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.MatchPhraseQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.index.query.MultiMatchQueryBuilder; +import org.opensearch.index.query.PrefixQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; +import org.opensearch.index.query.QueryStringQueryBuilder; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.index.query.RegexpQueryBuilder; +import org.opensearch.index.query.ScriptQueryBuilder; +import org.opensearch.index.query.SimpleQueryStringBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.WildcardQueryBuilder; import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder; +import org.opensearch.telemetry.metrics.Counter; import org.opensearch.telemetry.metrics.tags.Tags; import java.util.HashMap; @@ -25,7 +52,7 @@ final class SearchQueryCategorizingVisitor implements QueryBuilderVisitor { private static final String LEVEL_TAG = "level"; private final int level; private final SearchQueryCounters searchQueryCounters; - private final Map, Consumer> queryHandlers; + private final Map, Counter> queryHandlers; public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) { this(searchQueryCounters, 0); @@ -37,152 +64,49 @@ private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level) this.queryHandlers = initializeQueryHandlers(); } - private Map, Consumer> initializeQueryHandlers() { - Map, Consumer> handlers = new HashMap<>(); - - handlers.put(BoolQueryBuilder.class, this::handleBoolQuery); - handlers.put(FunctionScoreQueryBuilder.class, this::handleFunctionScoreQuery); - handlers.put(MatchQueryBuilder.class, this::handleMatchQuery); - handlers.put(MatchPhraseQueryBuilder.class, this::handleMatchPhrasePrefixQuery); - handlers.put(MultiMatchQueryBuilder.class, this::handleMultiMatchQuery); - handlers.put(QueryStringQueryBuilder.class, this::handleQueryStringQuery); - handlers.put(RangeQueryBuilder.class, this::handleRangeQuery); - handlers.put(RegexpQueryBuilder.class, this::handleRegexpQuery); - handlers.put(TermQueryBuilder.class, this::handleTermQuery); - handlers.put(WildcardQueryBuilder.class, this::handleWildcardQuery); - handlers.put(BoostingQueryBuilder.class, this::handleBoostingQuery); - handlers.put(ConstantScoreQueryBuilder.class, this::handleConstantScoreQuery); - handlers.put(DisMaxQueryBuilder.class, this::handleDisMaxQuery); - handlers.put(DistanceFeatureQueryBuilder.class, this::handleDistanceFeatureQuery); - handlers.put(ExistsQueryBuilder.class, this::handleExistsQuery); - handlers.put(FieldMaskingSpanQueryBuilder.class, this::handleFieldMaskingSpanQuery); - handlers.put(FuzzyQueryBuilder.class, this::handleFuzzyQuery); - handlers.put(GeoBoundingBoxQueryBuilder.class, this::handleGeoBoundingBoxQuery); - handlers.put(GeoDistanceQueryBuilder.class, this::handleGeoDistanceQuery); - handlers.put(GeoPolygonQueryBuilder.class, this::handleGeoPolygonQuery); - handlers.put(GeoShapeQueryBuilder.class, this::handleGeoShapeQuery); - handlers.put(IntervalQueryBuilder.class, this::handleIntervalQuery); - handlers.put(MatchAllQueryBuilder.class, this::handleMatchAllQuery); - handlers.put(PrefixQueryBuilder.class, this::handlePrefixQuery); - handlers.put(ScriptQueryBuilder.class, this::handleScriptQuery); - handlers.put(SimpleQueryStringBuilder.class, this::handleSimpleQueryStringQuery); + private Map, Counter> initializeQueryHandlers() { + Map, Counter> handlers = new HashMap<>(); + + handlers.put(BoolQueryBuilder.class, searchQueryCounters.boolCounter); + handlers.put(FunctionScoreQueryBuilder.class, searchQueryCounters.functionScoreCounter); + handlers.put(MatchQueryBuilder.class, searchQueryCounters.matchCounter); + handlers.put(MatchPhraseQueryBuilder.class, searchQueryCounters.matchPhrasePrefixCounter); + handlers.put(MultiMatchQueryBuilder.class, searchQueryCounters.multiMatchCounter); + handlers.put(QueryStringQueryBuilder.class, searchQueryCounters.queryStringQueryCounter); + handlers.put(RangeQueryBuilder.class, searchQueryCounters.rangeCounter); + handlers.put(RegexpQueryBuilder.class, searchQueryCounters.regexCounter); + handlers.put(TermQueryBuilder.class, searchQueryCounters.termCounter); + handlers.put(WildcardQueryBuilder.class, searchQueryCounters.wildcardCounter); + handlers.put(BoostingQueryBuilder.class, searchQueryCounters.boostCounter); + handlers.put(ConstantScoreQueryBuilder.class, searchQueryCounters.constantScoreCounter); + handlers.put(DisMaxQueryBuilder.class, searchQueryCounters.disMaxCounter); + handlers.put(DistanceFeatureQueryBuilder.class, searchQueryCounters.distanceFeatureCounter); + handlers.put(ExistsQueryBuilder.class, searchQueryCounters.existsCounter); + handlers.put(FieldMaskingSpanQueryBuilder.class, searchQueryCounters.fieldMaskingSpanCounter); + handlers.put(FuzzyQueryBuilder.class, searchQueryCounters.fuzzyCounter); + handlers.put(GeoBoundingBoxQueryBuilder.class, searchQueryCounters.geoBoundingBoxCounter); + handlers.put(GeoDistanceQueryBuilder.class, searchQueryCounters.geoDistanceCounter); + handlers.put(GeoPolygonQueryBuilder.class, searchQueryCounters.geoPolygonCounter); + handlers.put(GeoShapeQueryBuilder.class, searchQueryCounters.geoShapeCounter); + handlers.put(IntervalQueryBuilder.class, searchQueryCounters.intervalCounter); + handlers.put(MatchAllQueryBuilder.class, searchQueryCounters.matchallCounter); + handlers.put(PrefixQueryBuilder.class, searchQueryCounters.prefixCounter); + handlers.put(ScriptQueryBuilder.class, searchQueryCounters.scriptCounter); + handlers.put(SimpleQueryStringBuilder.class, searchQueryCounters.simpleQueryStringCounter); return handlers; } public void accept(QueryBuilder qb) { - queryHandlers.getOrDefault(qb.getClass(), this::handleOtherQuery).accept(qb); + Counter counter = queryHandlers.get(qb.getClass()); + if (counter != null) { + counter.add(1, Tags.create().addTag(LEVEL_TAG, level)); + } else { + searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); + } } public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { return new SearchQueryCategorizingVisitor(searchQueryCounters, level + 1); } - - private void handleBoolQuery(QueryBuilder qb) { - searchQueryCounters.boolCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleFunctionScoreQuery(QueryBuilder qb) { - searchQueryCounters.functionScoreCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleMatchQuery(QueryBuilder qb) { - searchQueryCounters.matchCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleMatchPhrasePrefixQuery(QueryBuilder qb) { - searchQueryCounters.matchPhrasePrefixCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleMultiMatchQuery(QueryBuilder qb) { - searchQueryCounters.multiMatchCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleQueryStringQuery(QueryBuilder qb) { - searchQueryCounters.queryStringQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleRangeQuery(QueryBuilder qb) { - searchQueryCounters.rangeCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleRegexpQuery(QueryBuilder qb) { - searchQueryCounters.regexCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleTermQuery(QueryBuilder qb) { - searchQueryCounters.termCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleWildcardQuery(QueryBuilder qb) { - searchQueryCounters.wildcardCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleBoostingQuery(QueryBuilder qb) { - searchQueryCounters.boostCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleConstantScoreQuery(QueryBuilder qb) { - searchQueryCounters.constantScoreCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleDisMaxQuery(QueryBuilder qb) { - searchQueryCounters.disMaxCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleDistanceFeatureQuery(QueryBuilder qb) { - searchQueryCounters.distanceFeatureCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleExistsQuery(QueryBuilder qb) { - searchQueryCounters.existsCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleFieldMaskingSpanQuery(QueryBuilder qb) { - searchQueryCounters.fieldMaskingSpanCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleFuzzyQuery(QueryBuilder qb) { - searchQueryCounters.fuzzyCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleGeoBoundingBoxQuery(QueryBuilder qb) { - searchQueryCounters.geoBoundingBoxCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleGeoDistanceQuery(QueryBuilder qb) { - searchQueryCounters.geoDistanceCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleGeoPolygonQuery(QueryBuilder qb) { - searchQueryCounters.geoPolygonCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleGeoShapeQuery(QueryBuilder qb) { - searchQueryCounters.geoShapeCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleIntervalQuery(QueryBuilder qb) { - searchQueryCounters.intervalCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleMatchAllQuery(QueryBuilder qb) { - searchQueryCounters.matchallCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handlePrefixQuery(QueryBuilder qb) { - searchQueryCounters.prefixCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleScriptQuery(QueryBuilder qb) { - searchQueryCounters.scriptCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleSimpleQueryStringQuery(QueryBuilder qb) { - searchQueryCounters.simpleQueryStringCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } - - private void handleOtherQuery(QueryBuilder qb) { - searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } }