diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java index 2aaeab3afb95d..b8f6da0212e03 100644 --- 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 index c44a7ef6a397c..5c0956030e183 100644 --- 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 index 1b52ae2f03605..5fa268770bfe8 100644 --- 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 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 index b2764d29da80a..d27f020a603c9 100644 --- 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 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 index bd8ec62f6c43e..36e223d227b52 100644 --- 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 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 index 1162689a54689..f59d4a8e1b207 100644 --- 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 index 090f74c5be7fe..7d6c024846a1c 100644 --- 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 index af5a125f9dd95..4cae0f00d746d 100644 --- 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 index 0b513b35b32c6..4b0491d38f155 100644 --- 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> 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 index 32a19ea3e9b50..f46b2d7fc60f8 100644 --- 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 index bcbc64ddf386d..ac3ad8d75c028 100644 --- 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 index 96d03c91964e3..31564d7983629 100644 --- 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 index 30a1c29c29126..b4fcfd015dad2 100644 --- 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 index 59ec5b9d77fc8..d0415c7b24209 100644 --- 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 index fae1e318c66bd..79deadd7da38e 100644 --- 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 index 4d5a6dde61a70..fb32dfad76dc0 100644 --- 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); } }