DO NOT MERGE: Prototype COUNT using the RunAggregationQuery RPC

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/CountTest.java

@@ -0,0 +1,48 @@
+// 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
+//
+//      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.firebase.firestore;
+
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollectionWithDocs;
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
+import static com.google.firebase.firestore.testutil.TestUtil.map;
+import static org.junit.Assert.assertEquals;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.firebase.firestore.testutil.IntegrationTestUtil;
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+public class CountTest {
+
+  @After
+  public void tearDown() {
+    IntegrationTestUtil.tearDown();
+  }
+
+  @Test
+  public void count() {
+    CollectionReference collection =
+        testCollectionWithDocs(
+            map(
+                "a", map("k", "a"),
+                "b", map("k", "b"),
+                "c", map("k", "c")));
+
+    AggregateQuerySnapshot snapshot = waitFor(collection.count().get());
+    assertEquals(Long.valueOf(3), snapshot.get(AggregateField.count()));
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateField.java

@@ -0,0 +1,31 @@
+// 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
+//
+//      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.firebase.firestore;
+
+import androidx.annotation.NonNull;
+
+public abstract class AggregateField {
+
+  private AggregateField() {}
+
+  @NonNull
+  public static CountAggregateField count() {
+    return new CountAggregateField();
+  }
+
+  public static final class CountAggregateField extends AggregateField {
+    CountAggregateField() {}
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuery.java

@@ -0,0 +1,59 @@
+// 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
+//
+//      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.firebase.firestore;
+
+import androidx.annotation.NonNull;
+import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.TaskCompletionSource;
+import com.google.firebase.firestore.remote.Datastore;
+import com.google.firebase.firestore.util.Executors;
+
+public final class AggregateQuery {
+
+  private final Query query;
+
+  AggregateQuery(@NonNull Query query, @NonNull AggregateField aggregateField) {
+    this.query = query;
+    if (!(aggregateField instanceof AggregateField.CountAggregateField)) {
+      throw new IllegalArgumentException("unsupported aggregateField: " + aggregateField);
+    }
+  }
+
+  @NonNull
+  public Query getQuery() {
+    return query;
+  }
+
+  @NonNull
+  public Task<AggregateQuerySnapshot> get() {
+    Datastore datastore = query.firestore.getClient().getDatastore();
+    TaskCompletionSource<AggregateQuerySnapshot> tcs = new TaskCompletionSource<>();
+
+    datastore
+        .runCountQuery(query.query.toTarget())
+        .continueWith(
+            Executors.DIRECT_EXECUTOR,
+            task -> {
+              if (task.isSuccessful()) {
+                tcs.setResult(new AggregateQuerySnapshot(task.getResult()));
+              } else {
+                tcs.setException(task.getException());
+              }
+              return null;
+            });
+
+    return tcs.getTask();
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuerySnapshot.java

@@ -0,0 +1,32 @@
+// 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
+//
+//      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.firebase.firestore;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+public class AggregateQuerySnapshot {
+
+  private final long count;
+
+  AggregateQuerySnapshot(long count) {
+    this.count = count;
+  }
+
+  @Nullable
+  public Long get(@NonNull AggregateField.CountAggregateField field) {
+    return count;
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/Query.java

@@ -1223,6 +1223,11 @@ private void validateHasExplicitOrderByForLimitToLast() {
     }
   }
 
+  @NonNull
+  public AggregateQuery count() {
+    return new AggregateQuery(this, AggregateField.count());
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java

@@ -70,6 +70,7 @@ public final class FirestoreClient {
   private final BundleSerializer bundleSerializer;
   private final GrpcMetadataProvider metadataProvider;
 
+  private Datastore datastore;
   private Persistence persistence;
   private LocalStore localStore;
   private RemoteStore remoteStore;
@@ -255,7 +256,7 @@ private void initialize(Context context, User user, FirebaseFirestoreSettings se
     // completes.
     Logger.debug(LOG_TAG, "Initializing. user=%s", user.getUid());
 
-    Datastore datastore =
+    this.datastore =
         new Datastore(
             databaseInfo, asyncQueue, authProvider, appCheckProvider, context, metadataProvider);
     ComponentProvider.Configuration configuration =
@@ -302,6 +303,10 @@ public void loadBundle(InputStream bundleData, LoadBundleTask resultTask) {
     asyncQueue.enqueueAndForget(() -> syncEngine.loadBundle(bundleReader, resultTask));
   }
 
+  public Datastore getDatastore() {
+    return datastore;
+  }
+
   public Task<Query> getNamedQuery(String queryName) {
     verifyNotTerminated();
     TaskCompletionSource<Query> completionSource = new TaskCompletionSource<>();

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/Datastore.java

@@ -14,6 +14,7 @@
 
 package com.google.firebase.firestore.remote;
 
+import static com.google.firebase.firestore.util.Assert.hardAssert;
 import static com.google.firebase.firestore.util.Util.exceptionFromStatus;
 
 import android.content.Context;
@@ -25,17 +26,23 @@
 import com.google.firebase.firestore.auth.CredentialsProvider;
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.DatabaseInfo;
+import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.MutationResult;
 import com.google.firebase.firestore.util.AsyncQueue;
+import com.google.firestore.v1.AggregationResult;
 import com.google.firestore.v1.BatchGetDocumentsRequest;
 import com.google.firestore.v1.BatchGetDocumentsResponse;
 import com.google.firestore.v1.CommitRequest;
 import com.google.firestore.v1.CommitResponse;
 import com.google.firestore.v1.FirestoreGrpc;
+import com.google.firestore.v1.RunAggregationQueryRequest;
+import com.google.firestore.v1.RunAggregationQueryResponse;
+import com.google.firestore.v1.StructuredAggregationQuery;
+import com.google.firestore.v1.Value;
 import io.grpc.Status;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -215,6 +222,53 @@ public void onClose(Status status) {
     return completionSource.getTask();
   }
 
+  public Task<Long> runCountQuery(Target queryTarget) {
+    com.google.firestore.v1.Target.QueryTarget encodedQueryTarget =
+        serializer.encodeQueryTarget(queryTarget);
+
+    StructuredAggregationQuery.Builder structuredAggregationQuery =
+        StructuredAggregationQuery.newBuilder();
+    structuredAggregationQuery.setStructuredQuery(encodedQueryTarget.getStructuredQuery());
+
+    StructuredAggregationQuery.Aggregation.Builder aggregation =
+        StructuredAggregationQuery.Aggregation.newBuilder();
+    aggregation.setCount(StructuredAggregationQuery.Aggregation.Count.getDefaultInstance());
+    aggregation.setAlias("zzyzx_agg_alias_count");
+    structuredAggregationQuery.addAggregations(aggregation);
+
+    RunAggregationQueryRequest.Builder request = RunAggregationQueryRequest.newBuilder();
+    request.setParent(encodedQueryTarget.getParent());
+    request.setStructuredAggregationQuery(structuredAggregationQuery);
+
+    return channel
+        .runRpc(FirestoreGrpc.getRunAggregationQueryMethod(), request.build())
+        .continueWith(
+            workerQueue.getExecutor(),
+            task -> {
+              if (!task.isSuccessful()) {
+                if (task.getException() instanceof FirebaseFirestoreException
+                    && ((FirebaseFirestoreException) task.getException()).getCode()
+                        == FirebaseFirestoreException.Code.UNAUTHENTICATED) {
+                  channel.invalidateToken();
+                }
+                throw task.getException();
+              }
+              RunAggregationQueryResponse response = task.getResult();
+
+              AggregationResult aggregationResult = response.getResult();
+              Map<String, Value> aggregateFieldsByAlias = aggregationResult.getAggregateFieldsMap();
+              hardAssert(
+                  aggregateFieldsByAlias.size() == 1,
+                  "aggregateFieldsByAlias.size()==" + aggregateFieldsByAlias.size());
+              Value countValue = aggregateFieldsByAlias.get("zzyzx_agg_alias_count");
+              hardAssert(countValue != null, "countValue == null");
+              hardAssert(
+                  countValue.getValueTypeCase() == Value.ValueTypeCase.INTEGER_VALUE,
+                  "countValue.getValueTypeCase() == " + countValue.getValueTypeCase());
+              return countValue.getIntegerValue();
+            });
+  }
+
   /**
    * Determines whether the given status has an error code that represents a permanent error when
    * received in response to a non-write operation.

firebase-firestore/src/proto/google/firestore/v1/aggregation_result.proto

@@ -0,0 +1,42 @@
+// 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
+//
+//     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.
+
+syntax = "proto3";
+
+package google.firestore.v1;
+
+import "google/firestore/v1/document.proto";
+
+option csharp_namespace = "Google.Cloud.Firestore.V1";
+option go_package = "google.golang.org/genproto/googleapis/firestore/v1;firestore";
+option java_multiple_files = true;
+option java_outer_classname = "AggregationResultProto";
+option java_package = "com.google.firestore.v1";
+option objc_class_prefix = "GCFS";
+option php_namespace = "Google\\Cloud\\Firestore\\V1";
+option ruby_package = "Google::Cloud::Firestore::V1";
+
+// The result of a single bucket from a Firestore aggregation query.
+//
+// The keys of `aggregate_fields` are the same for all results in an aggregation
+// query, unlike document queries which can have different fields present for
+// each result.
+message AggregationResult {
+  // The result of the aggregation functions, ex: `COUNT(*) AS total_docs`.
+  //
+  // The key is the [alias][google.firestore.v1.StructuredAggregationQuery.Aggregation.alias]
+  // assigned to the aggregation function on input and the size of this map
+  // equals the number of aggregation functions in the query.
+  map<string, Value> aggregate_fields = 2;
+}