Capture query categorization metrics for additional query types and a…

…ggregations types (#11582)


Signed-off-by: Siddhant Deshmukh <[email protected]>

CHANGELOG.md

@@ -185,6 +185,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Automatically add scheme to discovery.ec2.endpoint ([#11512](https://github.com/opensearch-project/OpenSearch/pull/11512))
 - Restore support for Java 8 for RestClient ([#11562](https://github.com/opensearch-project/OpenSearch/pull/11562))
 - Add deleted doc count in _cat/shards ([#11678](https://github.com/opensearch-project/OpenSearch/pull/11678))
+- Capture information for additional query types and aggregation types ([#11582](https://github.com/opensearch-project/OpenSearch/pull/11582))
 
 ### Deprecated
 

server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java

@@ -0,0 +1,55 @@
+/*
+ * 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.
+ */
+
+package org.opensearch.action.search;
+
+import org.opensearch.search.aggregations.AggregationBuilder;
+import org.opensearch.search.aggregations.PipelineAggregationBuilder;
+import org.opensearch.telemetry.metrics.tags.Tags;
+
+import java.util.Collection;
+
+/**
+ * Increments the counters related to Aggregation Search Queries.
+ */
+public class SearchQueryAggregationCategorizer {
+
+    private static final String TYPE_TAG = "type";
+    private final SearchQueryCounters searchQueryCounters;
+
+    public SearchQueryAggregationCategorizer(SearchQueryCounters searchQueryCounters) {
+        this.searchQueryCounters = searchQueryCounters;
+    }
+
+    public void incrementSearchQueryAggregationCounters(Collection<AggregationBuilder> aggregatorFactories) {
+        for (AggregationBuilder aggregationBuilder : aggregatorFactories) {
+            incrementCountersRecursively(aggregationBuilder);
+        }
+    }
+
+    private void incrementCountersRecursively(AggregationBuilder aggregationBuilder) {
+        // Increment counters for the current aggregation
+        String aggregationType = aggregationBuilder.getType();
+        searchQueryCounters.aggCounter.add(1, Tags.create().addTag(TYPE_TAG, aggregationType));
+
+        // Recursively process sub-aggregations if any
+        Collection<AggregationBuilder> subAggregations = aggregationBuilder.getSubAggregations();
+        if (subAggregations != null && !subAggregations.isEmpty()) {
+            for (AggregationBuilder subAggregation : subAggregations) {
+                incrementCountersRecursively(subAggregation);
+            }
+        }
+
+        // Process pipeline aggregations
+        Collection<PipelineAggregationBuilder> pipelineAggregations = aggregationBuilder.getPipelineAggregations();
+        for (PipelineAggregationBuilder pipelineAggregation : pipelineAggregations) {
+            String pipelineAggregationType = pipelineAggregation.getType();
+            searchQueryCounters.aggCounter.add(1, Tags.create().addTag(TYPE_TAG, pipelineAggregationType));
+        }
+    }
+}

server/src/main/java/org/opensearch/action/search/SearchQueryCategorizer.java

@@ -32,13 +32,15 @@ final class SearchQueryCategorizer {
 
     final SearchQueryCounters searchQueryCounters;
 
+    final SearchQueryAggregationCategorizer searchQueryAggregationCategorizer;
+
     public SearchQueryCategorizer(MetricsRegistry metricsRegistry) {
         searchQueryCounters = new SearchQueryCounters(metricsRegistry);
+        searchQueryAggregationCategorizer = new SearchQueryAggregationCategorizer(searchQueryCounters);
     }
 
     public void categorize(SearchSourceBuilder source) {
         QueryBuilder topLevelQueryBuilder = source.query();
-
         logQueryShape(topLevelQueryBuilder);
         incrementQueryTypeCounters(topLevelQueryBuilder);
         incrementQueryAggregationCounters(source.aggregations());
@@ -56,9 +58,11 @@ private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
     }
 
     private void incrementQueryAggregationCounters(AggregatorFactories.Builder aggregations) {
-        if (aggregations != null) {
-            searchQueryCounters.aggCounter.add(1);
+        if (aggregations == null) {
+            return;
         }
+
+        searchQueryAggregationCategorizer.incrementSearchQueryAggregationCounters(aggregations.getAggregatorFactories());
     }
 
     private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) {

server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java

@@ -9,26 +9,14 @@
 package org.opensearch.action.search;
 
 import org.apache.lucene.search.BooleanClause;
-import org.opensearch.index.query.BoolQueryBuilder;
-import org.opensearch.index.query.MatchPhraseQueryBuilder;
-import org.opensearch.index.query.MatchQueryBuilder;
-import org.opensearch.index.query.MultiMatchQueryBuilder;
 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.TermQueryBuilder;
-import org.opensearch.index.query.WildcardQueryBuilder;
-import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
-import org.opensearch.telemetry.metrics.tags.Tags;
 
 /**
- * Class to visit the querybuilder tree and also track the level information.
+ * Class to visit the query builder tree and also track the level information.
  * Increments the counters related to Search Query type.
  */
 final class SearchQueryCategorizingVisitor implements QueryBuilderVisitor {
-    private static final String LEVEL_TAG = "level";
     private final int level;
     private final SearchQueryCounters searchQueryCounters;
 
@@ -42,29 +30,7 @@ private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level)
     }
 
     public void accept(QueryBuilder qb) {
-        if (qb instanceof BoolQueryBuilder) {
-            searchQueryCounters.boolCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof FunctionScoreQueryBuilder) {
-            searchQueryCounters.functionScoreCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof MatchQueryBuilder) {
-            searchQueryCounters.matchCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof MatchPhraseQueryBuilder) {
-            searchQueryCounters.matchPhrasePrefixCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof MultiMatchQueryBuilder) {
-            searchQueryCounters.multiMatchCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof QueryStringQueryBuilder) {
-            searchQueryCounters.queryStringQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof RangeQueryBuilder) {
-            searchQueryCounters.rangeCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof RegexpQueryBuilder) {
-            searchQueryCounters.regexCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof TermQueryBuilder) {
-            searchQueryCounters.termCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else if (qb instanceof WildcardQueryBuilder) {
-            searchQueryCounters.wildcardCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else {
-            searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        }
+        searchQueryCounters.incrementCounter(qb, level);
     }
 
     public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {