Add level attribute to QueryBuilderVisitor and as a tag in Counters

Signed-off-by: Siddhant Deshmukh <[email protected]>
deshsidd · Oct 5, 2023 · dccc631 · dccc631
1 parent 2fdce7c
commit dccc631
Show file tree

Hide file tree

Showing 16 changed files with 68 additions and 65 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
@@ -27,6 +27,7 @@
 import org.opensearch.index.query.WildcardQueryBuilder;
 import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.telemetry.tracing.attributes.Attributes;
 
 public class SearchQueryCategorizor {
 
@@ -39,43 +40,43 @@ public static void categorize(SearchSourceBuilder source) {
 
         // Get and log the query shape
         QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
-        topLevelQueryBuilder.visit(shapeVisitor);
+        topLevelQueryBuilder.visit(shapeVisitor, 0);
         String queryShapeJson = shapeVisitor.prettyPrintTree("  ");
         log.info("Query shape : " + queryShapeJson);
 
         // Increment the query counters using Metric Framework
         QueryBuilderVisitor queryBuilderVisitor = new QueryBuilderVisitor() {
             @Override
-            public void accept(QueryBuilder qb) {
+            public void accept(QueryBuilder qb, int level) {
                 // 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 AggregationQ) {
-                    searchQueryCounters.aggCounter.add(1);
+                    searchQueryCounters.aggCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof BoolQueryBuilder) {
-                    searchQueryCounters.boolCounter.add(1);
+                    searchQueryCounters.boolCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof FunctionScoreQueryBuilder) {
-                    searchQueryCounters.functionScoreCounter.add(1);
+                    searchQueryCounters.functionScoreCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof MatchQueryBuilder) {
-                    searchQueryCounters.matchCounter.add(1);
+                    searchQueryCounters.matchCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof MatchPhraseQueryBuilder) {
-                    searchQueryCounters.matchPhrasePrefixCounter.add(1);
+                    searchQueryCounters.matchPhrasePrefixCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof MultiMatchQueryBuilder) {
-                    searchQueryCounters.multiMatchCounter.add(1);
+                    searchQueryCounters.multiMatchCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof QueryStringQueryBuilder) {
-                    searchQueryCounters.queryStringQueryCounter.add(1);
+                    searchQueryCounters.queryStringQueryCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof RangeQueryBuilder) {
-                    searchQueryCounters.rangeCounter.add(1);
+                    searchQueryCounters.rangeCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof RegexpQueryBuilder) {
-                    searchQueryCounters.regexCounter.add(1);
+                    searchQueryCounters.regexCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof TermQueryBuilder) {
-                    searchQueryCounters.termCounter.add(1);
+                    searchQueryCounters.termCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else if (qb instanceof WildcardQueryBuilder) {
-                    searchQueryCounters.wildcardCounter.add(1);
+                    searchQueryCounters.wildcardCounter.add(1, Attributes.create().addAttribute("level", level));
                 } else {
-                    searchQueryCounters.otherQueryCounter.add(1);
+                    searchQueryCounters.otherQueryCounter.add(1, Attributes.create().addAttribute("level", level));
                 }
             }
 
@@ -84,7 +85,7 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
                 return this;
             }
         };
-        topLevelQueryBuilder.visit(queryBuilderVisitor);
+        topLevelQueryBuilder.visit(queryBuilderVisitor, 0);
     }
 
 }
diff --git a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java
@@ -429,30 +429,30 @@ private static boolean rewriteClauses(
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (mustClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST);
             for (QueryBuilder mustClause : mustClauses) {
-                mustClause.visit(subVisitor);
+                mustClause.visit(subVisitor, level+1);
             }
         }
         if (shouldClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.SHOULD);
             for (QueryBuilder shouldClause : shouldClauses) {
-                shouldClause.visit(subVisitor);
+                shouldClause.visit(subVisitor, level+1);
             }
         }
         if (mustNotClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST_NOT);
             for (QueryBuilder mustNotClause : mustNotClauses) {
-                mustNotClause.visit(subVisitor);
+                mustNotClause.visit(subVisitor, level+1);
             }
         }
         if (filterClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.FILTER);
             for (QueryBuilder filterClause : filterClauses) {
-                filterClause.visit(subVisitor);
+                filterClause.visit(subVisitor, level+1);
             }
         }
 

diff --git a/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java
@@ -255,13 +255,13 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (positiveQuery != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery);
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery, level+1);
         }
         if (negativeQuery != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery);
+            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery, level+1);
         }
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java
@@ -186,9 +186,9 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
-        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
+        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder, level+1);
     }
 
 }
diff --git a/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java
@@ -249,12 +249,12 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (queries.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
             for (QueryBuilder subQb : queries) {
-                subVisitor.accept(subQb);
+                subVisitor.accept(subQb, level+1);
             }
         }
     }

diff --git a/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java
@@ -210,8 +210,8 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder, level+1);
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilder.java b/server/src/main/java/org/opensearch/index/query/QueryBuilder.java
@@ -99,9 +99,10 @@ default QueryBuilder rewrite(QueryRewriteContext queryShardContext) throws IOExc
     /**
      * Recurse through the QueryBuilder tree, visiting any child QueryBuilder.
      * @param visitor a query builder visitor to be called by each query builder in the tree.
+     * @param level level of the current query builder based on the QueryBuilder tree. 0 is top/root level.
      */
-    default void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    default void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
     };
 
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java b/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java
@@ -18,8 +18,9 @@ public interface QueryBuilderVisitor {
     /**
      * Accept method is called when the visitor accepts the queryBuilder object to be traversed in the query tree.
      * @param qb is a queryBuilder object which is accepeted by the visitor.
+     * @param level level of the current query builder based on the QueryBuilder tree. 0 is top/root level.
      */
-    void accept(QueryBuilder qb);
+    void accept(QueryBuilder qb, int level);
 
     /**
      * Fetches the child sub visitor from the main QueryBuilderVisitor Object.
@@ -35,7 +36,7 @@ public interface QueryBuilderVisitor {
      */
     QueryBuilderVisitor NO_OP_VISITOR = new QueryBuilderVisitor() {
         @Override
-        public void accept(QueryBuilder qb) {
+        public void accept(QueryBuilder qb, int level) {
             // Do nothing
         }
 

diff --git a/server/src/main/java/org/opensearch/index/query/QueryShapeVisitor.java b/server/src/main/java/org/opensearch/index/query/QueryShapeVisitor.java
@@ -22,7 +22,7 @@ public class QueryShapeVisitor implements QueryBuilderVisitor {
     private final Map<BooleanClause.Occur, List<QueryShapeVisitor>> childVisitors = new EnumMap<>(BooleanClause.Occur.class);
 
     @Override
-    public void accept(QueryBuilder qb) {
+    public void accept(QueryBuilder qb, int level) {
         queryType.set(qb.getName());
     }
 
@@ -36,10 +36,10 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
         QueryBuilderVisitor childVisitorWrapper = new QueryBuilderVisitor() {
             QueryShapeVisitor currentChild;
             @Override
-            public void accept(QueryBuilder qb) {
+            public void accept(QueryBuilder qb, int level) {
                 currentChild = new QueryShapeVisitor();
                 childVisitorList.add(currentChild);
-                currentChild.accept(qb);
+                currentChild.accept(qb, level+1);
             }
 
             @Override

diff --git a/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java
@@ -191,9 +191,9 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big, level+1);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little, level+1);
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java
@@ -189,8 +189,8 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(matchBuilder);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(matchBuilder, level+1);
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java
@@ -216,10 +216,10 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (multiTermQueryBuilder != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(multiTermQueryBuilder);
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(multiTermQueryBuilder, level+1);
         }
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java
@@ -301,12 +301,12 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (this.clauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.MUST);
             for (QueryBuilder subQb : this.clauses) {
-                subVisitor.accept(subQb);
+                subVisitor.accept(subQb, level+1);
             }
         }
     }

diff --git a/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java
@@ -287,14 +287,14 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (include != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(include);
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(include, level+1);
         }
 
         if (exclude != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.MUST_NOT).accept(exclude);
+            visitor.getChildVisitor(BooleanClause.Occur.MUST_NOT).accept(exclude, level+1);
         }
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java
@@ -191,12 +191,12 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
         if (clauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
             for (QueryBuilder subQb : this.clauses) {
-                subVisitor.accept(subQb);
+                subVisitor.accept(subQb, level+1);
             }
         }
     }

diff --git a/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java
@@ -200,9 +200,9 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor) {
-        visitor.accept(this);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    public void visit(QueryBuilderVisitor visitor, int level) {
+        visitor.accept(this, level);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big, level+1);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little, level+1);
     }
 }