Address review comments and add complex query unit test

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>

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

@@ -24,7 +24,7 @@
  * Class to categorize the search queries based on the type and increment the relevant counters.
  * Class also logs the query shape.
  */
-public class SearchQueryCategorizer {
+public final class SearchQueryCategorizer {
 
     private static final Logger log = LogManager.getLogger(SearchQueryCategorizer.class);
 

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

@@ -28,7 +28,7 @@
  * Increments the counters related to Search Query type.
  */
 public class SearchQueryCategorizingVisitor implements QueryBuilderVisitor {
-    public static final String LEVEL_TAG = "level";
+    private static final String LEVEL_TAG = "level";
     private final int level;
     private final SearchQueryCounters searchQueryCounters;
 

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

@@ -40,77 +40,77 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) {
         this.aggCounter = metricsRegistry.createCounter(
             "search.query.type.agg.count",
             "Counter for the number of top level agg search queries",
-            "0"
+            "1"
         );
         this.boolCounter = metricsRegistry.createCounter(
             "search.query.type.bool.count",
             "Counter for the number of top level and nested bool search queries",
-            "0"
+            "1"
         );
         this.functionScoreCounter = metricsRegistry.createCounter(
             "search.query.type.functionscore.count",
             "Counter for the number of top level and nested function score search queries",
-            "0"
+            "1"
         );
         this.matchCounter = metricsRegistry.createCounter(
             "search.query.type.match.count",
             "Counter for the number of top level and nested match search queries",
-            "0"
+            "1"
         );
         this.matchPhrasePrefixCounter = metricsRegistry.createCounter(
             "search.query.type.matchphrase.count",
             "Counter for the number of top level and nested match phrase prefix search queries",
-            "0"
+            "1"
         );
         this.multiMatchCounter = metricsRegistry.createCounter(
             "search.query.type.multimatch.count",
             "Counter for the number of top level and nested multi match search queries",
-            "0"
+            "1"
         );
         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",
-            "0"
+            "1"
         );
         this.queryStringQueryCounter = metricsRegistry.createCounter(
             "search.query.type.querystringquery.count",
             "Counter for the number of top level and nested queryStringQuery search queries",
-            "0"
+            "1"
         );
         this.rangeCounter = metricsRegistry.createCounter(
             "search.query.type.range.count",
             "Counter for the number of top level and nested range search queries",
-            "0"
+            "1"
         );
         this.regexCounter = metricsRegistry.createCounter(
             "search.query.type.regex.count",
             "Counter for the number of top level and nested regex search queries",
-            "0"
+            "1"
         );
         this.skippedCounter = metricsRegistry.createCounter(
             "search.query.type.skipped.count",
             "Counter for the number queries skipped due to error",
-            "0"
+            "1"
         );
         this.sortCounter = metricsRegistry.createCounter(
             "search.query.type.sort.count",
             "Counter for the number of top level sort search queries",
-            "0"
+            "1"
         );
         this.termCounter = metricsRegistry.createCounter(
             "search.query.type.term.count",
             "Counter for the number of top level and nested term search queries",
-            "0"
+            "1"
         );
         this.totalCounter = metricsRegistry.createCounter(
             "search.query.type.total.count",
             "Counter for the number of top level and nested search queries",
-            "0"
+            "1"
         );
         this.wildcardCounter = metricsRegistry.createCounter(
             "search.query.type.wildcard.count",
             "Counter for the number of top level and nested wildcard search queries",
-            "0"
+            "1"
         );
     }
 }

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

@@ -185,7 +185,7 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
 
     private volatile boolean isRequestStatsEnabled;
 
-    private volatile boolean searchQueryCategorizationEnabled;
+    private volatile boolean searchQueryMetricsEnabled;
 
     private final SearchRequestStats searchRequestStats;
 
@@ -227,14 +227,14 @@ public TransportSearchAction(
         clusterService.getClusterSettings().addSettingsUpdateConsumer(SEARCH_REQUEST_STATS_ENABLED, this::setIsRequestStatsEnabled);
         this.searchRequestStats = searchRequestStats;
         this.metricsRegistry = metricsRegistry;
-        this.searchQueryCategorizationEnabled = clusterService.getClusterSettings().get(SEARCH_QUERY_METRICS_ENABLED_SETTING);
+        this.searchQueryMetricsEnabled = clusterService.getClusterSettings().get(SEARCH_QUERY_METRICS_ENABLED_SETTING);
         clusterService.getClusterSettings()
-            .addSettingsUpdateConsumer(SEARCH_QUERY_METRICS_ENABLED_SETTING, this::setIsSearchQueryCategorizationEnabled);
+            .addSettingsUpdateConsumer(SEARCH_QUERY_METRICS_ENABLED_SETTING, this::setSearchQueryMetricsEnabled);
     }
 
