googleapis
diff --git a/‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuery.java
Lines changed: 182 additions & 43 deletions b/‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuery.java
Lines changed: 182 additions & 43 deletions
diff --git a/‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/ExecutionStats.java
Lines changed: 62 additions & 0 deletions b/‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/ExecutionStats.java
Lines changed: 62 additions & 0 deletions
@@ -25,6 +25,7 @@
 import com.google.api.gax.rpc.StreamController;
 import com.google.cloud.Timestamp;
 import com.google.cloud.firestore.v1.FirestoreSettings;
+import com.google.common.collect.ImmutableMap;
 import com.google.firestore.v1.RunAggregationQueryRequest;
 import com.google.firestore.v1.RunAggregationQueryResponse;
 import com.google.firestore.v1.RunQueryRequest;
@@ -40,7 +41,6 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -49,9 +49,9 @@
 public class AggregateQuery {
   @Nonnull private final Query query;
 
-  @Nonnull private List<AggregateField> aggregateFieldList;
+  @Nonnull private final List<AggregateField> aggregateFieldList;
 
-  @Nonnull private Map<String, String> aliasMap;
+  @Nonnull private final Map<String, String> aliasMap;
 
   AggregateQuery(@Nonnull Query query, @Nonnull List<AggregateField> aggregateFields) {
     this.query = query;
@@ -75,6 +75,26 @@ public ApiFuture<AggregateQuerySnapshot> get() {
     return get(null, null);
   }
 
+  /**
+   * Plans and optionally executes this query. Returns an ApiFuture that will be resolved with the
+   * planner information, statistics from the query execution (if any), and the query results (if
+   * any).
+   *
+   * @return An ApiFuture that will be resolved with the planner information, statistics from the
+   *     query execution (if any), and the query results (if any).
+   */
+  @Nonnull
+  public ApiFuture<ExplainResults<AggregateQuerySnapshot>> explain(ExplainOptions options) {
+    AggregateQueryExplainResponseDeliverer responseDeliverer =
+        new AggregateQueryExplainResponseDeliverer(
+            /* transactionId= */ null,
+            /* readTime= */ null,
+            /* startTimeNanos= */ query.rpcContext.getClock().nanoTime(),
+            /* explainOptions= */ options);
+    runQuery(responseDeliverer);
+    return responseDeliverer.getFuture();
+  }
+
   @Nonnull
   ApiFuture<AggregateQuerySnapshot> get(
       @Nullable final ByteString transactionId, @Nullable com.google.protobuf.Timestamp readTime) {
@@ -85,25 +105,34 @@ ApiFuture<AggregateQuerySnapshot> get(
     return responseDeliverer.getFuture();
   }
 
-  private void runQuery(AggregateQueryResponseDeliverer responseDeliverer) {
+  private <T> void runQuery(ResponseDeliverer<T> responseDeliverer) {
     RunAggregationQueryRequest request =
-        toProto(responseDeliverer.transactionId, responseDeliverer.readTime);
-    AggregateQueryResponseObserver responseObserver =
-        new AggregateQueryResponseObserver(responseDeliverer);
+        toProto(
+            responseDeliverer.getTransactionId(),
+            responseDeliverer.getReadTime(),
+            responseDeliverer.getExplainOptions());
+    AggregateQueryResponseObserver<T> responseObserver =
+        new AggregateQueryResponseObserver<T>(responseDeliverer);
     ServerStreamingCallable<RunAggregationQueryRequest, RunAggregationQueryResponse> callable =
         query.rpcContext.getClient().runAggregationQueryCallable();
     query.rpcContext.streamRequest(request, responseObserver, callable);
   }
 
-  private final class AggregateQueryResponseDeliverer {
+  @Nonnull
+  private Map<String, Value> convertServerAggregateFieldsMapToClientAggregateFieldsMap(
+      @Nonnull Map<String, Value> data) {
+    ImmutableMap.Builder<String, Value> builder = ImmutableMap.builder();
+    data.forEach((serverAlias, value) -> builder.put(aliasMap.get(serverAlias), value));
+    return builder.build();
+  }
 
-    @Nullable private final ByteString transactionId;
-    @Nullable private final com.google.protobuf.Timestamp readTime;
+  private abstract static class ResponseDeliverer<T> {
+    private final @Nullable ByteString transactionId;
+    private final @Nullable com.google.protobuf.Timestamp readTime;
     private final long startTimeNanos;
-    private final SettableApiFuture<AggregateQuerySnapshot> future = SettableApiFuture.create();
-    private final AtomicBoolean isFutureCompleted = new AtomicBoolean(false);
+    private final SettableApiFuture<T> future = SettableApiFuture.create();
 
-    AggregateQueryResponseDeliverer(
+    ResponseDeliverer(
         @Nullable ByteString transactionId,
         @Nullable com.google.protobuf.Timestamp readTime,
         long startTimeNanos) {
@@ -112,52 +141,148 @@ private final class AggregateQueryResponseDeliverer {
       this.startTimeNanos = startTimeNanos;
     }
 
-    ApiFuture<AggregateQuerySnapshot> getFuture() {
+    @Nullable
+    ByteString getTransactionId() {
+      return transactionId;
+    }
+
+    @Nullable
+    com.google.protobuf.Timestamp getReadTime() {
+      return readTime;
+    }
+
+    long getStartTimeNanos() {
+      return startTimeNanos;
+    }
+
+    @Nullable
+    ExplainOptions getExplainOptions() {
+      return null;
+    }
+
+    ApiFuture<T> getFuture() {
       return future;
     }
 
-    void deliverResult(@Nonnull Map<String, Value> data, Timestamp readTime) {
-      if (isFutureCompleted.compareAndSet(false, true)) {
-        Map<String, Value> mappedData = new HashMap<>();
-        data.forEach((serverAlias, value) -> mappedData.put(aliasMap.get(serverAlias), value));
-        future.set(new AggregateQuerySnapshot(AggregateQuery.this, readTime, mappedData));
-      }
+    protected void setFuture(T value) {
+      future.set(value);
     }
 
     void deliverError(Throwable throwable) {
-      if (isFutureCompleted.compareAndSet(false, true)) {
-        future.setException(throwable);
+      future.setException(throwable);
+    }
+
+    abstract void deliverResult(
+        @Nullable Map<String, Value> serverData,
+        Timestamp readTime,
+        @Nullable ExplainMetrics metrics);
+  }
+
+  private class AggregateQueryResponseDeliverer extends ResponseDeliverer<AggregateQuerySnapshot> {
+    AggregateQueryResponseDeliverer(
+        @Nullable ByteString transactionId,
+        @Nullable com.google.protobuf.Timestamp readTime,
+        long startTimeNanos) {
+      super(transactionId, readTime, startTimeNanos);
+    }
+
+    @Override
+    void deliverResult(
+        @Nullable Map<String, Value> serverData,
+        Timestamp readTime,
+        @Nullable ExplainMetrics metrics) {
+      if (serverData == null) {
+        deliverError(new RuntimeException("Did not receive any aggregate query results."));
+        return;
       }
+      setFuture(
+          new AggregateQuerySnapshot(
+              AggregateQuery.this,
+              readTime,
+              convertServerAggregateFieldsMapToClientAggregateFieldsMap(serverData)));
     }
   }
 
-  private final class AggregateQueryResponseObserver
-      implements ResponseObserver<RunAggregationQueryResponse> {
+  private final class AggregateQueryExplainResponseDeliverer
+      extends ResponseDeliverer<ExplainResults<AggregateQuerySnapshot>> {
+    private final @Nullable ExplainOptions explainOptions;
 
-    private final AggregateQueryResponseDeliverer responseDeliverer;
-    private StreamController streamController;
+    AggregateQueryExplainResponseDeliverer(
+        @Nullable ByteString transactionId,
+        @Nullable com.google.protobuf.Timestamp readTime,
+        long startTimeNanos,
+        @Nullable ExplainOptions explainOptions) {
+      super(transactionId, readTime, startTimeNanos);
+      this.explainOptions = explainOptions;
+    }
 
-    AggregateQueryResponseObserver(AggregateQueryResponseDeliverer responseDeliverer) {
-      this.responseDeliverer = responseDeliverer;
+    @Override
+    @Nullable
+    ExplainOptions getExplainOptions() {
+      return explainOptions;
     }
 
     @Override
-    public void onStart(StreamController streamController) {
-      this.streamController = streamController;
+    void deliverResult(
+        @Nullable Map<String, Value> serverData,
+        Timestamp readTime,
+        @Nullable ExplainMetrics metrics) {
+      // The server is required to provide ExplainMetrics for explain queries.
+      if (metrics == null) {
+        deliverError(new RuntimeException("Did not receive any metrics for explain query."));
+        return;
+      }
+      AggregateQuerySnapshot snapshot =
+          serverData == null
+              ? null
+              : new AggregateQuerySnapshot(
+                  AggregateQuery.this,
+                  readTime,
+                  convertServerAggregateFieldsMapToClientAggregateFieldsMap(serverData));
+      setFuture(new ExplainResults<>(metrics, snapshot));
+    }
+  }
+
+  private final class AggregateQueryResponseObserver<T>
+      implements ResponseObserver<RunAggregationQueryResponse> {
+    private final ResponseDeliverer<T> responseDeliverer;
+    private Timestamp readTime = Timestamp.MAX_VALUE;
+    @Nullable private Map<String, Value> aggregateFieldsMap = null;
+    @Nullable private ExplainMetrics metrics = null;
+
+    AggregateQueryResponseObserver(ResponseDeliverer<T> responseDeliverer) {
+      this.responseDeliverer = responseDeliverer;
     }
 
+    private boolean isExplainQuery() {
+      return this.responseDeliverer.getExplainOptions() != null;
+    }
+
+    @Override
+    public void onStart(StreamController streamController) {}
+
     @Override
     public void onResponse(RunAggregationQueryResponse response) {
-      // Close the stream to avoid it dangling, since we're not expecting any more responses.
-      streamController.cancel();
+      if (response.hasReadTime()) {
+        readTime = Timestamp.fromProto(response.getReadTime());
+      }
+
+      if (response.hasResult()) {
+        aggregateFieldsMap = response.getResult().getAggregateFieldsMap();
+      }
 
-      // Extract the aggregations and read time from the RunAggregationQueryResponse.
-      Timestamp readTime = Timestamp.fromProto(response.getReadTime());
+      if (response.hasExplainMetrics()) {
+        metrics = new ExplainMetrics(response.getExplainMetrics());
+      }
 
-      // Deliver the result; even though the `RunAggregationQuery` RPC is a "streaming" RPC, meaning
-      // that `onResponse()` can be called multiple times, it _should_ only be called once. But even
-      // if it is called more than once, `responseDeliverer` will drop superfluous results.
-      responseDeliverer.deliverResult(response.getResult().getAggregateFieldsMap(), readTime);
+      if (!isExplainQuery()) {
+        // Deliver the result; even though the `RunAggregationQuery` RPC is a "streaming" RPC,
+        // meaning that `onResponse()` can be called multiple times, it _should_ only be called
+        // once for non-explain queries. But even if it is called more than once,
+        // `responseDeliverer` will drop superfluous results. For explain queries, there will
+        // be more than one response, and the last response will contain the metrics.
+        onComplete();
+      }
     }
 
     @Override
@@ -170,17 +295,26 @@ public void onError(Throwable throwable) {
     }
 
     private boolean shouldRetry(Throwable throwable) {
+      // Do not retry EXPLAIN requests because it'd be executing
+      // multiple queries. This means stats would have to be aggregated,
+      // and that may not even make sense for many statistics.
+      if (isExplainQuery()) {
+        return false;
+      }
+
       Set<StatusCode.Code> retryableCodes =
           FirestoreSettings.newBuilder().runAggregationQuerySettings().getRetryableCodes();
       return query.shouldRetryQuery(
           throwable,
-          responseDeliverer.transactionId,
-          responseDeliverer.startTimeNanos,
+          responseDeliverer.getTransactionId(),
+          responseDeliverer.getStartTimeNanos(),
           retryableCodes);
     }
 
     @Override
-    public void onComplete() {}
+    public void onComplete() {
+      responseDeliverer.deliverResult(aggregateFieldsMap, readTime, metrics);
+    }
   }
 
   /**
@@ -191,13 +325,14 @@ public void onComplete() {}
    */
   @Nonnull
   public RunAggregationQueryRequest toProto() {
-    return toProto(null, null);
+    return toProto(/* transactionId= */ null, /* readTime= */ null, /* explainOptions= */ null);
   }
 
   @Nonnull
   RunAggregationQueryRequest toProto(
       @Nullable final ByteString transactionId,
-      @Nullable final com.google.protobuf.Timestamp readTime) {
+      @Nullable final com.google.protobuf.Timestamp readTime,
+      @Nullable ExplainOptions explainOptions) {
     RunQueryRequest runQueryRequest = query.toProto();
 
     RunAggregationQueryRequest.Builder request = RunAggregationQueryRequest.newBuilder();
@@ -209,6 +344,10 @@ RunAggregationQueryRequest toProto(
       request.setReadTime(readTime);
     }
 
+    if (explainOptions != null) {
+      request.setExplainOptions(explainOptions.toProto());
+    }
+
     StructuredAggregationQuery.Builder structuredAggregationQuery =
         request.getStructuredAggregationQueryBuilder();
     structuredAggregationQuery.setStructuredQuery(runQueryRequest.getStructuredQuery());
 
@@ -0,0 +1,62 @@
+/*
+ * Copyright 2024 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
+ *
+ *       http://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.firestore;
+
+import java.time.Duration;
+import java.util.Map;
+import javax.annotation.Nonnull;
+
+/** A ExecutionStats contains information about the execution of a query. */
+public final class ExecutionStats {
+  private final long resultsReturned;
+  private final @Nonnull Duration executionDuration;
+  private final long readOperations;
+  private final @Nonnull Map<String, Object> debugStats;
+
+  ExecutionStats(com.google.firestore.v1.ExecutionStats proto) {
+    this.resultsReturned = proto.getResultsReturned();
+    this.executionDuration =
+        Duration.ofSeconds(
+            proto.getExecutionDuration().getSeconds(), proto.getExecutionDuration().getNanos());
+    this.readOperations = proto.getReadOperations();
+    this.debugStats = UserDataConverter.decodeStruct(proto.getDebugStats());
+  }
+
+  /** Returns the number of query results. */
+  public long getResultsReturned() {
+    return resultsReturned;
+  }
+
+  /** Returns the total execution time of the query. */
+  @Nonnull
+  public Duration getExecutionDuration() {
+    return executionDuration;
+  }
+
+  /** Returns the number of read operations that occurred when executing the query. */
+  public long getReadOperations() {
+    return readOperations;
+  }
+
+  /**
+   * Returns a map that contains additional statistics related to query execution. Note: The content
+   * of this map are subject to change.
+   */
+  @Nonnull
+  public Map<String, Object> getDebugStats() {
+    return debugStats;
+  }
+}