feat: implement query profile (#1365)

* feat: implement query profile * javadoc fixes * add betaapi in a few more places * add some extra tests * add some more ITs * add some more fine grained tests * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · Mar 25, 2024 · 2515ed6 · 2515ed6
1 parent 9818aee
commit 2515ed6
Show file tree

Hide file tree

Showing 26 changed files with 1,071 additions and 48 deletions.
diff --git a/google-cloud-datastore/clirr-ignored-differences.xml b/google-cloud-datastore/clirr-ignored-differences.xml
@@ -1,16 +1,58 @@
 <?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.ExplainOptions, 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.ExplainOptions, com.google.cloud.datastore.ReadOption[])</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.ExplainOptions)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/QueryResults</className>
+        <method>java.util.Optional getExplainMetrics()</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.ExplainOptions)</method>
+        <differenceType>7012</differenceType>
+    </difference>
+
+    <!-- changed parmeters on internal only class-->
+    <difference>
+        <className>com/google/cloud/datastore/ReadOption$QueryConfig</className>
+        <method>com.google.cloud.datastore.ReadOption$QueryConfig create(com.google.cloud.datastore.Query, java.util.List)</method>
+        <to>*com.google.datastore.v1.ExplainOptions*</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>
+    <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>
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResults.java
@@ -17,11 +17,14 @@
 
 import static com.google.api.client.util.Preconditions.checkNotNull;
 