-    private void setIsSearchQueryCategorizationEnabled(boolean isSearchQueryCategorizationEnabled) {
-        this.searchQueryCategorizationEnabled = isSearchQueryCategorizationEnabled;
-        if (this.searchQueryCategorizationEnabled && this.searchQueryCategorizer == null) {
+    private void setSearchQueryMetricsEnabled(boolean searchQueryMetricsEnabled) {
+        this.searchQueryMetricsEnabled = searchQueryMetricsEnabled;
+        if ((this.searchQueryMetricsEnabled == true ) && this.searchQueryCategorizer == null) {
             this.searchQueryCategorizer = new SearchQueryCategorizer(metricsRegistry);
         }
     }
@@ -515,7 +515,7 @@ private void executeRequest(
             return;
         }
 
-        if (searchQueryCategorizationEnabled) {
+        if (searchQueryMetricsEnabled) {
             try {
                 searchQueryCategorizer.categorize(searchRequest.source());
             } catch (Exception e) {

server/src/main/java/org/opensearch/index/query/QueryShapeVisitor.java

@@ -33,7 +33,7 @@ public void accept(QueryBuilder qb) {
     public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
         // Should get called once per Occur value
         if (childVisitors.containsKey(occur)) {
-            throw new IllegalStateException("getChildVisitor already called for " + occur);
+            throw new IllegalStateException("child visitor already called for " + occur);
         }
         final List<QueryShapeVisitor> childVisitorList = new ArrayList<>();
         QueryBuilderVisitor childVisitorWrapper = new QueryBuilderVisitor() {
@@ -55,7 +55,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
         return childVisitorWrapper;
     }
 
-    public String toJson() {
+    String toJson() {
         StringBuilder outputBuilder = new StringBuilder("{\"type\":\"").append(queryType.get()).append("\"");
         for (Map.Entry<BooleanClause.Occur, List<QueryShapeVisitor>> entry : childVisitors.entrySet()) {
             outputBuilder.append(",\"").append(entry.getKey().name().toLowerCase(Locale.ROOT)).append("\"[");

server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java

@@ -20,6 +20,7 @@
 import org.opensearch.index.query.TermQueryBuilder;
 import org.opensearch.index.query.WildcardQueryBuilder;
 import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
+import org.opensearch.search.aggregations.bucket.range.RangeAggregationBuilder;
 import org.opensearch.search.aggregations.bucket.terms.MultiTermsAggregationBuilder;
 import org.opensearch.search.aggregations.support.MultiTermsValuesSourceConfig;
 import org.opensearch.search.builder.SearchSourceBuilder;
@@ -202,4 +203,25 @@ public void testWildcardQuery() {
 
         Mockito.verify(searchQueryCategorizer.searchQueryCounters.wildcardCounter).add(eq(1.0d), any(Tags.class));
     }
+
+    public void testComplexQuery() {
+        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+        sourceBuilder.size(50);
+
+        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field", "value2");
+        MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("tags", "php");
+        RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field", "text");
+        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder().must(termQueryBuilder).filter(matchQueryBuilder).should(regexpQueryBuilder);
+        sourceBuilder.query(boolQueryBuilder);
+        sourceBuilder.aggregation(new RangeAggregationBuilder("agg1").field("num"));
+
+
+        searchQueryCategorizer.categorize(sourceBuilder);
+
+        Mockito.verify(searchQueryCategorizer.searchQueryCounters.termCounter).add(eq(1.0d), any(Tags.class));
+        Mockito.verify(searchQueryCategorizer.searchQueryCounters.matchCounter).add(eq(1.0d), any(Tags.class));
+        Mockito.verify(searchQueryCategorizer.searchQueryCounters.regexCounter).add(eq(1.0d), any(Tags.class));
+        Mockito.verify(searchQueryCategorizer.searchQueryCounters.boolCounter).add(eq(1.0d), any(Tags.class));
+        Mockito.verify(searchQueryCategorizer.searchQueryCounters.aggCounter).add(eq(1.0d));
+    }
 }