Make count queries publicly available for use (#6608)

Author: milaGGL

.changeset/wicked-tomatoes-grow.md

@@ -0,0 +1,6 @@
+---
+'@firebase/firestore': minor
+'firebase': minor
+---
+
+Added `getCountFromServer()` (`getCount()` in the Lite SDK), which fetches the number of documents in the result set without actually downloading the documents.

common/api-review/firestore-lite.api.md

@@ -17,6 +17,35 @@ export type AddPrefixToKeys<Prefix extends string, T extends Record<string, unkn
     [K in keyof T & string as `${Prefix}.${K}`]+?: T[K];
 };
 
+// @public
+export class AggregateField<T> {
+    type: string;
+}
+
+// @public
+export type AggregateFieldType = AggregateField<number>;
+
+// @public
+export class AggregateQuerySnapshot<T extends AggregateSpec> {
+    data(): AggregateSpecData<T>;
+    readonly query: Query<unknown>;
+    readonly type = "AggregateQuerySnapshot";
+}
+
+// @public
+export function aggregateQuerySnapshotEqual<T extends AggregateSpec>(left: AggregateQuerySnapshot<T>, right: AggregateQuerySnapshot<T>): boolean;
+
+// @public
+export interface AggregateSpec {
+    // (undocumented)
+    [field: string]: AggregateFieldType;
+}
+
+// @public
+export type AggregateSpecData<T extends AggregateSpec> = {
+    [P in keyof T]: T[P] extends AggregateField<infer U> ? U : never;
+};
+
 // @public
 export function arrayRemove(...elements: unknown[]): FieldValue;
 
@@ -169,6 +198,11 @@ export class GeoPoint {
     };
 }
 
+// @public
+export function getCount(query: Query<unknown>): Promise<AggregateQuerySnapshot<{
+    count: AggregateField<number>;
+}>>;
+
 // @public
 export function getDoc<T>(reference: DocumentReference<T>): Promise<DocumentSnapshot<T>>;
 

common/api-review/firestore.api.md

@@ -17,6 +17,35 @@ export type AddPrefixToKeys<Prefix extends string, T extends Record<string, unkn
     [K in keyof T & string as `${Prefix}.${K}`]+?: T[K];
 };
 
+// @public
+export class AggregateField<T> {
+    type: string;
+}
+
+// @public
+export type AggregateFieldType = AggregateField<number>;
+
+// @public
+export class AggregateQuerySnapshot<T extends AggregateSpec> {
+    data(): AggregateSpecData<T>;
+    readonly query: Query<unknown>;
+    readonly type = "AggregateQuerySnapshot";
+}
+
+// @public
+export function aggregateQuerySnapshotEqual<T extends AggregateSpec>(left: AggregateQuerySnapshot<T>, right: AggregateQuerySnapshot<T>): boolean;
+
+// @public
+export interface AggregateSpec {
+    // (undocumented)
+    [field: string]: AggregateFieldType;
+}
+
+// @public
+export type AggregateSpecData<T extends AggregateSpec> = {
+    [P in keyof T]: T[P] extends AggregateField<infer U> ? U : never;
+};
+
 // @public
 export function arrayRemove(...elements: unknown[]): FieldValue;
 
@@ -209,6 +238,11 @@ export class GeoPoint {
     };
 }
 
+// @public
+export function getCountFromServer(query: Query<unknown>): Promise<AggregateQuerySnapshot<{
+    count: AggregateField<number>;
+}>>;
+
 // @public
 export function getDoc<T>(reference: DocumentReference<T>): Promise<DocumentSnapshot<T>>;
 

packages/firestore/lite/index.ts

@@ -27,6 +27,19 @@
 import { registerFirestore } from './register';
 registerFirestore();
 
