feat: query profiling

googleapis · kolea2 · Mar 25, 2024 · Dec 13, 2023 · Jan 4, 2024 · Jan 18, 2024
commit 9d0a267e25b1086cfab04eefc5d2ae112cdb31aa
@@ -13,4 +13,60 @@
  <to>*QueryConfig*</to>
  <differenceType>7005</differenceType>
  </difference>
+
+ <!-- added a default method to an interface-->
+ <difference>
+ <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/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>com.google.cloud.datastore.models.ResultSetStats 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>
 </differences>
@@ -19,6 +19,7 @@
 
 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;
@@ -33,12 +34,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 +61,13 @@ public int size() {
  return this.aggregationResults.size();
  }
 
+ /*
+ * Returns the ResultSetStats if QueryMode is set to EXPLAIN or EXPLAIN_ANALYZE. Otherwise, returns null.
+ */
+ public ResultSetStats getResultSetStats() {
+ return this.resultSetStats;
+ }
+
  @InternalApi
  public AggregationResult get(int index) {
  return this.aggregationResults.get(index);
@@ -70,11 +87,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);
  }
 }
@@ -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,29 @@ 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 +533,32 @@ 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.");
+ }
 }
@@ -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(

@@ -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,15 @@ 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.");
+ }
 }
@@ -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();

@@ -16,6 +16,8 @@
 
 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;
 
@@ -70,4 +72,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 ResultSetStats getResultSetStats() {
+ throw new UnsupportedOperationException("not implemented");
+ }
 }