Search Query Categorizor initial skeleton using QueryBuilderVisitor

deshsidd · Sep 27, 2023 · d1a1f2c · d1a1f2c
1 parent 8807d7a
commit d1a1f2c
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java
@@ -0,0 +1,49 @@
+/*
+ * 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.MatchQueryBuilder;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
+import org.opensearch.index.query.QueryStringQueryBuilder;
+import org.opensearch.search.builder.SearchSourceBuilder;
+
+public class SearchQueryCategorizor {
+
+    public static void categorize(SearchSourceBuilder source) {
+        QueryBuilder topLevelQueryBuilder = source.query();
+        QueryBuilderVisitor queryBuilderVisitor = new QueryBuilderVisitor() {
+            @Override
+            public void accept(QueryBuilder qb) {
+                // This method will be called for every QueryBuilder node in the tree.
+                // The tree referred to here is the tree of querybuilders for the incoming search
+                // query with the topLevelQueryBuilder as the root.
+
+                // Increment counter for current QueryBuilder using Metric Framework.
+                if (qb instanceof BoolQueryBuilder) {
+                    // Increment counter for Bool using Metric Framework.
+                } else if (qb instanceof MatchQueryBuilder) {
+                    // Increment counter for Match using Metric Framework.
+                } else if (qb instanceof QueryStringQueryBuilder) {
+                    // Increment counter for QueryStringQuery using Metric Framework.
+                }
+                // Similar for other builders
+            }
+
+            @Override
+            public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+                return this;
+            }
+        };
+        topLevelQueryBuilder.visit(queryBuilderVisitor);
+    }
+
+}
diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -32,6 +32,7 @@
 
 package org.opensearch.action.search;
 
+import org.apache.lucene.search.BooleanClause;
 import org.opensearch.action.OriginalIndices;
 import org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequest;
 import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsGroup;
@@ -72,6 +73,8 @@
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.core.indices.breaker.CircuitBreakerService;
 import org.opensearch.core.tasks.TaskId;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
 import org.opensearch.index.query.Rewriteable;
 import org.opensearch.search.SearchPhaseResult;
 import org.opensearch.search.SearchService;
@@ -432,6 +435,8 @@ private void executeRequest(
             return;
         }
 
+        SearchQueryCategorizor.categorize(searchRequest.source());
+
         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