+export {
+  aggregateQuerySnapshotEqual,
+  getCount
+} from '../src/lite-api/aggregate';
+
+export {
+  AggregateField,
+  AggregateFieldType,
+  AggregateSpec,
+  AggregateSpecData,
+  AggregateQuerySnapshot
+} from '../src/lite-api/aggregate_types';
+
 export { FirestoreSettings as Settings } from '../src/lite-api/settings';
 
 export {

packages/firestore/src/api.ts

@@ -15,6 +15,19 @@
  * limitations under the License.
  */
 
+export {
+  aggregateQuerySnapshotEqual,
+  getCountFromServer
+} from './api/aggregate';
+
+export {
+  AggregateField,
+  AggregateFieldType,
+  AggregateSpec,
+  AggregateSpecData,
+  AggregateQuerySnapshot
+} from './lite-api/aggregate_types';
+
 export { FieldPath, documentId } from './api/field_path';
 
 export {

packages/firestore/src/api/aggregate.ts

@@ -17,23 +17,43 @@
 
 import { Query } from '../api';
 import { firestoreClientRunCountQuery } from '../core/firestore_client';
-import { AggregateField, AggregateQuerySnapshot } from '../lite-api/aggregate';
+import {
+  AggregateField,
+  AggregateQuerySnapshot
+} from '../lite-api/aggregate_types';
 import { cast } from '../util/input_validation';
 
 import { ensureFirestoreConfigured, Firestore } from './database';
+import { ExpUserDataWriter } from './reference_impl';
+
+export { aggregateQuerySnapshotEqual } from '../lite-api/aggregate';
 
 /**
- * Executes the query and returns the results as a `AggregateQuerySnapshot` from the
- * server. Returns an error if the network is not available.
+ * Calculates the number of documents in the result set of the given query,
+ * without actually downloading the documents.
+ *
+ * Using this function to count the documents is efficient because only the
+ * final count, not the documents' data, is downloaded. This function can even
+ * count the documents if the result set would be prohibitively large to
+ * download entirely (e.g. thousands of documents).
  *
- * @param query - The `Query` to execute.
+ * The result received from the server is presented, unaltered, without
+ * considering any local state. That is, documents in the local cache are not
+ * taken into consideration, neither are local modifications not yet
+ * synchronized with the server. Previously-downloaded results, if any, are not
+ * used: every request using this source necessarily involves a round trip to
+ * the server.
  *
- * @returns A `Promise` that will be resolved with the results of the query.
+ * @param query - The query whose result set size to calculate.
+ * @returns A Promise that will be resolved with the count; the count can be
+ * retrieved from `snapshot.data().count`, where `snapshot` is the
+ * `AggregateQuerySnapshot` to which the returned Promise resolves.
  */
 export function getCountFromServer(
   query: Query<unknown>
 ): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
   const firestore = cast(query.firestore, Firestore);
   const client = ensureFirestoreConfigured(firestore);
-  return firestoreClientRunCountQuery(client, query);
+  const userDataWriter = new ExpUserDataWriter(firestore);
+  return firestoreClientRunCountQuery(client, query, userDataWriter);
 }

packages/firestore/src/core/count_query_runner.ts

@@ -0,0 +1,70 @@
+/**
+ * @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 { AbstractUserDataWriter, Query } from '../api';
+import {
+  AggregateField,
+  AggregateQuerySnapshot
+} from '../lite-api/aggregate_types';
+import { Value } from '../protos/firestore_proto_api';
+import { Datastore, invokeRunAggregationQueryRpc } from '../remote/datastore';
+import { hardAssert } from '../util/assert';
+
+/**
+ * CountQueryRunner encapsulates the logic needed to run the count aggregation
+ * queries.
+ */
+export class CountQueryRunner {
+  constructor(
+    private readonly query: Query<unknown>,
+    private readonly datastore: Datastore,
+    private readonly userDataWriter: AbstractUserDataWriter
+  ) {}
+
+  run(): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+    return invokeRunAggregationQueryRpc(this.datastore, this.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]) =>
+            this.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<number> }>(
+            this.query,
+            {
+              count: countValue
+            }
+          )
+        );
+      }
+    );
+  }
+}

packages/firestore/src/core/firestore_client.ts

@@ -17,17 +17,17 @@
 
 import { GetOptions } from '@firebase/firestore-types';
 
+import {
+  AbstractUserDataWriter,
+  AggregateField,
+  AggregateQuerySnapshot
+} from '../api';
 import { LoadBundleTask } from '../api/bundle';
 import {
   CredentialChangeListener,
   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 {
@@ -68,6 +68,7 @@ import {
   OfflineComponentProvider,
   OnlineComponentProvider
 } from './component_provider';
+import { CountQueryRunner } from './count_query_runner';
 import { DatabaseId, DatabaseInfo } from './database_info';
 import {
   addSnapshotsInSyncListener,
@@ -510,7 +511,8 @@ export function firestoreClientTransaction<T>(
 
 export function firestoreClientRunCountQuery(
   client: FirestoreClient,
-  query: LiteQuery<unknown>
+  query: LiteQuery<unknown>,
+  userDataWriter: AbstractUserDataWriter
 ): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
   const deferred = new Deferred<
     AggregateQuerySnapshot<{ count: AggregateField<number> }>
@@ -526,7 +528,12 @@ export function firestoreClientRunCountQuery(
           )
         );
       } else {
-        const result = await getCount(query);
+        const datastore = await getDatastore(client);
+        const result = new CountQueryRunner(
+          query,
+          datastore,
+          userDataWriter
+        ).run();
         deferred.resolve(result);
       }
     } catch (e) {