feat: implement query profile

google-cloud-datastore/clirr-ignored-differences.xml

@@ -1,16 +1,73 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!-- see http://mojo.codehaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
 <differences>
-    <!-- renaming internal only class-->
+    <!--  added a default method to an interface-->
     <difference>
-        <className>com/google/cloud/datastore/ReadOption$QueryAndReadOptions</className>
-        <method>*</method>
-        <differenceType>8001</differenceType>
+        <className>com/google/cloud/datastore/Datastore</className>
+        <method>com.google.cloud.datastore.QueryResults run(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.QueryProfile$QueryMode, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7012</differenceType>
     </difference>
     <difference>
-        <className>com/google/cloud/datastore/execution/request/AggregationQueryRequestProtoPreparer</className>
-        <method>*QueryAndReadOptions*</method>
-        <to>*QueryConfig*</to>
+        <className>com/google/cloud/datastore/Datastore</className>
+        <method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.QueryProfile$QueryMode, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/QueryResults</className>
+        <method>java.util.Optional getResultSetStats()</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/DatastoreReader</className>
+        <method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.models.QueryProfile$QueryMode)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/DatastoreReader</className>
+        <method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery,  com.google.cloud.datastore.models.QueryProfile$QueryMode)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/Transaction</className>
+        <method>com.google.cloud.datastore.QueryResults run(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.QueryProfile$QueryMode)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+
+    <!-- added to internal only interface-->
+    <difference>
+        <className>com/google/cloud/datastore/execution/QueryExecutor</className>
+        <method>java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.models.QueryProfile$QueryMode, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7012</differenceType>
+    </difference>
+
+    <!-- updated internal only interface-->
+    <difference>
+        <className>com/google/cloud/datastore/execution/AggregationQueryExecutor</className>
+        <method>com.google.cloud.datastore.AggregationResults execute(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7004</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/execution/AggregationQueryExecutor</className>
+        <method>java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7004</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/execution/QueryExecutor</className>
+        <method>java.lang.Object execute(com.google.cloud.datastore.Query, com.google.cloud.datastore.ReadOption[])</method>
+        <differenceType>7004</differenceType>
+    </difference>
+
+    <!-- updated internal only class and method-->
+    <difference>
+        <className>com/google/cloud/datastore/ReadOption$QueryConfig</className>
+        <method>*QueryConfig*</method>
+        <to>*com.google.datastore.v1.QueryMode*</to>
         <differenceType>7005</differenceType>
     </difference>
+    <difference>
+        <className>com/google/cloud/datastore/ReadOption$QueryConfig</className>
+        <method>com.google.cloud.datastore.ReadOption$QueryConfig create(com.google.cloud.datastore.Query)</method>
+        <differenceType>7004</differenceType>
+    </difference>
+
 </differences>

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

@@ -19,9 +19,11 @@
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.Timestamp;
+import com.google.cloud.datastore.models.ResultSetStats;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 
 /**
  * The result of an {@link AggregationQuery} query submission. Contains a List<{@link
@@ -33,12 +35,21 @@ public class AggregationResults implements Iterable<AggregationResult> {
 
   private final List<AggregationResult> aggregationResults;
   private final Timestamp readTime;
+  private final ResultSetStats resultSetStats;
 
   public AggregationResults(List<AggregationResult> aggregationResults, Timestamp readTime) {
+    this(aggregationResults, readTime, null);
+  }
+
+  public AggregationResults(
+      List<AggregationResult> aggregationResults,
+      Timestamp readTime,
+      ResultSetStats resultSetStats) {
     checkNotNull(aggregationResults, "Aggregation results cannot be null");
     checkNotNull(readTime, "readTime cannot be null");
     this.aggregationResults = aggregationResults;
     this.readTime = readTime;
+    this.resultSetStats = resultSetStats;
   }
 
   /** Returns {@link Iterator} for underlying List&lt;{@link AggregationResult}&gt;. */
@@ -51,6 +62,13 @@ public int size() {
     return this.aggregationResults.size();
   }
 
+  /*
+   * Returns the ResultSetStats if QueryMode is set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, returns null.
+   */
+  public Optional<ResultSetStats> getResultSetStats() {
+    return Optional.ofNullable(this.resultSetStats);
+  }
+
   @InternalApi
   public AggregationResult get(int index) {
     return this.aggregationResults.get(index);
@@ -70,11 +88,13 @@ public boolean equals(Object o) {
       return false;
     }
     AggregationResults that = (AggregationResults) o;
-    return Objects.equals(aggregationResults, that.aggregationResults);
+    return Objects.equals(aggregationResults, that.aggregationResults)
+        && Objects.equals(readTime, that.readTime)
+        && Objects.equals(resultSetStats, that.resultSetStats);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(aggregationResults);
+    return Objects.hash(aggregationResults, readTime, resultSetStats);
   }
 }

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