+import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.cloud.Timestamp;
+import com.google.cloud.datastore.models.ExplainMetrics;
 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&lt;{@link
@@ -34,11 +37,22 @@ public class AggregationResults implements Iterable<AggregationResult> {
   private final List<AggregationResult> aggregationResults;
   private final Timestamp readTime;
 
-  public AggregationResults(List<AggregationResult> aggregationResults, Timestamp readTime) {
+  private final ExplainMetrics explainMetrics;
+
+  @BetaApi
+  public AggregationResults(
+      List<AggregationResult> aggregationResults,
+      Timestamp readTime,
+      ExplainMetrics explainMetrics) {
     checkNotNull(aggregationResults, "Aggregation results cannot be null");
     checkNotNull(readTime, "readTime cannot be null");
     this.aggregationResults = aggregationResults;
     this.readTime = readTime;
+    this.explainMetrics = explainMetrics;
+  }
+
+  public AggregationResults(List<AggregationResult> aggregationResults, Timestamp readTime) {
+    this(aggregationResults, readTime, null);
   }
 
   /** Returns {@link Iterator} for underlying List&lt;{@link AggregationResult}&gt;. */
@@ -51,6 +65,14 @@ public int size() {
     return this.aggregationResults.size();
   }
 
+  /**
+   * Returns {@code ExplainMetrics} if {@code ExplainOptions} were enabled. Otherwise, returns null.
+   */
+  @BetaApi
+  public Optional<ExplainMetrics> getExplainMetrics() {
+    return Optional.ofNullable(this.explainMetrics);
+  }
+
   @InternalApi
   public AggregationResult get(int index) {
     return this.aggregationResults.get(index);
@@ -70,11 +92,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(explainMetrics, that.explainMetrics);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(aggregationResults);
+    return Objects.hash(aggregationResults, readTime, explainMetrics);
   }
 }
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/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.ExplainOptions;
 import com.google.datastore.v1.TransactionOptions;
 import java.util.Iterator;
 import java.util.List;
@@ -462,6 +464,28 @@ interface TransactionCallable<T> {
    */
   <T> QueryResults<T> run(Query<T> query, ReadOption... options);
 
+  /**
+   * Submits a {@link Query} with specified {@link com.google.cloud.datastore.models.ExplainOptions}
+   * 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, ExplainOptions.newBuilder().setAnalyze(true).build());
+   * }</pre>
+   *
+   * @throws DatastoreException upon failure
+   */
+  @BetaApi
+  default <T> QueryResults<T> run(
+      Query<T> query, ExplainOptions explainOptions, 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 +532,33 @@ 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.ExplainOptions} 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, ExplainOptions.newBuilder().setAnalyze(true).build());
+   * }</pre>
+   *
+   * @throws DatastoreException upon failure
+   * @return {@link AggregationResults}
+   */
+  @BetaApi
+  default AggregationResults runAggregation(
+      AggregationQuery query, ExplainOptions explainOptions, ReadOption... options) {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
 }
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
@@ -16,6 +16,7 @@
 
 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;
@@ -30,6 +31,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
+import com.google.datastore.v1.ExplainOptions;
 import com.google.datastore.v1.ReadOptions;
 import com.google.datastore.v1.ReserveIdsRequest;
 import com.google.datastore.v1.TransactionOptions;
@@ -182,28 +184,54 @@ public <T> T runInTransaction(
 
   @Override
   public <T> QueryResults<T> run(Query<T> query) {
-    return run(Optional.empty(), query);
+    return run(Optional.empty(), query, null);
   }
 
   @Override
   public <T> QueryResults<T> run(Query<T> query, ReadOption... options) {
-    return run(toReadOptionsPb(options), query);
+    return run(toReadOptionsPb(options), query, null);
+  }
+
+  @Override
+  @BetaApi
+  public <T> QueryResults<T> run(
+      Query<T> query,
+      com.google.cloud.datastore.models.ExplainOptions explainOptions,
+      ReadOption... options) {
+    return run(toReadOptionsPb(options), query, explainOptions.toPb());
   }
 
   @SuppressWarnings("unchecked")
-  <T> QueryResults<T> run(Optional<ReadOptions> readOptionsPb, Query<T> query) {
+  <T> QueryResults<T> run(
+      Optional<ReadOptions> readOptionsPb, Query<T> query, ExplainOptions explainOptions) {
     return new QueryResultsImpl<T>(
-        this, readOptionsPb, (RecordQuery<T>) query, query.getNamespace());
+        this, readOptionsPb, (RecordQuery<T>) query, query.getNamespace(), explainOptions);
   }
 
   @Override
   public AggregationResults runAggregation(AggregationQuery query) {
-    return aggregationQueryExecutor.execute(query);
+    return aggregationQueryExecutor.execute(query, null);
   }
 
   @Override
   public AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
-    return aggregationQueryExecutor.execute(query, options);
+    return aggregationQueryExecutor.execute(query, null, options);
+  }
+
+  @Override
+  @BetaApi
+  public AggregationResults runAggregation(
+      AggregationQuery query, com.google.cloud.datastore.models.ExplainOptions explainOptions) {
+    return aggregationQueryExecutor.execute(query, explainOptions);
+  }
+
+  @Override
+  @BetaApi
+  public AggregationResults runAggregation(
+      AggregationQuery query,
+      com.google.cloud.datastore.models.ExplainOptions explainOptions,
+      ReadOption... options) {
+    return aggregationQueryExecutor.execute(query, explainOptions, options);
   }
 
   com.google.datastore.v1.RunQueryResponse runQuery(

diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
@@ -16,6 +16,8 @@
 
 package com.google.cloud.datastore;
 
+import com.google.api.core.BetaApi;
+import com.google.cloud.datastore.models.ExplainOptions;
 import java.util.Iterator;
 import java.util.List;
 
@@ -62,4 +64,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.ExplainOptions} and returns {@link AggregationResults}.
+   *
+   * @throws DatastoreException upon failure
+   */
+  @BetaApi
+  default AggregationResults runAggregation(AggregationQuery query, ExplainOptions explainOptions) {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
 }
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
@@ -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();

diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/QueryResults.java b/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.ExplainMetrics;
 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,9 @@ public interface QueryResults<V> extends Iterator<V> {
 
   /** Returns MoreResults state of the query after the current batch. */
   QueryResultBatch.MoreResultsType getMoreResults();
+
+  @BetaApi
+  default Optional<ExplainMetrics> getExplainMetrics() {
+    throw new UnsupportedOperationException("Not implemented.");
+  }
 }