diff --git a/firebase-firestore/src/androidTest/java/com/google/firebase/firestore/CountTest.java b/firebase-firestore/src/androidTest/java/com/google/firebase/firestore/CountTest.java new file mode 100644 index 00000000000..78cf1e94fe9 --- /dev/null +++ b/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())); + } +} diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateField.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateField.java new file mode 100644 index 00000000000..7406d2fd5dc --- /dev/null +++ b/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() {} + } +} diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuery.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuery.java new file mode 100644 index 00000000000..1c6820705a9 --- /dev/null +++ b/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 get() { + Datastore datastore = query.firestore.getClient().getDatastore(); + TaskCompletionSource 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(); + } +} diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuerySnapshot.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/AggregateQuerySnapshot.java new file mode 100644 index 00000000000..dceff564bea --- /dev/null +++ b/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; + } +} diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/Query.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/Query.java index 9978ec0b3e5..61e5e6681d2 100644 --- a/firebase-firestore/src/main/java/com/google/firebase/firestore/Query.java +++ b/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) { diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java index 464cd5913f9..20d5d96da6c 100644 --- a/firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java +++ b/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 getNamedQuery(String queryName) { verifyNotTerminated(); TaskCompletionSource completionSource = new TaskCompletionSource<>(); diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/remote/Datastore.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/remote/Datastore.java index 46323e881b9..ffdce24be93 100644 --- a/firebase-firestore/src/main/java/com/google/firebase/firestore/remote/Datastore.java +++ b/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 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 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. diff --git a/firebase-firestore/src/proto/google/firestore/v1/aggregation_result.proto b/firebase-firestore/src/proto/google/firestore/v1/aggregation_result.proto new file mode 100644 index 00000000000..538e3fef5e4 --- /dev/null +++ b/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 aggregate_fields = 2; +} diff --git a/firebase-firestore/src/proto/google/firestore/v1/firestore.proto b/firebase-firestore/src/proto/google/firestore/v1/firestore.proto index d425edf9e0f..dda4721596c 100644 --- a/firebase-firestore/src/proto/google/firestore/v1/firestore.proto +++ b/firebase-firestore/src/proto/google/firestore/v1/firestore.proto @@ -18,6 +18,7 @@ syntax = "proto3"; package google.firestore.v1; import "google/api/annotations.proto"; +import "google/firestore/v1/aggregation_result.proto"; import "google/firestore/v1/common.proto"; import "google/firestore/v1/document.proto"; import "google/firestore/v1/query.proto"; @@ -136,6 +137,29 @@ service Firestore { }; } + // Runs an aggregation query. + // + // Rather than producing [Document][google.firestore.v1.Document] results like [Firestore.RunQuery][google.firestore.v1.Firestore.RunQuery], + // this API allows running an aggregation to produce a series of + // [AggregationResult][google.firestore.v1.AggregationResult] server-side. + // + // High-Level Example: + // + // ``` + // -- Return the number of documents in table given a filter. + // SELECT COUNT(*) FROM ( SELECT * FROM k where a = true ); + // ``` + rpc RunAggregationQuery(RunAggregationQueryRequest) returns (stream RunAggregationQueryResponse) { + option (google.api.http) = { + post: "/v1/{parent=projects/*/databases/*/documents}:runAggregationQuery" + body: "*" + additional_bindings { + post: "/v1/{parent=projects/*/databases/*/documents/*/**}:runAggregationQuery" + body: "*" + } + }; + } + // Streams batches of document updates and deletes, in order. rpc Write(stream WriteRequest) returns (stream WriteResponse) { option (google.api.http) = { @@ -485,6 +509,62 @@ message RunQueryResponse { int32 skipped_results = 4; } +// The request for [Firestore.RunAggregationQuery][google.firestore.v1.Firestore.RunAggregationQuery]. +message RunAggregationQueryRequest { + // Required. The parent resource name. In the format: + // `projects/{project_id}/databases/{database_id}/documents` or + // `projects/{project_id}/databases/{database_id}/documents/{document_path}`. + // For example: + // `projects/my-project/databases/my-database/documents` or + // `projects/my-project/databases/my-database/documents/chatrooms/my-chatroom` + string parent = 1; + + // The query to run. + oneof query_type { + // An aggregation query. + StructuredAggregationQuery structured_aggregation_query = 2; + } + + // The consistency mode for the query, defaults to strong consistency. + oneof consistency_selector { + // Run the aggregation within an already active transaction. + // + // The value here is the opaque transaction ID to execute the query in. + bytes transaction = 4; + + // Starts a new transaction as part of the query, defaulting to read-only. + // + // The new transaction ID will be returned as the first response in the + // stream. + TransactionOptions new_transaction = 5; + + // Executes the query at the given timestamp. + // + // Requires: + // + // * Cannot be more than 270 seconds in the past. + google.protobuf.Timestamp read_time = 6; + } +} + +// The response for [Firestore.RunAggregationQuery][google.firestore.v1.Firestore.RunAggregationQuery]. +message RunAggregationQueryResponse { + // A single aggregation result. + // + // Not present when reporting partial progress or when the query produced + // zero results. + AggregationResult result = 1; + + // The transaction that was started as part of this request. + // + // Only present on the first response when the request requested to start + // a new transaction. + bytes transaction = 2; + + // The time at which the aggregate value is valid for. + google.protobuf.Timestamp read_time = 3; +} + // The request for [Firestore.Write][google.firestore.v1.Firestore.Write]. // // The first request creates a stream, or resumes an existing one from a token. diff --git a/firebase-firestore/src/proto/google/firestore/v1/query.proto b/firebase-firestore/src/proto/google/firestore/v1/query.proto index 25d53238130..d76642c4c21 100644 --- a/firebase-firestore/src/proto/google/firestore/v1/query.proto +++ b/firebase-firestore/src/proto/google/firestore/v1/query.proto @@ -288,6 +288,57 @@ message StructuredQuery { google.protobuf.Int32Value limit = 5; } +message StructuredAggregationQuery { + // Defines a aggregation that produces a single result. + message Aggregation { + // Count of documents that match the query. + // + // The `COUNT(*)` aggregation function operates on the entire document + // so it does not require a field reference. + message Count { + // Optional. Optional constraint on the maximum number of documents to count. + // + // This provides a way to set an upper bound on the number of documents + // to scan, limiting latency and cost. + // + // High-Level Example: + // + // ``` + // SELECT COUNT_UP_TO(1000) FROM ( SELECT * FROM k ); + // ``` + // + // Requires: + // + // * Must be greater than zero when present. + int32 up_to = 1; + } + + // The type of aggregation to perform, required. + oneof operator { + // Count aggregator. + Count count = 1; + } + + // Required. The name of the field to store the result of the aggregation into. + // + // Requires: + // + // * Must be present. + // * Must be unique across all aggregation aliases. + // * Conform to existing [document field name][google.firestore.v1.Document.fields] limitations. + string alias = 7; + } + + // The base query to aggregate over. + oneof query_type { + // Nested structured query. + StructuredQuery structured_query = 1; + } + + // Optional. Series of aggregations to apply on top of the `structured_query`. + repeated Aggregation aggregations = 3; +} + // A position in a query result set. message Cursor { // The values that represent a position, in the order they appear in