Fixed review comments

Author: rahulmane-goog

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java

@@ -637,15 +637,6 @@ AggregationResults runAggregation(
   /** Returns true if this background resource has been shut down. */
   boolean isClosed();
 
-  /**
-   * Returns a new Datastore client with the specified request tags added.
-   *
-   * @param requestTags the request tags to append to existing ones
-   */
-  default Datastore withRequestTags(String... requestTags) {
-    return withRequestTags(java.util.Arrays.asList(requestTags));
-  }
-
   /**
    * Returns a new Datastore client with the specified request tags added.
    *
@@ -655,6 +646,6 @@ default Datastore withRequestTags(List<String> requestTags) {
     ImmutableList.Builder<String> builder = ImmutableList.builder();
     builder.addAll(getOptions().getRequestTags());
     builder.addAll(requestTags);
-    return getOptions().toBuilder().setTags(builder.build()).build().getService();
+    return getOptions().toBuilder().setRequestTags(builder.build()).build().getService();
   }
 }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

@@ -381,13 +381,14 @@ <T> QueryResults<T> run(
       Query<T> query,
       ExplainOptions explainOptions,
       RequestOptions requestOptions) {
-    return new QueryResultsImpl<T>(
-        this,
-        readOptionsPb,
-        (RecordQuery<T>) query,
-        query.getNamespace(),
-        explainOptions,
-        requestOptions);
+    return new QueryResultsImpl.Builder<T>()
+        .setDatastore(this)
+        .setReadOptionsPb(readOptionsPb)
+        .setQuery((RecordQuery<T>) query)
+        .setNamespace(query.getNamespace())
+        .setExplainOptions(explainOptions)
+        .setRequestOptions(requestOptions)
+        .build();
   }
 
   @Override

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java

@@ -256,12 +256,26 @@ public Builder setDatabaseId(String databaseId) {
       return this;
     }
 
-    public Builder setTags(ImmutableList<String> requestTags) {
+    /**
+     * Sets the request tags to be associated with all requests sent by this client.
+     *
+     * @param requestTags the list of request tags to set
+     * @return the builder object
+     */
+    public Builder setRequestTags(ImmutableList<String> requestTags) {
+      Preconditions.checkNotNull(requestTags, "Request tags cannot be null");
       this.requestTags = requestTags;
       return this;
     }
 
-    public Builder setTags(String... requestTags) {
+    /**
+     * Sets the request tags to be associated with all requests sent by this client.
+     *
+     * @param requestTags the request tags to set
+     * @return the builder object
+     */
+    public Builder setRequestTags(String... requestTags) {
+      Preconditions.checkNotNull(requestTags, "Request tags cannot be null");
       this.requestTags = ImmutableList.copyOf(requestTags);
       return this;
     }
@@ -386,6 +400,11 @@ public String getDatabaseId() {
     return this.databaseId;
   }
 
+  /**
+   * Returns the request tags to be associated with all requests sent by this client.
+   *
+   * @return the request tags
+   */
   public ImmutableList<String> getRequestTags() {
     return requestTags;
   }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java

@@ -52,35 +52,18 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
   private ExplainMetrics explainMetrics;
   private final RequestOptions requestOptions;
 
-  /** Creates a QueryResultsImpl. */
-  QueryResultsImpl(
-      DatastoreImpl datastore,
-      Optional<ReadOptions> readOptionsPb,
-      RecordQuery<T> query,
-      String namespace,
-      ExplainOptions explainOptions) {
-    this(datastore, readOptionsPb, query, namespace, explainOptions, null);
-  }
-
-  /** Creates a QueryResultsImpl with RequestOptions. */
-  QueryResultsImpl(
-      DatastoreImpl datastore,
-      Optional<ReadOptions> readOptionsPb,
-      RecordQuery<T> query,
-      String namespace,
-      ExplainOptions explainOptions,
-      RequestOptions requestOptions) {
-    this.datastore = datastore;
-    this.readOptionsPb = readOptionsPb;
-    this.query = query;
-    queryResultType = query.getType();
-    this.explainOptions = explainOptions;
-    this.requestOptions = requestOptions;
+  private QueryResultsImpl(Builder<T> builder) {
+    this.datastore = builder.datastore;
+    this.readOptionsPb = builder.readOptionsPb;
+    this.query = builder.query;
+    queryResultType = builder.query.getType();
+    this.explainOptions = builder.explainOptions;
+    this.requestOptions = builder.requestOptions;
     PartitionId.Builder pbBuilder = PartitionId.newBuilder();
     pbBuilder.setProjectId(datastore.getOptions().getProjectId());
     pbBuilder.setDatabaseId(datastore.getOptions().getDatabaseId());
-    if (namespace != null) {
-      pbBuilder.setNamespaceId(namespace);
+    if (builder.namespace != null) {
+      pbBuilder.setNamespaceId(builder.namespace);
     } else if (datastore.getOptions().getNamespace() != null) {
       pbBuilder.setNamespaceId(datastore.getOptions().getNamespace());
     }
@@ -93,6 +76,51 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
     }
   }
 
+  static class Builder<T> {
+    private DatastoreImpl datastore;
+    private Optional<ReadOptions> readOptionsPb = Optional.empty();
+    private RecordQuery<T> query;
+    private String namespace;
+    private ExplainOptions explainOptions;
+    private RequestOptions requestOptions;
+
+    Builder<T> setDatastore(DatastoreImpl datastore) {
+      this.datastore = datastore;
+      return this;
+    }
+
+    Builder<T> setReadOptionsPb(Optional<ReadOptions> readOptionsPb) {
+      this.readOptionsPb = readOptionsPb;
+      return this;
+    }
+
+    Builder<T> setQuery(RecordQuery<T> query) {
+      this.query = query;
+      return this;
+    }
+
+    Builder<T> setNamespace(String namespace) {
+      this.namespace = namespace;
+      return this;
+    }
+
+    Builder<T> setExplainOptions(ExplainOptions explainOptions) {
+      this.explainOptions = explainOptions;
+      return this;
+    }
+
+    Builder<T> setRequestOptions(RequestOptions requestOptions) {
+      this.requestOptions = requestOptions;
+      return this;
+    }
+
+    QueryResultsImpl<T> build() {
+      Preconditions.checkNotNull(datastore, "datastore cannot be null");
+      Preconditions.checkNotNull(query, "query cannot be null");
+      return new QueryResultsImpl<>(this);
+    }
+  }
+
   private void sendRequest() {
     RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
     readOptionsPb.ifPresent(requestPb::setReadOptions);

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/ReadOption.java

@@ -19,6 +19,7 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.cloud.Timestamp;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.datastore.v1.ExplainOptions;
 import com.google.datastore.v1.RequestOptions;
@@ -152,14 +153,6 @@ public static class QueryConfig<Q extends Query<?>> {
     ExplainOptions explainOptions;
     RequestOptions requestOptions;
 
-    private QueryConfig(Q query, ExplainOptions explainOptions, List<ReadOption> readOptions) {
-      this(query, explainOptions, readOptions, null);
-    }
-
-    private QueryConfig(Q query, ExplainOptions explainOptions) {
-      this(query, explainOptions, Collections.emptyList(), null);
-    }
-
     private QueryConfig(
         Q query,
         ExplainOptions explainOptions,
@@ -189,20 +182,69 @@ public RequestOptions getRequestOptions() {
 
     public static <Q extends Query<?>> QueryConfig<Q> create(
         Q query, ExplainOptions explainOptions) {
-      return new QueryConfig<>(query, explainOptions);
+      return QueryConfig.<Q>newBuilder().setQuery(query).setExplainOptions(explainOptions).build();
     }
 
     public static <Q extends Query<?>> QueryConfig<Q> create(
         Q query, ExplainOptions explainOptions, List<ReadOption> readOptions) {
-      return new QueryConfig<>(query, explainOptions, readOptions);
+      return QueryConfig.<Q>newBuilder()
+          .setQuery(query)
+          .setExplainOptions(explainOptions)
+          .setReadOptions(readOptions)
+          .build();
     }
 
     public static <Q extends Query<?>> QueryConfig<Q> create(
         Q query,
         ExplainOptions explainOptions,
         List<ReadOption> readOptions,
         RequestOptions requestOptions) {
-      return new QueryConfig<>(query, explainOptions, readOptions, requestOptions);
+      return QueryConfig.<Q>newBuilder()
+          .setQuery(query)
+          .setExplainOptions(explainOptions)
+          .setReadOptions(readOptions)
+          .setRequestOptions(requestOptions)
+          .build();
+    }
+
+    /** Creates a new builder for {@link QueryConfig}. */
+    public static <Q extends Query<?>> Builder<Q> newBuilder() {
+      return new Builder<>();
+    }
+
+    /** Builder for {@link QueryConfig}. */
+    public static class Builder<Q extends Query<?>> {
+      private Q query;
+      private List<ReadOption> readOptions = Collections.emptyList();
+      private ExplainOptions explainOptions;
+      private RequestOptions requestOptions;
+
+      private Builder() {}
+
+      public Builder<Q> setQuery(Q query) {
+        this.query = query;
+        return this;
+      }
+
+      public Builder<Q> setReadOptions(List<ReadOption> readOptions) {
+        this.readOptions = readOptions;
+        return this;
+      }
+
+      public Builder<Q> setExplainOptions(ExplainOptions explainOptions) {
+        this.explainOptions = explainOptions;
+        return this;
+      }
+
+      public Builder<Q> setRequestOptions(RequestOptions requestOptions) {
+        this.requestOptions = requestOptions;
+        return this;
+      }
+
+      public QueryConfig<Q> build() {
+        Preconditions.checkNotNull(query, "query cannot be null");
+        return new QueryConfig<>(query, explainOptions, readOptions, requestOptions);
+      }
     }
   }
 }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RequestOptionsHelper.java

@@ -16,9 +16,11 @@
 
 package com.google.cloud.datastore;
 
+import com.google.api.core.InternalApi;
 import com.google.datastore.v1.RequestOptions;
 
 /** Helper class for building and merging Datastore request options. */
+@InternalApi
 public final class RequestOptionsHelper {
 
   private RequestOptionsHelper() {}

java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AbstractDatastoreTest.java

@@ -1402,7 +1402,7 @@ private Predicate<RunAggregationQueryRequest> aggregationQueryWithAlias(String a
 
   @Test
   public void testRunQueryWithInstanceLevelRequestTags() {
-    DatastoreOptions optionsWithTags = options.toBuilder().setTags("instance-tag").build();
+    DatastoreOptions optionsWithTags = options.toBuilder().setRequestTags("instance-tag").build();
     Datastore datastoreWithTags = optionsWithTags.getService();
 
     PartitionId partitionId =
@@ -1423,7 +1423,9 @@ public void testRunQueryWithInstanceLevelRequestTags() {
             .setPartitionId(partitionId)
             .setQuery(queryPb)
             .setRequestOptions(
-                com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("instance-tag").build())
+                com.google.datastore.v1.RequestOptions.newBuilder()
+                    .addRequestTags("instance-tag")
+                    .build())
             .build();
 
     RunQueryResponse response =
@@ -1535,6 +1537,38 @@ public void testRunQueryWithRequestOptions() {
     EasyMock.verify(rpcFactoryMock, rpcMock);
   }
 
+  @Test
+  public void testRunAggregationQueryWithInstanceLevelRequestTags() {
+    DatastoreOptions optionsWithTags = options.toBuilder().setRequestTags("instance-tag").build();
+    Datastore datastoreWithTags = optionsWithTags.getService();
+
+    com.google.datastore.v1.RequestOptions requestOptions =
+        com.google.datastore.v1.RequestOptions.newBuilder().addRequestTags("test-tag").build();
+
+    RunAggregationQueryResponse aggregationQueryResponse = placeholderAggregationQueryResponse();
+
+    EasyMock.expect(
+            rpcMock.runAggregationQuery(
+                matches(
+                    aggregationQueryWithAliasAndRequestOptions(
+                        "total_count", requestOptions, false, false))))
+        .andReturn(aggregationQueryResponse);
+
+    EasyMock.replay(rpcFactoryMock, rpcMock);
+
+    EntityQuery selectAllQuery = Query.newEntityQueryBuilder().build();
+    AggregationQuery getCountQuery =
+        Query.newAggregationQueryBuilder()
+            .addAggregation(count().as("total_count"))
+            .over(selectAllQuery)
+            .build();
+
+    AggregationResult result = getOnlyElement(datastoreWithTags.runAggregation(getCountQuery));
+    assertThat(result.getLong("total_count")).isEqualTo(209L);
+
+    EasyMock.verify(rpcFactoryMock, rpcMock);
+  }
+
   @Test
   public void testRunAggregationQueryWithRequestOptions() {
     com.google.datastore.v1.RequestOptions requestOptions =

java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreOptionsTest.java

@@ -204,6 +204,20 @@ public void testNamespace() {
     assertEquals("ns1", options.setNamespace("ns1").build().getNamespace());
   }
 
+  @Test
+  public void testRequestTags() {
+    assertTrue(options.build().getRequestTags().isEmpty());
+    assertThat(options.setRequestTags("tag1", "tag2").build().getRequestTags())
+        .containsExactly("tag1", "tag2")
+        .inOrder();
+    assertThat(
+            options
+                .setRequestTags(com.google.common.collect.ImmutableList.of("tag3"))
+                .build()
+                .getRequestTags())
+        .containsExactly("tag3");
+  }
+
   @Test
   public void testDatastore() {
     assertSame(datastoreRpc, options.build().getRpc());