@@ -16,7 +16,9 @@
 
 package com.google.cloud.datastore;
 
+import com.google.api.core.BetaApi;
 import com.google.cloud.Service;
+import com.google.cloud.datastore.models.QueryProfile;
 import com.google.datastore.v1.TransactionOptions;
 import java.util.Iterator;
 import java.util.List;
@@ -462,6 +464,30 @@ interface TransactionCallable<T> {
    */
   <T> QueryResults<T> run(Query<T> query, ReadOption... options);
 
+  /**
+   * Submits a {@link Query} with specified {@link
+   * com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns its result. {@link
+   * ReadOption}s can be specified if desired.
+   *
+   * <p>Example of running a query to find all entities of one kind.
+   *
+   * <pre>{@code
+   * String kind = "my_kind";
+   * StructuredQuery<Entity> query = Query.newEntityQueryBuilder()
+   *     .setKind(kind)
+   *     .build();
+   * QueryResults<Entity> results = datastore.run(query, QueryMode.EXPLAIN_ANALYZE);
+   * ResultSetStats resultSetStats = results.getResultSetStats();
+   * }</pre>
+   *
+   * @throws DatastoreException upon failure
+   */
+  @BetaApi
+  default <T> QueryResults<T> run(
+      Query<T> query, QueryProfile.QueryMode queryMode, ReadOption... options) {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
+
   /**
    * Submits a {@link AggregationQuery} and returns {@link AggregationResults}. {@link ReadOption}s
    * can be specified if desired.
@@ -508,4 +534,34 @@ interface TransactionCallable<T> {
   default AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
     throw new UnsupportedOperationException("Not implemented.");
   }
+
+  /**
+   * Submits a {@link AggregationQuery} with specified {@link
+   * com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns {@link
+   * AggregationResults}. {@link ReadOption}s can be specified if desired.
+   *
+   * <p>Example of running an {@link AggregationQuery} to find the count of entities of one kind.
+   *
+   * <p>{@link StructuredQuery} example:
+   *
+   * <pre>{@code
+   * EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
+   *    .setKind("Task")
+   *    .build();
+   * AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
+   *    .addAggregation(count().as("total_count"))
+   *    .over(selectAllQuery)
+   *    .build();
+   * AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery, QueryMode.EXPLAIN_ANALYZE);
+   * ResultSetStats aggregationStats = aggregationResults.getResultSetStats();
+   * }</pre>
+   *
+   * @throws DatastoreException upon failure
+   * @return {@link AggregationResults}
+   */
+  @BetaApi
+  default AggregationResults runAggregation(
+      AggregationQuery query, QueryProfile.QueryMode queryMode, ReadOption... options) {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
 }

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

@@ -16,20 +16,23 @@
 
 package com.google.cloud.datastore;
 
+import com.google.api.core.BetaApi;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.BaseService;
 import com.google.cloud.ExceptionHandler;
 import com.google.cloud.RetryHelper;
 import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.datastore.execution.AggregationQueryExecutor;
+import com.google.cloud.datastore.models.QueryProfile;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
+import com.google.datastore.v1.QueryMode;
 import com.google.datastore.v1.ReadOptions;
 import com.google.datastore.v1.ReserveIdsRequest;
 import com.google.datastore.v1.TransactionOptions;
@@ -182,28 +185,50 @@ public <T> T runInTransaction(
 
   @Override
   public <T> QueryResults<T> run(Query<T> query) {
-    return run(Optional.empty(), query);
+    return run(Optional.empty(), query, QueryMode.NORMAL);
   }
 
   @Override
   public <T> QueryResults<T> run(Query<T> query, ReadOption... options) {
-    return run(toReadOptionsPb(options), query);
+    return run(toReadOptionsPb(options), query, QueryMode.NORMAL);
+  }
+
+  @Override
+  @BetaApi
+  public <T> QueryResults<T> run(
+      Query<T> query, QueryProfile.QueryMode queryMode, ReadOption... options) {
+    return run(toReadOptionsPb(options), query, queryMode.toPb());
   }
 
   @SuppressWarnings("unchecked")
-  <T> QueryResults<T> run(Optional<ReadOptions> readOptionsPb, Query<T> query) {
+  <T> QueryResults<T> run(
+      Optional<ReadOptions> readOptionsPb, Query<T> query, QueryMode queryMode) {
     return new QueryResultsImpl<T>(
-        this, readOptionsPb, (RecordQuery<T>) query, query.getNamespace());
+        this, readOptionsPb, (RecordQuery<T>) query, query.getNamespace(), queryMode);
   }
 
   @Override
   public AggregationResults runAggregation(AggregationQuery query) {
-    return aggregationQueryExecutor.execute(query);
+    return aggregationQueryExecutor.execute(query, QueryProfile.QueryMode.NORMAL);
   }
 
   @Override
   public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
-    return aggregationQueryExecutor.execute(query, options);
+    return aggregationQueryExecutor.execute(query, QueryProfile.QueryMode.NORMAL, options);
+  }
+
+  @Override
+  @BetaApi
+  public AggregationResults runAggregation(
+      AggregationQuery query, QueryProfile.QueryMode queryMode) {
+    return aggregationQueryExecutor.execute(query, queryMode);
+  }
+
+  @Override
+  @BetaApi
+  public AggregationResults runAggregation(
+      AggregationQuery query, QueryProfile.QueryMode queryMode, ReadOption... options) {
+    return aggregationQueryExecutor.execute(query, queryMode, options);
   }
 
   com.google.datastore.v1.RunQueryResponse runQuery(

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

@@ -16,6 +16,7 @@
 
 package com.google.cloud.datastore;
 
+import com.google.cloud.datastore.models.QueryProfile;
 import java.util.Iterator;
 import java.util.List;
 
@@ -62,4 +63,16 @@ public interface DatastoreReader {
   default AggregationResults runAggregation(AggregationQuery query) {
     throw new UnsupportedOperationException("Not implemented.");
   }
+
+  /**
+   * Submits a {@link AggregationQuery} with a specified {@link
+   * com.google.cloud.datastore.models.QueryProfile.QueryMode} and returns {@link
+   * AggregationResults}.
+   *
+   * @throws DatastoreException upon failure
+   */
+  default AggregationResults runAggregation(
+      AggregationQuery query, QueryProfile.QueryMode queryMode) {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
 }

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

@@ -16,6 +16,7 @@
 
 package com.google.cloud.datastore;
 
+import com.google.api.core.InternalApi;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.MoreObjects.ToStringHelper;
 import com.google.common.collect.Maps;
@@ -117,6 +118,11 @@ public Class<V> resultClass() {
       return resultClass;
     }
 
+    @InternalApi
+    public com.google.datastore.v1.EntityResult.ResultType getQueryType() {
+      return this.queryType;
+    }
+
     @Override
     public int hashCode() {
       return resultClass.hashCode();

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

@@ -16,8 +16,11 @@
 
 package com.google.cloud.datastore;
 
+import com.google.api.core.BetaApi;
+import com.google.cloud.datastore.models.ResultSetStats;
 import com.google.datastore.v1.QueryResultBatch;
 import java.util.Iterator;
+import java.util.Optional;
 
 /**
  * The result of a Google Cloud Datastore query submission. When the result is not typed it is
@@ -70,4 +73,12 @@ public interface QueryResults<V> extends Iterator<V> {
 
   /** Returns MoreResults state of the query after the current batch. */
   QueryResultBatch.MoreResultsType getMoreResults();
+
+  /*
+   * Returns the {@link ResultSetStats} if {@link QueryMode} was set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, it will return null.
+   */
+  @BetaApi
+  default Optional<ResultSetStats> getResultSetStats() {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
 }