Backport/backport 10110 to 2.x

Signed-off-by: Varun Jain <varunudr@amazon.com>
opensearch-project · Sep 21, 2023 · ef84d3c · ef84d3c
1 parent 4017359
commit ef84d3c
Show file tree

Hide file tree

Showing 30 changed files with 395 additions and 17 deletions.
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 - Add metrics for thread_pool task wait time ([#9681](https://github.com/opensearch-project/OpenSearch/pull/9681))
 - Async blob read support for S3 plugin ([#9694](https://github.com/opensearch-project/OpenSearch/pull/9694))
+- 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 JNA version from 5.5 to 5.13 ([#9963](https://github.com/opensearch-project/OpenSearch/pull/9963))
@@ -34,4 +36,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Security
 
-[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
diff --git a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java b/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) {
+        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);
+            }
+        }
+
+    }
+
 }
diff --git a/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java b/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) {
+        visitor.accept(this);
+        if (positiveQuery != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery);
+        }
+        if (negativeQuery != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery);
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java b/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) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder);
+    }
+
 }
diff --git a/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java b/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) {
+        visitor.accept(this);
+        if (queries.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
+            for (QueryBuilder subQb : queries) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java b/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) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilder.java b/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) {
+        visitor.accept(this);
+    };
+
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java b/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java
@@ -0,0 +1,48 @@
+/*
+ * 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;
+
+/**
+ * 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);
+
+    /**
+     * 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;
+        }
+    };
+
+}
diff --git a/server/src/main/java/org/opensearch/index/query/SpanContainingQueryBuilder.java b/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) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanFirstQueryBuilder.java b/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) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(matchBuilder);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanMultiTermQueryBuilder.java b/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) {
+        visitor.accept(this);
+        if (multiTermQueryBuilder != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(multiTermQueryBuilder);
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanNearQueryBuilder.java b/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) {
+        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

diff --git a/server/src/main/java/org/opensearch/index/query/SpanNotQueryBuilder.java b/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) {
+        visitor.accept(this);
+        if (include != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(include);
+        }
+
+        if (exclude != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST_NOT).accept(exclude);
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanOrQueryBuilder.java b/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) {
+        visitor.accept(this);
+        if (clauses.isEmpty() == false) {
+            QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
+            for (QueryBuilder subQb : this.clauses) {
+                subVisitor.accept(subQb);
+            }
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/SpanWithinQueryBuilder.java b/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) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(big);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(little);
+    }
 }