[Feature] Visitor design pattern in QueryBuilder

CHANGELOG.md

@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add events correlation engine plugin ([#6854](https://github.com/opensearch-project/OpenSearch/issues/6854))
 - Introduce new dynamic cluster setting to control slice computation for concurrent segment search ([#9107](https://github.com/opensearch-project/OpenSearch/pull/9107))
 - Implement on behalf of token passing for extensions ([#8679](https://github.com/opensearch-project/OpenSearch/pull/8679))
+- Implement Visitor Design pattern in QueryBuilder to enable the capability to traverse through the complex QueryBuilder tree. ([#10110](https://github.com/opensearch-project/OpenSearch/pull/10110))
 
 ### Dependencies
 - Bump `log4j-core` from 2.18.0 to 2.19.0
@@ -107,4 +108,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Security
 
 [Unreleased 3.0]: https://github.com/opensearch-project/OpenSearch/compare/2.x...HEAD
-[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.11...2.x
+[Unreleased 2.x]: https://github.com/opensearch-project/OpenSearch/compare/2.11...2.x

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

@@ -427,4 +427,35 @@ private static boolean rewriteClauses(
         }
         return changed;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (mustClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST);
+            for (QueryBuilder mustClause : mustClauses) {
+                mustClause.visit(subVisitor);
+            }
+        }
+        if (shouldClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.SHOULD);
+            for (QueryBuilder shouldClause : shouldClauses) {
+                shouldClause.visit(subVisitor);
+            }
+        }
+        if (mustNotClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST_NOT);
+            for (QueryBuilder mustNotClause : mustNotClauses) {
+                mustNotClause.visit(subVisitor);
+            }
+        }
+        if (filterClauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.FILTER);
+            for (QueryBuilder filterClause : filterClauses) {
+                filterClause.visit(subVisitor);
+            }
+        }
+
+    }
+
 }

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

@@ -33,6 +33,7 @@
 package org.opensearch.index.query;
 
 import org.apache.lucene.queries.function.FunctionScoreQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -252,4 +253,15 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
         InnerHitContextBuilder.extractInnerHits(positiveQuery, innerHits);
         InnerHitContextBuilder.extractInnerHits(negativeQuery, innerHits);
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (positiveQuery != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery);
+        }
+        if (negativeQuery != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery);
+        }
+    }
 }

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

@@ -32,6 +32,7 @@
 
 package org.opensearch.index.query;
 
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
@@ -183,4 +184,11 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
     protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> innerHits) {
         InnerHitContextBuilder.extractInnerHits(filterBuilder, innerHits);
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder);
+    }
+
 }

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

@@ -32,6 +32,7 @@
 
 package org.opensearch.index.query;
 
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.DisjunctionMaxQuery;
 import org.apache.lucene.search.Query;
 import org.opensearch.common.lucene.search.Queries;
@@ -246,4 +247,15 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
             InnerHitContextBuilder.extractInnerHits(query, innerHits);
         }
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (queries.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
+            for (QueryBuilder subQb : queries) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
 }

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.FieldMaskingSpanQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -207,4 +208,10 @@ protected boolean doEquals(FieldMaskingSpanQueryBuilder other) {
     public String getWriteableName() {
         return SPAN_FIELD_MASKING_FIELD.getPreferredName();
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder);
+    }
 }

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

@@ -95,4 +95,13 @@ public interface QueryBuilder extends NamedWriteable, ToXContentObject, Rewritea
     default QueryBuilder rewrite(QueryRewriteContext queryShardContext) throws IOException {
         return this;
     }
+
+    /**
+     * 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.
+     */
+    default void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+    };
+
 }

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

