Categorize search queries by type and log query shape

CHANGELOG.md

@@ -91,6 +91,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 - Per request phase latency ([#10351](https://github.com/opensearch-project/OpenSearch/issues/10351))
 - [Remote Store] Add repository stats for remote store([#10567](https://github.com/opensearch-project/OpenSearch/pull/10567))
+- Add search query categorizer ([#10255](https://github.com/opensearch-project/OpenSearch/pull/10255))
 
 ### Dependencies
 - Bump `com.google.api.grpc:proto-google-common-protos` from 2.10.0 to 2.25.1 ([#10208](https://github.com/opensearch-project/OpenSearch/pull/10208), [#10298](https://github.com/opensearch-project/OpenSearch/pull/10298))

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

@@ -0,0 +1,81 @@
+/*
+ * 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.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
+import org.opensearch.index.query.QueryShapeVisitor;
+import org.opensearch.search.aggregations.AggregatorFactories;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.sort.SortBuilder;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+import org.opensearch.telemetry.metrics.tags.Tags;
+
+import java.util.List;
+import java.util.ListIterator;
+
+/**
+ * Class to categorize the search queries based on the type and increment the relevant counters.
+ * Class also logs the query shape.
+ */
+final class SearchQueryCategorizer {
+
+    private static final Logger log = LogManager.getLogger(SearchQueryCategorizer.class);
+
+    final SearchQueryCounters searchQueryCounters;
+
+    public SearchQueryCategorizer(MetricsRegistry metricsRegistry) {
+        searchQueryCounters = new SearchQueryCounters(metricsRegistry);
+    }
+
+    public void categorize(SearchSourceBuilder source) {
+        QueryBuilder topLevelQueryBuilder = source.query();
+
+        logQueryShape(topLevelQueryBuilder);
+        incrementQueryTypeCounters(topLevelQueryBuilder);
+        incrementQueryAggregationCounters(source.aggregations());
+        incrementQuerySortCounters(source.sorts());
+    }
+
+    private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
+        if (sorts != null && sorts.size() > 0) {
+            for (ListIterator<SortBuilder<?>> it = sorts.listIterator(); it.hasNext();) {
+                SortBuilder sortBuilder = it.next();
+                String sortOrder = sortBuilder.order().toString();
+                searchQueryCounters.sortCounter.add(1, Tags.create().addTag("sort_order", sortOrder));
+            }
+        }
+    }
+
+    private void incrementQueryAggregationCounters(AggregatorFactories.Builder aggregations) {
+        if (aggregations != null) {
+            searchQueryCounters.aggCounter.add(1);
+        }
+    }
+
+    private void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) {
+        if (topLevelQueryBuilder == null) {
+            return;
+        }
+        QueryBuilderVisitor searchQueryVisitor = new SearchQueryCategorizingVisitor(searchQueryCounters);
+        topLevelQueryBuilder.visit(searchQueryVisitor);
+    }
+
+    private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
+        if (topLevelQueryBuilder == null) {
+            return;
+        }
+        QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
+        topLevelQueryBuilder.visit(shapeVisitor);
+        log.debug("Query shape : {}", shapeVisitor.prettyPrintTree("  "));
+    }
+
+}

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

@@ -0,0 +1,73 @@
+/*
+ * 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.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.
+ * 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;
+
+    public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) {
+        this(searchQueryCounters, 0);
+    }
+
+    private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level) {
+        this.searchQueryCounters = counters;
+        this.level = 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));
+        }
+    }
+
+    public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+        return new SearchQueryCategorizingVisitor(searchQueryCounters, level + 1);
+    }
+}

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

@@ -0,0 +1,117 @@
+/*
+ * 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.telemetry.metrics.Counter;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+
+/**
+ * Class contains all the Counters related to search query types.
+ */
+final class SearchQueryCounters {
+    private static final String UNIT = "1";
+    private final MetricsRegistry metricsRegistry;
+
+    // Counters related to Query types
+    public final Counter aggCounter;
+    public final Counter boolCounter;
+    public final Counter functionScoreCounter;
+    public final Counter matchCounter;
+    public final Counter matchPhrasePrefixCounter;
+    public final Counter multiMatchCounter;
+    public final Counter otherQueryCounter;
+    public final Counter queryStringQueryCounter;
+    public final Counter rangeCounter;
+    public final Counter regexCounter;
+
+    public final Counter sortCounter;
+    public final Counter skippedCounter;
+    public final Counter termCounter;
+    public final Counter totalCounter;
+    public final Counter wildcardCounter;
+
+    public SearchQueryCounters(MetricsRegistry metricsRegistry) {
+        this.metricsRegistry = metricsRegistry;
+        this.aggCounter = metricsRegistry.createCounter(
+            "search.query.type.agg.count",
+            "Counter for the number of top level agg search queries",
+            UNIT
+        );
+        this.boolCounter = metricsRegistry.createCounter(
+            "search.query.type.bool.count",
+            "Counter for the number of top level and nested bool search queries",
+            UNIT
+        );
+        this.functionScoreCounter = metricsRegistry.createCounter(
+            "search.query.type.functionscore.count",
+            "Counter for the number of top level and nested function score search queries",
+            UNIT
+        );
+        this.matchCounter = metricsRegistry.createCounter(
+            "search.query.type.match.count",
+            "Counter for the number of top level and nested match search queries",
+            UNIT
+        );
+        this.matchPhrasePrefixCounter = metricsRegistry.createCounter(
+            "search.query.type.matchphrase.count",
+            "Counter for the number of top level and nested match phrase prefix search queries",
+            UNIT
+        );
+        this.multiMatchCounter = metricsRegistry.createCounter(
+            "search.query.type.multimatch.count",
+            "Counter for the number of top level and nested multi match search queries",
+            UNIT
+        );
+        this.otherQueryCounter = metricsRegistry.createCounter(
+            "search.query.type.other.count",
+            "Counter for the number of top level and nested search queries that do not match any other categories",
+            UNIT
+        );
+        this.queryStringQueryCounter = metricsRegistry.createCounter(
+            "search.query.type.querystringquery.count",
+            "Counter for the number of top level and nested queryStringQuery search queries",
+            UNIT
+        );
+        this.rangeCounter = metricsRegistry.createCounter(
+            "search.query.type.range.count",
+            "Counter for the number of top level and nested range search queries",
+            UNIT
+        );
+        this.regexCounter = metricsRegistry.createCounter(
+            "search.query.type.regex.count",
+            "Counter for the number of top level and nested regex search queries",
+            UNIT
+        );
+        this.skippedCounter = metricsRegistry.createCounter(
+            "search.query.type.skipped.count",
+            "Counter for the number queries skipped due to error",
+            UNIT
+        );
+        this.sortCounter = metricsRegistry.createCounter(
+            "search.query.type.sort.count",
+            "Counter for the number of top level sort search queries",
+            UNIT
+        );
+        this.termCounter = metricsRegistry.createCounter(
+            "search.query.type.term.count",
+            "Counter for the number of top level and nested term search queries",
+            UNIT
+        );
+        this.totalCounter = metricsRegistry.createCounter(
+            "search.query.type.total.count",
+            "Counter for the number of top level and nested search queries",
+            UNIT
+        );
+        this.wildcardCounter = metricsRegistry.createCounter(
+            "search.query.type.wildcard.count",
+            "Counter for the number of top level and nested wildcard search queries",
+            UNIT
+        );
+    }
+}

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

@@ -88,6 +88,7 @@
 import org.opensearch.search.profile.SearchProfileShardResults;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.Task;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.RemoteClusterAware;
 import org.opensearch.transport.RemoteClusterService;
@@ -137,6 +138,13 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
         Property.NodeScope
     );
 
+    public static final Setting<Boolean> SEARCH_QUERY_METRICS_ENABLED_SETTING = Setting.boolSetting(
+        "search.query.metrics.enabled",
+        false,
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
+
     // cluster level setting for timeout based search cancellation. If search request level parameter is present then that will take
     // precedence over the cluster setting value
     public static final String SEARCH_CANCEL_AFTER_TIME_INTERVAL_SETTING_KEY = "search.cancel_after_time_interval";
@@ -177,8 +185,14 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
 
     private volatile boolean isRequestStatsEnabled;
 
+    private volatile boolean searchQueryMetricsEnabled;
+
     private final SearchRequestStats searchRequestStats;
 
+    private final MetricsRegistry metricsRegistry;
+
+    private SearchQueryCategorizer searchQueryCategorizer;
+
     @Inject
     public TransportSearchAction(
         NodeClient client,
@@ -193,7 +207,8 @@ public TransportSearchAction(
         IndexNameExpressionResolver indexNameExpressionResolver,
         NamedWriteableRegistry namedWriteableRegistry,
         SearchPipelineService searchPipelineService,
-        SearchRequestStats searchRequestStats
+        SearchRequestStats searchRequestStats,
+        MetricsRegistry metricsRegistry
     ) {
         super(SearchAction.NAME, transportService, actionFilters, (Writeable.Reader<SearchRequest>) SearchRequest::new);
         this.client = client;
@@ -211,6 +226,17 @@ public TransportSearchAction(
         this.isRequestStatsEnabled = clusterService.getClusterSettings().get(SEARCH_REQUEST_STATS_ENABLED);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(SEARCH_REQUEST_STATS_ENABLED, this::setIsRequestStatsEnabled);
         this.searchRequestStats = searchRequestStats;
+        this.metricsRegistry = metricsRegistry;
+        this.searchQueryMetricsEnabled = clusterService.getClusterSettings().get(SEARCH_QUERY_METRICS_ENABLED_SETTING);
+        clusterService.getClusterSettings()
+            .addSettingsUpdateConsumer(SEARCH_QUERY_METRICS_ENABLED_SETTING, this::setSearchQueryMetricsEnabled);
+    }
+
+    private void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) {
+        this.searchQueryMetricsEnabled = searchQueryMetricsEnabled;
+        if ((this.searchQueryMetricsEnabled == true) && this.searchQueryCategorizer == null) {
+            this.searchQueryCategorizer = new SearchQueryCategorizer(metricsRegistry);
+        }
     }
 
     private void setIsRequestStatsEnabled(boolean isRequestStatsEnabled) {
@@ -489,6 +515,14 @@ private void executeRequest(
             return;
         }
 
+        if (searchQueryMetricsEnabled) {
+            try {
+                searchQueryCategorizer.categorize(searchRequest.source());
+            } catch (Exception e) {
+                logger.error("Error while trying to categorize the query.", e);
+            }
+        }
+
         ActionListener<SearchSourceBuilder> rewriteListener = ActionListener.wrap(source -> {
             if (source != searchRequest.source()) {
                 // only set it if it changed - we don't allow null values to be set but it might be already null. this way we catch

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -377,6 +377,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 TransportSearchAction.SEARCH_CANCEL_AFTER_TIME_INTERVAL_SETTING,
                 TransportSearchAction.SEARCH_REQUEST_STATS_ENABLED,
                 TransportSearchAction.SEARCH_PHASE_TOOK_ENABLED,
+                TransportSearchAction.SEARCH_QUERY_METRICS_ENABLED_SETTING,
                 RemoteClusterService.REMOTE_CLUSTER_SKIP_UNAVAILABLE,
                 SniffConnectionStrategy.REMOTE_CONNECTIONS_PER_CLUSTER,
                 RemoteClusterService.REMOTE_INITIAL_CONNECTION_TIMEOUT_SETTING,