Fix that `Query` cannot be updated after it's set in `Search` (#7710)

oap-server/server-library/library-elasticsearch-client/src/main/java/org/apache/skywalking/library/elasticsearch/requests/search/MatchQuery.java

@@ -38,10 +38,14 @@ public final class MatchQuery extends Query {
         public void serialize(final MatchQuery value, final JsonGenerator gen,
                               final SerializerProvider provider)
             throws IOException {
-            gen.writeFieldName("match");
             gen.writeStartObject();
             {
-                gen.writeStringField(value.getName(), value.getText());
+                gen.writeFieldName("match");
+                gen.writeStartObject();
+                {
+                    gen.writeStringField(value.getName(), value.getText());
+                }
+                gen.writeEndObject();
             }
             gen.writeEndObject();
         }

oap-server/server-library/library-elasticsearch-client/src/main/java/org/apache/skywalking/library/elasticsearch/requests/search/SearchBuilder.java

@@ -24,12 +24,14 @@ import org.apache.skywalking.library.elasticsearch.requests.search.aggregation.A
 import org.apache.skywalking.library.elasticsearch.requests.search.aggregation.AggregationBuilder;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
 public final class SearchBuilder {
     private Integer from;
     private Integer size;
     private Query query;
+    private QueryBuilder queryBuilder;
     private ImmutableList.Builder<Sort> sort;
     private ImmutableMap.Builder<String, Aggregation> aggregations;
 
@@ -58,13 +60,15 @@ public final class SearchBuilder {
     }
 
     public SearchBuilder query(Query query) {
-        requireNonNull(query, "query");
-        this.query = query;
+        ensureQueryIsNotSet();
+        this.query = requireNonNull(query, "query");
         return this;
     }
 
     public SearchBuilder query(QueryBuilder queryBuilder) {
-        return query(queryBuilder.build());
+        ensureQueryIsNotSet();
+        this.queryBuilder = requireNonNull(queryBuilder, "queryBuilder");
+        return this;
     }
 
     public SearchBuilder aggregation(Aggregation aggregation) {
@@ -92,6 +96,14 @@ public final class SearchBuilder {
         } else {
             aggregations = aggregations().build();
         }
+        final Query query;
+        if (this.query != null) {
+            query = this.query;
+        } else if (queryBuilder != null) {
+            query = queryBuilder.build();
+        } else {
+            query = null;
+        }
 
         return new Search(
             from, size, query, sorts, aggregations
@@ -111,4 +123,10 @@ public final class SearchBuilder {
         }
         return aggregations;
     }
+
+    private void ensureQueryIsNotSet() {
+        final String errMsg = "query and queryBuilder can not be set simultaneously";
+        checkState(query == null, errMsg);
+        checkState(queryBuilder == null, errMsg);
+    }
 }

oap-server/server-library/library-elasticsearch-client/src/test/java/org/apache/skywalking/library/elasticsearch/requests/search/SearchBuilderTest.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.skywalking.library.elasticsearch.requests.search;
+
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class SearchBuilderTest {
+    @Test
+    public void searchQueryShouldBeUpdatableAfterSet() {
+        final BoolQueryBuilder queryBuilder = Query.bool();
+        final SearchBuilder searchBuilder = Search.builder().query(queryBuilder);
+        queryBuilder.must(Query.term("t", "v"));
+        queryBuilder.should(Query.term("t", "v"));
+        queryBuilder.mustNot(Query.term("t2", "v2"));
+        queryBuilder.shouldNot(Query.term("t2", "v2"));
+        queryBuilder.shouldNot(Query.term("t2", "v2"));
+
+        final BoolQuery query = (BoolQuery) searchBuilder.build().getQuery();
+        assertThat(query.getMust()).hasSize(1);
+        assertThat(query.getShould()).hasSize(1);
+        assertThat(query.getMustNot()).hasSize(1);
+        assertThat(query.getShouldNot()).hasSize(2);
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void searchQueryBuilderShouldNotBeSetMultipleTimes() {
+        final BoolQueryBuilder queryBuilder = Query.bool();
+        final SearchBuilder searchBuilder = Search.builder().query(queryBuilder);
+        searchBuilder.query(Query.bool());
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void searchQueryShouldNotBeSetMultipleTimes() {
+        final SearchBuilder searchBuilder = Search.builder().query(Query.bool().build());
+        searchBuilder.query(Query.bool().build());
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void searchQueryAndBuilderShouldNotBeSetSimultaneously() {
+        final SearchBuilder searchBuilder = Search.builder().query(Query.bool().build());
+        searchBuilder.query(Query.bool());
+    }
+}

pom.xml

@@ -209,6 +209,7 @@
         <jmh.version>1.21</jmh.version>
         <gmaven-plugin.version>1.5</gmaven-plugin.version>
         <checkstyle.fails.on.error>true</checkstyle.fails.on.error>
+        <assertj-core.version>3.20.2</assertj-core.version>
 
     </properties>
 
@@ -233,6 +234,10 @@
             <artifactId>objenesis</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.assertj</groupId>
+            <artifactId>assertj-core</artifactId>
+        </dependency>
 
         <dependency>
             <groupId>org.openjdk.jmh</groupId>
@@ -287,6 +292,12 @@
                 <version>${objenesis.version}</version>
                 <scope>test</scope>
             </dependency>
+            <dependency>
+                <groupId>org.assertj</groupId>
+                <artifactId>assertj-core</artifactId>
+                <version>${assertj-core.version}</version>
+                <scope>test</scope>
+            </dependency>
             <dependency>
                 <groupId>org.openjdk.jmh</groupId>
                 <artifactId>jmh-generator-annprocess</artifactId>