diff --git a/common/api-review/firestore.api.md b/common/api-review/firestore.api.md index e77f2353747..35e21d31786 100644 --- a/common/api-review/firestore.api.md +++ b/common/api-review/firestore.api.md @@ -17,10 +17,9 @@ export type AddPrefixToKeys { - // (undocumented) - readonly query: Query; +// @public +export class AggregateQuery { + readonly query: Query; // (undocumented) readonly type = "AggregateQuery"; } @@ -28,7 +27,7 @@ export class AggregateQuery { // @public (undocumented) export function aggregateQueryEqual(left: AggregateQuery, right: AggregateQuery): boolean; -// @public (undocumented) +// @public export class AggregateQuerySnapshot { // (undocumented) getCount(): number | null; @@ -93,8 +92,8 @@ export function connectFirestoreEmulator(firestore: Firestore, host: string, por mockUserToken?: EmulatorMockTokenOptions | string; }): void; -// @public (undocumented) -export function countQuery(query: Query): AggregateQuery; +// @public +export function countQuery(query: Query): AggregateQuery; // @public export function deleteDoc(reference: DocumentReference): Promise; diff --git a/packages/firestore/src/lite-api/aggregate.ts b/packages/firestore/src/lite-api/aggregate.ts index 0afd82a5996..49770bb8293 100644 --- a/packages/firestore/src/lite-api/aggregate.ts +++ b/packages/firestore/src/lite-api/aggregate.ts @@ -15,40 +15,89 @@ * limitations under the License. */ -import { DocumentData, Query, queryEqual } from './reference'; +import { Value } from '../protos/firestore_proto_api'; +import { invokeRunAggregationQueryRpc } from '../remote/datastore'; +import { hardAssert } from '../util/assert'; +import { cast } from '../util/input_validation'; -export class AggregateQuery { +import { getDatastore } from './components'; +import { Firestore } from './database'; +import { Query, queryEqual } from './reference'; +import { LiteUserDataWriter } from './reference_impl'; + +/** + * An `AggregateQuery` computes some aggregation statistics from the result set of + * a base `Query`. + */ +export class AggregateQuery { readonly type = 'AggregateQuery'; - readonly query: Query; + /** + * The query on which you called `countQuery` in order to get this `AggregateQuery`. + */ + readonly query: Query; /** @hideconstructor */ - constructor(query: Query) { + constructor(query: Query) { this.query = query; } } +/** + * An `AggregateQuerySnapshot` contains results of a `AggregateQuery`. + */ export class AggregateQuerySnapshot { readonly type = 'AggregateQuerySnapshot'; readonly query: AggregateQuery; /** @hideconstructor */ - constructor(query: AggregateQuery, readonly _count: number) { + constructor(query: AggregateQuery, private readonly _count: number) { this.query = query; } + /** + * @returns The result of a document count aggregation. Returns null if no count aggregation is + * available in the result. + */ getCount(): number | null { return this._count; } } -export function countQuery(query: Query): AggregateQuery { +/** + * Creates an `AggregateQuery` counting the number of documents matching this query. + * + * @returns An `AggregateQuery` object that can be used to count the number of documents in + * the result set of this query. + */ +export function countQuery(query: Query): AggregateQuery { return new AggregateQuery(query); } export function getAggregateFromServerDirect( query: AggregateQuery ): Promise { - return Promise.resolve(new AggregateQuerySnapshot(query, 42)); + const firestore = cast(query.query.firestore, Firestore); + const datastore = getDatastore(firestore); + const userDataWriter = new LiteUserDataWriter(firestore); + + return invokeRunAggregationQueryRpc(datastore, 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 count = counts[0]; + hardAssert( + typeof count === 'number', + 'Count aggeragte field value is not a number: ' + count + ); + + return Promise.resolve(new AggregateQuerySnapshot(query, count)); + }); } export function aggregateQueryEqual( 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/remote/datastore.ts b/packages/firestore/src/remote/datastore.ts index 0d96661f6b1..febae1187dc 100644 --- a/packages/firestore/src/remote/datastore.ts +++ b/packages/firestore/src/remote/datastore.ts @@ -18,14 +18,18 @@ import { CredentialsProvider } from '../api/credentials'; import { User } from '../auth/user'; import { Query, queryToTarget } from '../core/query'; +import { AggregateQuery } from '../lite-api/aggregate'; import { Document } from '../model/document'; import { DocumentKey } from '../model/document_key'; 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 +49,8 @@ import { JsonProtoSerializer, toMutation, toName, - toQueryTarget + toQueryTarget, + toRunAggregationQueryRequest } from './serializer'; /** @@ -232,6 +237,29 @@ export async function invokeRunQueryRpc( ); } +export async function invokeRunAggregationQueryRpc( + datastore: Datastore, + aggregateQuery: AggregateQuery +): Promise { + const datastoreImpl = debugCast(datastore, DatastoreImpl); + const request = toRunAggregationQueryRequest( + datastoreImpl.serializer, + queryToTarget(aggregateQuery.query._query) + ); + const response = await datastoreImpl.invokeStreamingRPC< + ProtoRunAggregationQueryRequest, + ProtoRunAggregationQueryResponse + >('RunAggregationQuery', request.parent!, { + structuredAggregationQuery: request.structuredAggregationQuery + }); + 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..67fffbdde92 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'; @@ -78,7 +79,6 @@ export abstract class RestConnection implements Connection { const headers = {}; this.modifyHeadersForRequest(headers, authToken, appCheckToken); - return this.performRPCRequest(rpcName, url, headers, req).then( response => { logDebug(LOG_TAG, 'Received: ', response); 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/count.test.ts b/packages/firestore/test/integration/api/count.test.ts deleted file mode 100644 index 1a4bac61ced..00000000000 --- a/packages/firestore/test/integration/api/count.test.ts +++ /dev/null @@ -1,53 +0,0 @@ -/** - * @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 { expect } from 'chai'; -import { - countQuery, - getAggregateFromServerDirect, - query -} from '../util/firebase_export'; -import { apiDescribe, withTestCollection } from '../util/helpers'; - -apiDescribe('Aggregation COUNT query:', (persistence: boolean) => { - it('empty collection count equals to 0', () => { - const testDocs = {}; - return withTestCollection(persistence, testDocs, collection => { - const countQuery_ = countQuery(query(collection)); - return getAggregateFromServerDirect(countQuery_).then(snapshot => { - expect(snapshot.getCount()).to.equal(0); - }); - }); - }); - - it('test collection count equals to 6', () => { - const testDocs = { - a: { k: 'a' }, - b: { k: 'b' }, - c: { k: 'c' }, - d: { k: 'd' }, - e: { k: 'e' }, - f: { k: 'f' } - }; - return withTestCollection(persistence, testDocs, collection => { - const countQuery_ = countQuery(query(collection)); - return getAggregateFromServerDirect(countQuery_).then(snapshot => { - expect(snapshot.getCount()).to.equal(6); - }); - }); - }); -}); diff --git a/packages/firestore/test/lite/integration.test.ts b/packages/firestore/test/lite/integration.test.ts index deaa9a0236f..398175c195f 100644 --- a/packages/firestore/test/lite/integration.test.ts +++ b/packages/firestore/test/lite/integration.test.ts @@ -20,6 +20,12 @@ import { initializeApp } from '@firebase/app'; import { expect, use } from 'chai'; import chaiAsPromised from 'chai-as-promised'; +import { + countQuery, + getAggregateFromServerDirect, + aggregateQueryEqual, + aggregateQuerySnapshotEqual +} from '../../src/lite-api/aggregate'; import { Bytes } from '../../src/lite-api/bytes'; import { Firestore, @@ -2038,3 +2044,134 @@ describe('withConverter() support', () => { }); }); }); + +describe('countQuery()', () => { + const testDocs = [ + { author: 'authorA', title: 'titleA' }, + { author: 'authorA', title: 'titleB' }, + { author: 'authorB', title: 'titleC' }, + { author: 'authorB', title: 'titleD' }, + { author: 'authorB', title: 'titleE' } + ]; + + it('AggregateQuery and AggregateQuerySnapshot inherits the original query', () => { + return withTestCollection(async coll => { + const query_ = query(coll); + const countQuery_ = countQuery(query_); + expect(countQuery_.query).to.equal(query_); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.query).to.equal(countQuery_); + expect(snapshot.query.query).to.equal(query_); + }); + }); + + it('empty test collection count', () => { + return withTestCollection(async coll => { + const countQuery_ = countQuery(query(coll)); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.getCount()).to.equal(0); + }); + }); + + it('test collection count with 5 docs', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const countQuery_ = countQuery(query(collection)); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.getCount()).to.equal(5); + }); + }); + + it('test collection count with filter', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query_ = query(collection, where('author', '==', 'authorA')); + const countQuery_ = countQuery(query_); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.getCount()).to.equal(2); + }); + }); + + it('test collection count with filter and a small limit size', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query_ = query( + collection, + where('author', '==', 'authorA'), + limit(1) + ); + const countQuery_ = countQuery(query_); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.getCount()).to.equal(1); + }); + }); + + it('test collection count with filter and a large limit size', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query_ = query( + collection, + where('author', '==', 'authorA'), + limit(3) + ); + const countQuery_ = countQuery(query_); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.getCount()).to.equal(2); + }); + }); + + it('test collection count with converter on query', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query_ = query( + collection, + where('author', '==', 'authorA') + ).withConverter(postConverter); + const countQuery_ = countQuery(query_); + const snapshot = await getAggregateFromServerDirect(countQuery_); + expect(snapshot.getCount()).to.equal(2); + }); + }); + + it('aggregateQueryEqual on same queries', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query1 = query(collection, where('author', '==', 'authorA')); + const query2 = query(collection, where('author', '==', 'authorA')); + const countQuery1 = countQuery(query1); + const countQuery2 = countQuery(query2); + expect(aggregateQueryEqual(countQuery1, countQuery2)).to.be.true; + }); + }); + + it('aggregateQueryEqual on different queries', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query1 = query(collection, where('author', '==', 'authorA')); + const query2 = query(collection, where('author', '==', 'authorB')); + const countQuery1 = countQuery(query1); + const countQuery2 = countQuery(query2); + expect(aggregateQueryEqual(countQuery1, countQuery2)).to.be.false; + }); + }); + + it('aggregateQuerySnapshotEqual on same queries', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query1 = query(collection, where('author', '==', 'authorA')); + const query2 = query(collection, where('author', '==', 'authorA')); + const countQuery1A = countQuery(query1); + const countQuery1B = countQuery(query1); + const countQuery2 = countQuery(query2); + const snapshot1A = await getAggregateFromServerDirect(countQuery1A); + const snapshot1B = await getAggregateFromServerDirect(countQuery1B); + const snapshot2 = await getAggregateFromServerDirect(countQuery2); + expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot1B)).to.be.true; + expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot2)).to.be.true; + }); + }); + + it('aggregateQuerySnapshotEqual on different queries', () => { + return withTestCollectionAndInitialData(testDocs, async collection => { + const query1 = query(collection, where('author', '==', 'authorA')); + const query2 = query(collection, where('author', '==', 'authorB')); + const countQuery1 = countQuery(query1); + const countQuery2 = countQuery(query2); + const snapshot1 = await getAggregateFromServerDirect(countQuery1); + const snapshot2 = await getAggregateFromServerDirect(countQuery2); + expect(aggregateQuerySnapshotEqual(snapshot1, snapshot2)).to.be.false; + }); + }); +});