diff --git a/.changeset/hot-insects-wink.md b/.changeset/hot-insects-wink.md new file mode 100644 index 00000000000..30c0bc3eebc --- /dev/null +++ b/.changeset/hot-insects-wink.md @@ -0,0 +1,5 @@ +--- +'@firebase/firestore': minor +--- + +Implement count query for internal use. diff --git a/packages/firestore/src/api/aggregate.ts b/packages/firestore/src/api/aggregate.ts new file mode 100644 index 00000000000..501b5437573 --- /dev/null +++ b/packages/firestore/src/api/aggregate.ts @@ -0,0 +1,39 @@ +/** + * @license + * 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. + */ + +import { Query } from '../api'; +import { firestoreClientRunCountQuery } from '../core/firestore_client'; +import { AggregateField, AggregateQuerySnapshot } from '../lite-api/aggregate'; +import { cast } from '../util/input_validation'; + +import { ensureFirestoreConfigured, Firestore } from './database'; + +/** + * Executes the query and returns the results as a `AggregateQuerySnapshot` from the + * server. Returns an error if the network is not available. + * + * @param query - The `Query` to execute. + * + * @returns A `Promise` that will be resolved with the results of the query. + */ +export function getCountFromServer( + query: Query +): Promise }>> { + const firestore = cast(query.firestore, Firestore); + const client = ensureFirestoreConfigured(firestore); + return firestoreClientRunCountQuery(client, query); +} diff --git a/packages/firestore/src/core/firestore_client.ts b/packages/firestore/src/core/firestore_client.ts index 68cb92b9a73..6b9e410c53c 100644 --- a/packages/firestore/src/core/firestore_client.ts +++ b/packages/firestore/src/core/firestore_client.ts @@ -23,6 +23,12 @@ import { CredentialsProvider } from '../api/credentials'; import { User } from '../auth/user'; +import { + AggregateField, + AggregateQuerySnapshot, + getCount +} from '../lite-api/aggregate'; +import { Query as LiteQuery } from '../lite-api/reference'; import { LocalStore } from '../local/local_store'; import { localStoreExecuteQuery, @@ -38,6 +44,7 @@ import { toByteStreamReader } from '../platform/byte_stream_reader'; import { newSerializer, newTextEncoder } from '../platform/serializer'; import { Datastore } from '../remote/datastore'; import { + canUseNetwork, RemoteStore, remoteStoreDisableNetwork, remoteStoreEnableNetwork, @@ -501,6 +508,34 @@ export function firestoreClientTransaction( return deferred.promise; } +export function firestoreClientRunCountQuery( + client: FirestoreClient, + query: LiteQuery +): Promise }>> { + const deferred = new Deferred< + AggregateQuerySnapshot<{ count: AggregateField }> + >(); + client.asyncQueue.enqueueAndForget(async () => { + try { + const remoteStore = await getRemoteStore(client); + if (!canUseNetwork(remoteStore)) { + deferred.reject( + new FirestoreError( + Code.UNAVAILABLE, + 'Failed to get count result because the client is offline.' + ) + ); + } else { + const result = await getCount(query); + deferred.resolve(result); + } + } catch (e) { + deferred.reject(e as Error); + } + }); + return deferred.promise; +} + async function readDocumentFromCache( localStore: LocalStore, docKey: DocumentKey, diff --git a/packages/firestore/src/lite-api/aggregate.ts b/packages/firestore/src/lite-api/aggregate.ts new file mode 100644 index 00000000000..0fb0c15da94 --- /dev/null +++ b/packages/firestore/src/lite-api/aggregate.ts @@ -0,0 +1,153 @@ +/** + * @license + * 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. + */ + +import { deepEqual } from '@firebase/util'; + +import { Value } from '../protos/firestore_proto_api'; +import { invokeRunAggregationQueryRpc } from '../remote/datastore'; +import { hardAssert } from '../util/assert'; +import { cast } from '../util/input_validation'; + +import { getDatastore } from './components'; +import { Firestore } from './database'; +import { Query, queryEqual } from './reference'; +import { LiteUserDataWriter } from './reference_impl'; + +/** + * An `AggregateField`that captures input type T. + */ +// eslint-disable-next-line @typescript-eslint/no-unused-vars +export class AggregateField { + type = 'AggregateField'; +} + +/** + * Creates and returns an aggregation field that counts the documents in the result set. + * @returns An `AggregateField` object with number input type. + */ +export function count(): AggregateField { + return new AggregateField(); +} + +/** + * The union of all `AggregateField` types that are returned from the factory + * functions. + */ +export type AggregateFieldType = ReturnType; + +/** + * A type whose values are all `AggregateField` objects. + * This is used as an argument to the "getter" functions, and the snapshot will + * map the same names to the corresponding values. + */ +export interface AggregateSpec { + [field: string]: AggregateFieldType; +} + +/** + * A type whose keys are taken from an `AggregateSpec` type, and whose values + * are the result of the aggregation performed by the corresponding + * `AggregateField` from the input `AggregateSpec`. + */ +export type AggregateSpecData = { + [P in keyof T]: T[P] extends AggregateField ? U : never; +}; + +/** + * An `AggregateQuerySnapshot` contains the results of running an aggregate query. + */ +export class AggregateQuerySnapshot { + readonly type = 'AggregateQuerySnapshot'; + + /** @hideconstructor */ + constructor( + readonly query: Query, + private readonly _data: AggregateSpecData + ) {} + + /** + * The results of the requested aggregations. The keys of the returned object + * will be the same as those of the `AggregateSpec` object specified to the + * aggregation method, and the values will be the corresponding aggregation + * result. + * + * @returns The aggregation statistics result of running a query. + */ + data(): AggregateSpecData { + return this._data; + } +} + +/** + * Counts the number of documents in the result set of the given query, ignoring + * any locally-cached data and any locally-pending writes and simply surfacing + * whatever the server returns. If the server cannot be reached then the + * returned promise will be rejected. + * + * @param query - The `Query` to execute. + * + * @returns An `AggregateQuerySnapshot` that contains the number of documents. + */ +export function getCount( + query: Query +): Promise }>> { + const firestore = cast(query.firestore, Firestore); + const datastore = getDatastore(firestore); + const userDataWriter = new LiteUserDataWriter(firestore); + return invokeRunAggregationQueryRpc(datastore, query._query).then(result => { + hardAssert( + result[0] !== undefined, + 'Aggregation fields are missing from result.' + ); + + const counts = Object.entries(result[0]) + .filter(([key, value]) => key === 'count_alias') + .map(([key, value]) => userDataWriter.convertValue(value as Value)); + + const countValue = counts[0]; + + hardAssert( + typeof countValue === 'number', + 'Count aggregate field value is not a number: ' + countValue + ); + + return Promise.resolve( + new AggregateQuerySnapshot<{ count: AggregateField }>(query, { + count: countValue + }) + ); + }); +} + +/** + * Compares two `AggregateQuerySnapshot` instances for equality. + * Two `AggregateQuerySnapshot` instances are considered "equal" if they have + * the same underlying query, and the same data. + * + * @param left - The `AggregateQuerySnapshot` to compare. + * @param right - The `AggregateQuerySnapshot` to compare. + * + * @returns true if the AggregateQuerySnapshots are equal. + */ +export function aggregateQuerySnapshotEqual( + left: AggregateQuerySnapshot, + right: AggregateQuerySnapshot +): boolean { + return ( + queryEqual(left.query, right.query) && deepEqual(left.data(), right.data()) + ); +} diff --git a/packages/firestore/src/platform/node/grpc_connection.ts b/packages/firestore/src/platform/node/grpc_connection.ts index f8b15440ffa..b5d77a11bf7 100644 --- a/packages/firestore/src/platform/node/grpc_connection.ts +++ b/packages/firestore/src/platform/node/grpc_connection.ts @@ -85,6 +85,12 @@ export class GrpcConnection implements Connection { // We cache stubs for the most-recently-used token. private cachedStub: GeneratedGrpcStub | null = null; + get shouldResourcePathBeIncludedInRequest(): boolean { + // Both `invokeRPC()` and `invokeStreamingRPC()` ignore their `path` arguments, and expect + // the "path" to be part of the given `request`. + return true; + } + constructor(protos: grpc.GrpcObject, private databaseInfo: DatabaseInfo) { // eslint-disable-next-line @typescript-eslint/no-explicit-any this.firestore = (protos as any)['google']['firestore']['v1']; diff --git a/packages/firestore/src/protos/firestore_proto_api.ts b/packages/firestore/src/protos/firestore_proto_api.ts index 8634f2a7d4a..e7dfe0a88b2 100644 --- a/packages/firestore/src/protos/firestore_proto_api.ts +++ b/packages/firestore/src/protos/firestore_proto_api.ts @@ -334,6 +334,32 @@ export declare namespace firestoreV1ApiClientInterfaces { readTime?: string; skippedResults?: number; } + interface RunAggregationQueryRequest { + parent?: string; + structuredAggregationQuery?: StructuredAggregationQuery; + transaction?: string; + newTransaction?: TransactionOptions; + readTime?: string; + } + interface RunAggregationQueryResponse { + result?: AggregationResult; + transaction?: string; + readTime?: string; + } + interface AggregationResult { + aggregateFields?: ApiClientObjectMap; + } + interface StructuredAggregationQuery { + structuredQuery?: StructuredQuery; + aggregations?: Aggregation[]; + } + interface Aggregation { + count?: Count; + alias?: string; + } + interface Count { + upTo?: number; + } interface Status { code?: number; message?: string; @@ -479,6 +505,10 @@ export declare type RunQueryRequest = firestoreV1ApiClientInterfaces.RunQueryRequest; export declare type RunQueryResponse = firestoreV1ApiClientInterfaces.RunQueryResponse; +export declare type RunAggregationQueryRequest = + firestoreV1ApiClientInterfaces.RunAggregationQueryRequest; +export declare type RunAggregationQueryResponse = + firestoreV1ApiClientInterfaces.RunAggregationQueryResponse; export declare type Status = firestoreV1ApiClientInterfaces.Status; export declare type StructuredQuery = firestoreV1ApiClientInterfaces.StructuredQuery; diff --git a/packages/firestore/src/protos/google/firestore/v1/aggregation_result.proto b/packages/firestore/src/protos/google/firestore/v1/aggregation_result.proto new file mode 100644 index 00000000000..538e3fef5e4 --- /dev/null +++ b/packages/firestore/src/protos/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/packages/firestore/src/protos/google/firestore/v1/firestore.proto b/packages/firestore/src/protos/google/firestore/v1/firestore.proto index b149a7634eb..aefbe71699f 100644 --- a/packages/firestore/src/protos/google/firestore/v1/firestore.proto +++ b/packages/firestore/src/protos/google/firestore/v1/firestore.proto @@ -19,6 +19,7 @@ package google.firestore.v1; import "google/api/annotations.proto"; import "google/api/client.proto"; import "google/api/field_behavior.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"; @@ -133,6 +134,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: "*" + } + }; + } + // Partitions a query by returning partition cursors that can be used to run // the query in parallel. The returned partition cursors are split points that // can be used by RunQuery as starting/end points for the query results. @@ -522,6 +546,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 [(google.api.field_behavior) = REQUIRED]; + + // 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.PartitionQuery][google.firestore.v1.Firestore.PartitionQuery]. message PartitionQueryRequest { // Required. The parent resource name. In the format: diff --git a/packages/firestore/src/protos/google/firestore/v1/query.proto b/packages/firestore/src/protos/google/firestore/v1/query.proto index 304499847c4..1bb2b6cdd01 100644 --- a/packages/firestore/src/protos/google/firestore/v1/query.proto +++ b/packages/firestore/src/protos/google/firestore/v1/query.proto @@ -287,6 +287,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. + google.protobuf.Int64Value 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 diff --git a/packages/firestore/src/protos/protos.json b/packages/firestore/src/protos/protos.json index fa1bbd289f4..093e22c6451 100644 --- a/packages/firestore/src/protos/protos.json +++ b/packages/firestore/src/protos/protos.json @@ -919,6 +919,15 @@ "ruby_package": "Google::Cloud::Firestore::V1" }, "nested": { + "AggregationResult": { + "fields": { + "aggregateFields": { + "keyType": "string", + "type": "Value", + "id": 2 + } + } + }, "DocumentMask": { "fields": { "fieldPaths": { @@ -1269,6 +1278,29 @@ } ] }, + "RunAggregationQuery": { + "requestType": "RunAggregationQueryRequest", + "responseType": "RunAggregationQueryResponse", + "responseStream": true, + "options": { + "(google.api.http).post": "/v1/{parent=projects/*/databases/*/documents}:runAggregationQuery", + "(google.api.http).body": "*", + "(google.api.http).additional_bindings.post": "/v1/{parent=projects/*/databases/*/documents/*/**}:runAggregationQuery", + "(google.api.http).additional_bindings.body": "*" + }, + "parsedOptions": [ + { + "(google.api.http)": { + "post": "/v1/{parent=projects/*/databases/*/documents}:runAggregationQuery", + "body": "*", + "additional_bindings": { + "post": "/v1/{parent=projects/*/databases/*/documents/*/**}:runAggregationQuery", + "body": "*" + } + } + } + ] + }, "PartitionQuery": { "requestType": "PartitionQueryRequest", "responseType": "PartitionQueryResponse", @@ -1760,6 +1792,63 @@ } } }, + "RunAggregationQueryRequest": { + "oneofs": { + "queryType": { + "oneof": [ + "structuredAggregationQuery" + ] + }, + "consistencySelector": { + "oneof": [ + "transaction", + "newTransaction", + "readTime" + ] + } + }, + "fields": { + "parent": { + "type": "string", + "id": 1, + "options": { + "(google.api.field_behavior)": "REQUIRED" + } + }, + "structuredAggregationQuery": { + "type": "StructuredAggregationQuery", + "id": 2 + }, + "transaction": { + "type": "bytes", + "id": 4 + }, + "newTransaction": { + "type": "TransactionOptions", + "id": 5 + }, + "readTime": { + "type": "google.protobuf.Timestamp", + "id": 6 + } + } + }, + "RunAggregationQueryResponse": { + "fields": { + "result": { + "type": "AggregationResult", + "id": 1 + }, + "transaction": { + "type": "bytes", + "id": 2 + }, + "readTime": { + "type": "google.protobuf.Timestamp", + "id": 3 + } + } + }, "PartitionQueryRequest": { "oneofs": { "queryType": { @@ -2296,6 +2385,57 @@ } } }, + "StructuredAggregationQuery": { + "oneofs": { + "queryType": { + "oneof": [ + "structuredQuery" + ] + } + }, + "fields": { + "structuredQuery": { + "type": "StructuredQuery", + "id": 1 + }, + "aggregations": { + "rule": "repeated", + "type": "Aggregation", + "id": 3 + } + }, + "nested": { + "Aggregation": { + "oneofs": { + "operator": { + "oneof": [ + "count" + ] + } + }, + "fields": { + "count": { + "type": "Count", + "id": 1 + }, + "alias": { + "type": "string", + "id": 7 + } + }, + "nested": { + "Count": { + "fields": { + "upTo": { + "type": "google.protobuf.Int64Value", + "id": 1 + } + } + } + } + } + } + }, "Cursor": { "fields": { "values": { diff --git a/packages/firestore/src/remote/connection.ts b/packages/firestore/src/remote/connection.ts index 19a54d37c73..717f2e9dbd2 100644 --- a/packages/firestore/src/remote/connection.ts +++ b/packages/firestore/src/remote/connection.ts @@ -84,6 +84,14 @@ export interface Connection { appCheckToken: Token | null ): Stream; + /** + * Returns whether or not the implementation requires that the "path" of the resource + * (a document or a collection) be present in the request message. If true, then the + * request message must include the path. If false, then the request message must NOT + * include the path. + */ + readonly shouldResourcePathBeIncludedInRequest: boolean; + // TODO(mcg): subscribe to connection state changes. } diff --git a/packages/firestore/src/remote/datastore.ts b/packages/firestore/src/remote/datastore.ts index 0d96661f6b1..64661485218 100644 --- a/packages/firestore/src/remote/datastore.ts +++ b/packages/firestore/src/remote/datastore.ts @@ -24,8 +24,11 @@ import { Mutation } from '../model/mutation'; import { BatchGetDocumentsRequest as ProtoBatchGetDocumentsRequest, BatchGetDocumentsResponse as ProtoBatchGetDocumentsResponse, + RunAggregationQueryRequest as ProtoRunAggregationQueryRequest, + RunAggregationQueryResponse as ProtoRunAggregationQueryResponse, RunQueryRequest as ProtoRunQueryRequest, - RunQueryResponse as ProtoRunQueryResponse + RunQueryResponse as ProtoRunQueryResponse, + Value as ProtoValue } from '../protos/firestore_proto_api'; import { debugAssert, debugCast, hardAssert } from '../util/assert'; import { AsyncQueue } from '../util/async_queue'; @@ -45,7 +48,8 @@ import { JsonProtoSerializer, toMutation, toName, - toQueryTarget + toQueryTarget, + toRunAggregationQueryRequest } from './serializer'; /** @@ -232,6 +236,32 @@ export async function invokeRunQueryRpc( ); } +export async function invokeRunAggregationQueryRpc( + datastore: Datastore, + query: Query +): Promise { + const datastoreImpl = debugCast(datastore, DatastoreImpl); + const request = toRunAggregationQueryRequest( + datastoreImpl.serializer, + queryToTarget(query) + ); + + const parent = request.parent; + if (!datastoreImpl.connection.shouldResourcePathBeIncludedInRequest) { + delete request.parent; + } + const response = await datastoreImpl.invokeStreamingRPC< + ProtoRunAggregationQueryRequest, + ProtoRunAggregationQueryResponse + >('RunAggregationQuery', parent!, request, /*expectedResponseCount=*/ 1); + return ( + response + // Omit RunAggregationQueryResponse that only contain readTimes. + .filter(proto => !!proto.result) + .map(proto => proto.result!.aggregateFields!) + ); +} + export function newPersistentWriteStream( datastore: Datastore, queue: AsyncQueue, diff --git a/packages/firestore/src/remote/rest_connection.ts b/packages/firestore/src/remote/rest_connection.ts index 9c318486620..40ab436787a 100644 --- a/packages/firestore/src/remote/rest_connection.ts +++ b/packages/firestore/src/remote/rest_connection.ts @@ -37,6 +37,7 @@ const RPC_NAME_URL_MAPPING: StringMap = {}; RPC_NAME_URL_MAPPING['BatchGetDocuments'] = 'batchGet'; RPC_NAME_URL_MAPPING['Commit'] = 'commit'; RPC_NAME_URL_MAPPING['RunQuery'] = 'runQuery'; +RPC_NAME_URL_MAPPING['RunAggregationQuery'] = 'runAggregationQuery'; const RPC_URL_VERSION = 'v1'; @@ -54,6 +55,12 @@ export abstract class RestConnection implements Connection { protected readonly baseUrl: string; private readonly databaseRoot: string; + get shouldResourcePathBeIncludedInRequest(): boolean { + // Both `invokeRPC()` and `invokeStreamingRPC()` use their `path` arguments to determine + // where to run the query, and expect the `request` to NOT specify the "path". + return false; + } + constructor(private readonly databaseInfo: DatabaseInfo) { this.databaseId = databaseInfo.databaseId; const proto = databaseInfo.ssl ? 'https' : 'http'; diff --git a/packages/firestore/src/remote/serializer.ts b/packages/firestore/src/remote/serializer.ts index e3c5de5c5ed..21091d55a36 100644 --- a/packages/firestore/src/remote/serializer.ts +++ b/packages/firestore/src/remote/serializer.ts @@ -77,6 +77,7 @@ import { OrderDirection as ProtoOrderDirection, Precondition as ProtoPrecondition, QueryTarget as ProtoQueryTarget, + RunAggregationQueryRequest as ProtoRunAggregationQueryRequest, Status as ProtoStatus, Target as ProtoTarget, TargetChangeTargetChangeType as ProtoTargetChangeTargetChangeType, @@ -852,6 +853,26 @@ export function toQueryTarget( return result; } +export function toRunAggregationQueryRequest( + serializer: JsonProtoSerializer, + target: Target +): ProtoRunAggregationQueryRequest { + const queryTarget = toQueryTarget(serializer, target); + + return { + structuredAggregationQuery: { + aggregations: [ + { + count: {}, + alias: 'count_alias' + } + ], + structuredQuery: queryTarget.structuredQuery + }, + parent: queryTarget.parent + }; +} + export function convertQueryTargetToQuery(target: ProtoQueryTarget): Query { let path = fromQueryPath(target.parent!); diff --git a/packages/firestore/test/integration/api_internal/aggregation.test.ts b/packages/firestore/test/integration/api_internal/aggregation.test.ts new file mode 100644 index 00000000000..88d6caf17fd --- /dev/null +++ b/packages/firestore/test/integration/api_internal/aggregation.test.ts @@ -0,0 +1,126 @@ +/** + * @license + * 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. + */ + +/** This file is emporarily staying in api_internal folder till aggregate queries are public. */ +import { expect } from 'chai'; + +import { getCountFromServer } from '../../../src/api/aggregate'; +import { + collection, + collectionGroup, + doc, + disableNetwork, + query, + terminate, + where, + writeBatch, + DocumentData, + QueryDocumentSnapshot +} from '../util/firebase_export'; +import { + apiDescribe, + withEmptyTestCollection, + withTestCollection, + withTestDb +} from '../util/helpers'; +import { USE_EMULATOR } from '../util/settings'; + +(USE_EMULATOR ? apiDescribe : apiDescribe.skip)( + 'Count quries', + (persistence: boolean) => { + it('can run count query getCountFromServer', () => { + const testDocs = { + a: { author: 'authorA', title: 'titleA' }, + b: { author: 'authorB', title: 'titleB' } + }; + return withTestCollection(persistence, testDocs, async coll => { + const snapshot = await getCountFromServer(coll); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it("count query doesn't use converter", () => { + const testDocs = { + a: { author: 'authorA', title: 'titleA' }, + b: { author: 'authorB', title: 'titleB' } + }; + const throwingConverter = { + toFirestore(obj: never): DocumentData { + throw new Error('should never be called'); + }, + fromFirestore(snapshot: QueryDocumentSnapshot): never { + throw new Error('should never be called'); + } + }; + return withTestCollection(persistence, testDocs, async coll => { + const query_ = query( + coll, + where('author', '==', 'authorA') + ).withConverter(throwingConverter); + const snapshot = await getCountFromServer(query_); + expect(snapshot.data().count).to.equal(1); + }); + }); + + it('count query supports collection groups', () => { + return withTestDb(persistence, async db => { + const collectionGroupId = doc(collection(db, 'aggregateQueryTest')).id; + const docPaths = [ + `${collectionGroupId}/cg-doc1`, + `abc/123/${collectionGroupId}/cg-doc2`, + `zzz${collectionGroupId}/cg-doc3`, + `abc/123/zzz${collectionGroupId}/cg-doc4`, + `abc/123/zzz/${collectionGroupId}` + ]; + const batch = writeBatch(db); + for (const docPath of docPaths) { + batch.set(doc(db, docPath), { x: 1 }); + } + await batch.commit(); + const snapshot = await getCountFromServer( + collectionGroup(db, collectionGroupId) + ); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('getCountFromServer fails if firestore is terminated', () => { + return withEmptyTestCollection(persistence, async (coll, firestore) => { + await terminate(firestore); + expect(() => getCountFromServer(coll)).to.throw( + 'The client has already been terminated.' + ); + }); + }); + + it("terminate doesn't crash when there is count query in flight", () => { + return withEmptyTestCollection(persistence, async (coll, firestore) => { + void getCountFromServer(coll); + await terminate(firestore); + }); + }); + + it('getCountFromServer fails if user is offline', () => { + return withEmptyTestCollection(persistence, async (coll, firestore) => { + await disableNetwork(firestore); + await expect(getCountFromServer(coll)).to.be.eventually.rejectedWith( + 'Failed to get count result because the client is offline' + ); + }); + }); + } +); diff --git a/packages/firestore/test/integration/util/helpers.ts b/packages/firestore/test/integration/util/helpers.ts index f2584916daf..79dbacaafa0 100644 --- a/packages/firestore/test/integration/util/helpers.ts +++ b/packages/firestore/test/integration/util/helpers.ts @@ -290,6 +290,13 @@ export function withTestCollection( return withTestCollectionSettings(persistence, DEFAULT_SETTINGS, docs, fn); } +export function withEmptyTestCollection( + persistence: boolean, + fn: (collection: CollectionReference, db: Firestore) => Promise +): Promise { + return withTestCollection(persistence, {}, fn); +} + // TODO(mikelehen): Once we wipe the database between tests, we can probably // return the same collection every time. export function withTestCollectionSettings( diff --git a/packages/firestore/test/lite/integration.test.ts b/packages/firestore/test/lite/integration.test.ts index 402c6066e09..823d0b524b6 100644 --- a/packages/firestore/test/lite/integration.test.ts +++ b/packages/firestore/test/lite/integration.test.ts @@ -20,6 +20,10 @@ import { initializeApp } from '@firebase/app'; import { expect, use } from 'chai'; import chaiAsPromised from 'chai-as-promised'; +import { + aggregateQuerySnapshotEqual, + getCount +} from '../../src/lite-api/aggregate'; import { Bytes } from '../../src/lite-api/bytes'; import { Firestore, @@ -79,7 +83,8 @@ import { runTransaction } from '../../src/lite-api/transaction'; import { writeBatch } from '../../src/lite-api/write_batch'; import { DEFAULT_PROJECT_ID, - DEFAULT_SETTINGS + DEFAULT_SETTINGS, + USE_EMULATOR } from '../integration/util/settings'; import { @@ -2113,3 +2118,269 @@ describe('withConverter() support', () => { }); }); }); + +// eslint-disable-next-line no-restricted-properties +(USE_EMULATOR ? describe : describe.skip)('Count quries', () => { + it('AggregateQuerySnapshot inherits the original query', () => { + return withTestCollection(async coll => { + const query_ = query(coll); + const snapshot = await getCount(query_); + expect(snapshot.query).to.equal(query_); + }); + }); + + it('run count query on empty collection', () => { + return withTestCollection(async coll => { + const snapshot = await getCount(coll); + expect(snapshot.data().count).to.equal(0); + }); + }); + + it('run count query on collection with 3 docs', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const snapshot = await getCount(coll); + expect(snapshot.data().count).to.equal(3); + }); + }); + + it('run count query fails on invalid collection reference', () => { + return withTestDb(async db => { + const queryForRejection = collection(db, '__badpath__'); + await expect(getCount(queryForRejection)).to.eventually.be.rejectedWith( + 'Request failed with error: Bad Request' + ); + }); + }); + + it('count query supports filter', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, where('author', '==', 'authorA')); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('count query supports filter and a small limit size', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, where('author', '==', 'authorA'), limit(1)); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(1); + }); + }); + + it('count query supports filter and a large limit size', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, where('author', '==', 'authorA'), limit(3)); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('count query supports order by', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: null }, + { author: 'authorB' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, orderBy('title')); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(3); + }); + }); + + it('count query supports order by and startAt', () => { + const testDocs = [ + { id: 3, author: 'authorA', title: 'titleA' }, + { id: 1, author: 'authorA', title: 'titleB' }, + { id: 2, author: 'authorB', title: 'titleC' }, + { id: null, author: 'authorB', title: 'titleD' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, orderBy('id'), startAt(2)); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('count query supports order by and startAfter', () => { + const testDocs = [ + { id: 3, author: 'authorA', title: 'titleA' }, + { id: 1, author: 'authorA', title: 'titleB' }, + { id: 2, author: 'authorB', title: 'titleC' }, + { id: null, author: 'authorB', title: 'titleD' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, orderBy('id'), startAfter(2)); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(1); + }); + }); + + it('count query supports order by and endAt', () => { + const testDocs = [ + { id: 3, author: 'authorA', title: 'titleA' }, + { id: 1, author: 'authorA', title: 'titleB' }, + { id: 2, author: 'authorB', title: 'titleC' }, + { id: null, author: 'authorB', title: 'titleD' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, orderBy('id'), startAt(1), endAt(2)); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('count query supports order by and endBefore', () => { + const testDocs = [ + { id: 3, author: 'authorA', title: 'titleA' }, + { id: 1, author: 'authorA', title: 'titleB' }, + { id: 2, author: 'authorB', title: 'titleC' }, + { id: null, author: 'authorB', title: 'titleD' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query(coll, orderBy('id'), startAt(1), endBefore(2)); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(1); + }); + }); + + it("count query doesn't use converter", () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + const throwingConverter = { + toFirestore(obj: never): DocumentData { + throw new Error('should never be called'); + }, + fromFirestore(snapshot: QueryDocumentSnapshot): never { + throw new Error('should never be called'); + } + }; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query_ = query( + coll, + where('author', '==', 'authorA') + ).withConverter(throwingConverter); + const snapshot = await getCount(query_); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('count query supports collection groups', () => { + return withTestDb(async db => { + const collectionGroupId = doc(collection(db, 'countTest')).id; + const docPaths = [ + `${collectionGroupId}/cg-doc1`, + `abc/123/${collectionGroupId}/cg-doc2`, + `zzz${collectionGroupId}/cg-doc3`, + `abc/123/zzz${collectionGroupId}/cg-doc4`, + `abc/123/zzz/${collectionGroupId}` + ]; + const batch = writeBatch(db); + for (const docPath of docPaths) { + batch.set(doc(db, docPath), { x: 1 }); + } + await batch.commit(); + const snapshot = await getCount(collectionGroup(db, collectionGroupId)); + expect(snapshot.data().count).to.equal(2); + }); + }); + + it('aggregateQuerySnapshotEqual on same queries be truthy', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query1 = query(coll, where('author', '==', 'authorA')); + const query2 = query(coll, where('author', '==', 'authorA')); + const snapshot1A = await getCount(query1); + const snapshot1B = await getCount(query1); + const snapshot2 = await getCount(query2); + expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot1B)).to.be.true; + expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot2)).to.be.true; + }); + }); + + it('aggregateQuerySnapshotEqual on same queries with different documents size be falsy', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query1 = query(coll, where('author', '==', 'authorA')); + const snapshot1A = await getCount(query1); + await addDoc(coll, { author: 'authorA', title: 'titleD' }); + const query2 = query(coll, where('author', '==', 'authorA')); + const snapshot1B = await getCount(query1); + const snapshot2 = await getCount(query2); + expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot1B)).to.be.false; + expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot2)).to.be.false; + }); + }); + + it('aggregateQuerySnapshotEqual on different queries be falsy', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' }, + { author: 'authorB', title: 'titleD' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const query1 = query(coll, where('author', '==', 'authorA')); + const query2 = query(coll, where('author', '==', 'authorB')); + const snapshot1 = await getCount(query1); + const snapshot2 = await getCount(query2); + expect(aggregateQuerySnapshotEqual(snapshot1, snapshot2)).to.be.false; + }); + }); + + it('count query fails on a terminated Firestore', () => { + return withTestCollection(async coll => { + await terminate(coll.firestore); + expect(() => getCount(coll)).to.throw( + 'The client has already been terminated.' + ); + }); + }); + + it('terminate Firestore not effect count query in flight', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' } + ]; + return withTestCollectionAndInitialData(testDocs, async coll => { + const promise = getCount(coll); + await terminate(coll.firestore); + const snapshot = await promise; + expect(snapshot.data().count).to.equal(3); + }); + }); +}); diff --git a/packages/firestore/test/unit/remote/datastore.test.ts b/packages/firestore/test/unit/remote/datastore.test.ts index f57fcf6f38d..795bf59edaf 100644 --- a/packages/firestore/test/unit/remote/datastore.test.ts +++ b/packages/firestore/test/unit/remote/datastore.test.ts @@ -65,6 +65,8 @@ describe('Datastore', () => { ): Stream { throw new Error('MockConnection.openStream() must be replaced'); } + + shouldResourcePathBeIncludedInRequest: boolean = false; } class MockAuthCredentialsProvider extends EmptyAuthCredentialsProvider { diff --git a/packages/firestore/test/unit/specs/spec_test_components.ts b/packages/firestore/test/unit/specs/spec_test_components.ts index ab63a4576b5..1436efa3fcd 100644 --- a/packages/firestore/test/unit/specs/spec_test_components.ts +++ b/packages/firestore/test/unit/specs/spec_test_components.ts @@ -253,6 +253,8 @@ export class MockConnection implements Connection { constructor(private queue: AsyncQueue) {} + shouldResourcePathBeIncludedInRequest: boolean = false; + /** * Tracks the currently active watch targets as detected by the mock watch * stream, as a mapping from target ID to query Target.