Repository: googleapis/java-datastore
Pull request #875 (2022-10-13 06:20:31)

refactor: making query class generic 
This change decouples Query.java from the concerns related to handling entity records. It is required as the definition of a Query is changed now and it is not just limited to the the fact that every query returns entity records, With the introduction of count feature and aggregation query, we will have queries returning aggregation results.


_Note: This PR is intended to trim down the changes done as part of #823._

conventional-commit-lint-gcf[bot]: <!-- probot comment [1356997]-->
🤖 I detect that the PR title and the commit message differ and there's only one commit. To use the PR title for the commit history, you can use Github's automerge feature with squashing, or use `automerge` label. Good luck human!

 -- conventional-commit-lint bot
https://conventionalcommits.org/

Patch:
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index 9892e151..2d68a229 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -180,8 +180,10 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
     return run(toReadOptionsPb(options), query);
   }
 
+  @SuppressWarnings("unchecked")
   <T> QueryResults<T> run(com.google.datastore.v1.ReadOptions readOptionsPb, Query<T> query) {
-    return new QueryResultsImpl<>(this, readOptionsPb, query);
+    return new QueryResultsImpl<>(
+        this, readOptionsPb, (RecordQuery<T>) query, query.getNamespace());
   }
 
   com.google.datastore.v1.RunQueryResponse runQuery(
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java
index 2b99fd0a..bebb4df9 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GqlQuery.java
@@ -19,6 +19,7 @@ package com.google.cloud.datastore;
 import static com.google.cloud.datastore.Validator.validateNamespace;
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.api.core.InternalApi;
 import com.google.cloud.Timestamp;
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
@@ -71,7 +72,7 @@ import java.util.TreeMap;
  * @param <V> the type of the result values this query will produce
  * @see <a href="https://cloud.google.com/datastore/docs/apis/gql/gql_reference">GQL Reference</a>
  */
-public final class GqlQuery<V> extends Query<V> {
+public final class GqlQuery<V> extends Query<V> implements RecordQuery<V> {
 
   private static final long serialVersionUID = -5514894742849230793L;
 
@@ -80,6 +81,8 @@ public final class GqlQuery<V> extends Query<V> {
   private final ImmutableMap<String, Binding> namedBindings;
   private final ImmutableList<Binding> positionalBindings;
 
+  private final ResultType<V> resultType;
+
   static final class Binding implements Serializable {
 
     private static final long serialVersionUID = 2344746877591371548L;
@@ -423,7 +426,8 @@ public final class GqlQuery<V> extends Query<V> {
   }
 
   private GqlQuery(Builder<V> builder) {
-    super(builder.resultType, builder.namespace);
+    super(builder.namespace);
+    resultType = checkNotNull(builder.resultType);
     queryString = builder.queryString;
     allowLiteral = builder.allowLiteral;
     namedBindings = ImmutableMap.copyOf(builder.namedBindings);
@@ -461,9 +465,15 @@ public final class GqlQuery<V> extends Query<V> {
     return builder.build();
   }
 
+  @Override
+  public ResultType<V> getType() {
+    return resultType;
+  }
+
   @Override
   public String toString() {
-    return super.toStringHelper()
+    return toStringHelper()
+        .add("type", getType())
         .add("queryString", queryString)
         .add("allowLiteral", allowLiteral)
         .add("namedBindings", namedBindings)
@@ -507,13 +517,15 @@ public final class GqlQuery<V> extends Query<V> {
     return queryPb.build();
   }
 
+  @InternalApi
   @Override
-  void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) {
+  public void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) {
     requestPb.setGqlQuery(toPb());
   }
 
+  @InternalApi
   @Override
-  Query<V> nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) {
+  public RecordQuery<V> nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) {
     return StructuredQuery.<V>fromPb(getType(), getNamespace(), responsePb.getQuery())
         .nextQuery(responsePb);
   }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
index 00aa6f17..a0bed598 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Query.java
@@ -16,8 +16,6 @@
 
 package com.google.cloud.datastore;
 
-import static com.google.common.base.Preconditions.checkNotNull;
-
 import com.google.common.base.MoreObjects;
 import com.google.common.base.MoreObjects.ToStringHelper;
 import com.google.common.collect.Maps;
@@ -39,7 +37,6 @@ public abstract class Query<V> implements Serializable {
 
   private static final long serialVersionUID = 7967659059395653941L;
 
-  private final ResultType<V> resultType;
   private final String namespace;
 
   /**
@@ -156,27 +153,18 @@ public abstract class Query<V> implements Serializable {
     }
   }
 
-  Query(ResultType<V> resultType, String namespace) {
-    this.resultType = checkNotNull(resultType);
+  Query(String namespace) {
     this.namespace = namespace;
   }
 
-  ResultType<V> getType() {
-    return resultType;
-  }
-
   public String getNamespace() {
     return namespace;
   }
 
   ToStringHelper toStringHelper() {
-    return MoreObjects.toStringHelper(this).add("type", resultType).add("namespace", namespace);
+    return MoreObjects.toStringHelper(this).add("namespace", namespace);
   }
 
-  abstract void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb);
-
-  abstract Query<V> nextQuery(com.google.datastore.v1.RunQueryResponse responsePb);
-
   /**
    * Returns a new {@link GqlQuery} builder.
    *
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
index 9ed82298..222efa3b 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResultsImpl.java
@@ -30,7 +30,7 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
   private final com.google.datastore.v1.ReadOptions readOptionsPb;
   private final com.google.datastore.v1.PartitionId partitionIdPb;
   private final ResultType<T> queryResultType;
-  private Query<T> query;
+  private RecordQuery<T> query;
   private ResultType<?> actualResultType;
   private com.google.datastore.v1.RunQueryResponse runQueryResponsePb;
   private com.google.datastore.v1.Query mostRecentQueryPb;
@@ -40,7 +40,10 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
   private MoreResultsType moreResults;
 
   QueryResultsImpl(
-      DatastoreImpl datastore, com.google.datastore.v1.ReadOptions readOptionsPb, Query<T> query) {
+      DatastoreImpl datastore,
+      com.google.datastore.v1.ReadOptions readOptionsPb,
+      RecordQuery<T> query,
+      String namespace) {
     this.datastore = datastore;
     this.readOptionsPb = readOptionsPb;
     this.query = query;
@@ -48,8 +51,8 @@ class QueryResultsImpl<T> extends AbstractIterator<T> implements QueryResults<T>
     com.google.datastore.v1.PartitionId.Builder pbBuilder =
         com.google.datastore.v1.PartitionId.newBuilder();
     pbBuilder.setProjectId(datastore.getOptions().getProjectId());
-    if (query.getNamespace() != null) {
-      pbBuilder.setNamespaceId(query.getNamespace());
+    if (namespace != null) {
+      pbBuilder.setNamespaceId(namespace);
     } else if (datastore.getOptions().getNamespace() != null) {
       pbBuilder.setNamespaceId(datastore.getOptions().getNamespace());
     }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java
new file mode 100644
index 00000000..9dc96645
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RecordQuery.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed 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
+ *
+ *     https://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 com.google.cloud.datastore;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.datastore.Query.ResultType;
+
+/** An internal marker interface to represent {@link Query} that returns the entity records. */
+@InternalApi
+public interface RecordQuery<V> {
+
+  @InternalApi
+  ResultType<V> getType();
+
+  @InternalApi
+  void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb);
+
+  @InternalApi
+  RecordQuery<V> nextQuery(com.google.datastore.v1.RunQueryResponse responsePb);
+}
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
index 8e50d086..8d2974c1 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/StructuredQuery.java
@@ -26,6 +26,7 @@ import static com.google.cloud.datastore.TimestampValue.of;
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.api.core.ApiFunction;
+import com.google.api.core.InternalApi;
 import com.google.cloud.StringEnumType;
 import com.google.cloud.StringEnumValue;
 import com.google.cloud.Timestamp;
@@ -85,7 +86,7 @@ import java.util.Objects;
  * @see <a href="https://cloud.google.com/appengine/docs/java/datastore/queries">Datastore
  *     queries</a>
  */
-public abstract class StructuredQuery<V> extends Query<V> {
+public abstract class StructuredQuery<V> extends Query<V> implements RecordQuery<V> {
 
   private static final long serialVersionUID = 546838955624019594L;
   static final String KEY_PROPERTY_NAME = "__key__";
@@ -100,6 +101,8 @@ public abstract class StructuredQuery<V> extends Query<V> {
   private final int offset;
   private final Integer limit;
 
+  private final ResultType<V> resultType;
+
   public abstract static class Filter implements Serializable {
 
     private static final long serialVersionUID = -6443285436239990860L;
@@ -899,7 +902,8 @@ public abstract class StructuredQuery<V> extends Query<V> {
   }
 
   StructuredQuery(BuilderImpl<V, ?> builder) {
-    super(builder.resultType, builder.namespace);
+    super(builder.namespace);
+    resultType = checkNotNull(builder.resultType);
     kind = builder.kind;
     projection = ImmutableList.copyOf(builder.projection);
     filter = builder.filter;
@@ -914,6 +918,7 @@ public abstract class StructuredQuery<V> extends Query<V> {
   @Override
   public String toString() {
     return toStringHelper()
+        .add("type", getType())
         .add("kind", kind)
         .add("startCursor", startCursor)
         .add("endCursor", endCursor)
@@ -1013,13 +1018,19 @@ public abstract class StructuredQuery<V> extends Query<V> {
 
   public abstract Builder<V> toBuilder();
 
+  public ResultType<V> getType() {
+    return resultType;
+  }
+
+  @InternalApi
   @Override
-  void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) {
+  public void populatePb(com.google.datastore.v1.RunQueryRequest.Builder requestPb) {
     requestPb.setQuery(toPb());
   }
 
+  @InternalApi
   @Override
-  StructuredQuery<V> nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) {
+  public StructuredQuery<V> nextQuery(com.google.datastore.v1.RunQueryResponse responsePb) {
     Builder<V> builder = toBuilder();
     builder.setStartCursor(new Cursor(responsePb.getBatch().getEndCursor()));
     if (offset > 0 && responsePb.getBatch().getSkippedResults() < offset) {


Test patch:
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index fa077bc6..72067fd2 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -613,7 +613,7 @@ public class DatastoreTest {
     Entity entity5 = Entity.newBuilder(KEY5).set("value", "value").build();
     datastore.add(ENTITY3, entity4, entity5);
     List<RunQueryResponse> responses = new ArrayList<>();
-    Query<Key> query = Query.newKeyQueryBuilder().build();
+    RecordQuery<Key> query = Query.newKeyQueryBuilder().build();
     RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
     query.populatePb(requestPb);
     QueryResultBatch queryResultBatchPb =
@@ -722,7 +722,7 @@ public class DatastoreTest {
     datastore.add(ENTITY3, entity4, entity5);
     DatastoreRpc datastoreRpc = datastore.getOptions().getDatastoreRpcV1();
     List<RunQueryResponse> responses = new ArrayList<>();
-    Query<Entity> query = Query.newEntityQueryBuilder().build();
+    RecordQuery<Entity> query = Query.newEntityQueryBuilder().build();
     RunQueryRequest.Builder requestPb = RunQueryRequest.newBuilder();
     query.populatePb(requestPb);
     QueryResultBatch queryResultBatchPb =