@@ -0,0 +1,50 @@
+/*
+ * 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.index.query;
+
+import org.apache.lucene.search.BooleanClause;
+
+import java.io.IOException;
+
+/**
+ * QueryBuilderVisitor is an interface to define Visitor Object to be traversed in QueryBuilder tree.
+ */
+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.
+     */
+    void accept(QueryBuilder qb) throws IOException;
+
+    /**
+     * Fetches the child sub visitor from the main QueryBuilderVisitor Object.
+     * @param occur defines the occurrence of the result fetched from the search query in the final search result.
+     * @return a child queryBuilder Visitor Object.
+     */
+    QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur);
+
+    /**
+     *  NoopQueryVisitor is a default implementation of QueryBuilderVisitor.
+     *  When a user does not want to implement QueryBuilderVisitor and have to just pass an empty object then this class will be used.
+     *
+     */
+    QueryBuilderVisitor NO_OP_VISITOR = new QueryBuilderVisitor() {
+        @Override
+        public void accept(QueryBuilder qb) {
+            // Do nothing
+        }
+
+        @Override
+        public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+            return this;
+        }
+    };
+
+}

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanContainingQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -188,4 +189,11 @@ protected boolean doEquals(SpanContainingQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    }
 }

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanFirstQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -186,4 +187,10 @@ protected boolean doEquals(SpanFirstQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(matchBuilder);
+    }
 }

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

@@ -33,6 +33,7 @@
 
 import org.apache.lucene.queries.SpanMatchNoDocsQuery;
 import org.apache.lucene.queries.spans.SpanMultiTermQueryWrapper;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BoostQuery;
 import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
@@ -213,4 +214,12 @@ protected boolean doEquals(SpanMultiTermQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (multiTermQueryBuilder != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(multiTermQueryBuilder);
+        }
+    }
 }

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanNearQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -299,6 +300,17 @@ public String getWriteableName() {
         return NAME;
     }
 
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (this.clauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.MUST);
+            for (QueryBuilder subQb : this.clauses) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
+
     /**
      * SpanGapQueryBuilder enables gaps in a SpanNearQuery.
      * Since, SpanGapQuery is private to SpanNearQuery, SpanGapQueryBuilder cannot

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanNotQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -284,4 +285,16 @@ protected boolean doEquals(SpanNotQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (include != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(include);
+        }
+
+        if (exclude != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST_NOT).accept(exclude);
+        }
+    }
 }

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanOrQuery;
 import org.apache.lucene.queries.spans.SpanQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -188,4 +189,15 @@ protected boolean doEquals(SpanOrQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        if (clauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
+            for (QueryBuilder subQb : this.clauses) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
 }

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

@@ -34,6 +34,7 @@
 
 import org.apache.lucene.queries.spans.SpanQuery;
 import org.apache.lucene.queries.spans.SpanWithinQuery;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.Query;
 import org.opensearch.core.ParseField;
 import org.opensearch.core.common.ParsingException;
@@ -197,4 +198,11 @@ protected boolean doEquals(SpanWithinQueryBuilder other) {
     public String getWriteableName() {
         return NAME;
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) throws IOException {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    }
 }

server/src/test/java/org/opensearch/index/query/BoolQueryBuilderTests.java

@@ -456,4 +456,14 @@ public void testMustRewrite() throws IOException {
         IllegalStateException e = expectThrows(IllegalStateException.class, () -> boolQuery.toQuery(context));
         assertEquals("Rewrite first", e.getMessage());
     }
+
+    public void testVisit() {
+        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
+        boolQueryBuilder.should(new BoolQueryBuilder());
+        boolQueryBuilder.must(new BoolQueryBuilder());
+        boolQueryBuilder.mustNot(new BoolQueryBuilder());
+        boolQueryBuilder.filter(new BoolQueryBuilder());
+        IOException e = expectThrows(IOException.class, () -> boolQueryBuilder.visit(QueryBuilderVisitor.NO_OP_VISITOR));
+        assertEquals("Boosting Query Builder Traversed", e.getMessage());
+    }
 }

server/src/test/java/org/opensearch/index/query/BoostingQueryBuilderTests.java

@@ -153,4 +153,14 @@ public void testMustRewrite() throws IOException {
         e = expectThrows(IllegalStateException.class, () -> queryBuilder2.toQuery(context));
         assertEquals("Rewrite first", e.getMessage());
     }
+
+    public void testVisit() {
+        BoostingQueryBuilder builder = new BoostingQueryBuilder(
+            new TermQueryBuilder("unmapped_field", "value"),
+            new TermQueryBuilder(KEYWORD_FIELD_NAME, "other_value")
+        );
+
+        IOException e = expectThrows(IOException.class, () -> builder.visit(createTestVisitor()));
+        assertEquals("Boosting Query Builder Traversed", e.getMessage());
+    }
 }