Mila/count update api surface (#6589)

Author: milaGGL

packages/firebase/compat/index.d.ts

@@ -8227,11 +8227,15 @@ declare namespace firebase.storage {
 }
 
 declare namespace firebase.firestore {
+
+  /** Alias dynamic document field value types to any */
+  export type DocumentFieldValue = any;
+  
   /**
    * Document data (for use with `DocumentReference.set()`) consists of fields
    * mapped to values.
    */
-  export type DocumentData = { [field: string]: any };
+  export type DocumentData = { [field: string]: DocumentFieldValue };
 
   /**
    * Update data (for use with `DocumentReference.update()`) consists of field

packages/firestore-types/index.d.ts

@@ -17,7 +17,9 @@
 
 import { EmulatorMockTokenOptions } from '@firebase/util';
 
-export type DocumentData = { [field: string]: any };
+export type DocumentFieldValue = any;
+
+export type DocumentData = { [field: string]: DocumentFieldValue };
 
 export type UpdateData = { [fieldPath: string]: any };
 

packages/firestore/lite/index.ts

@@ -29,6 +29,16 @@ registerFirestore();
 
 export { FirestoreSettings as Settings } from '../src/lite-api/settings';
 
+
+export {
+  AggregateField,
+  AggregateSpec,
+  AggregateSpecData,
+  AggregateQuerySnapshot,
+  getCount,
+  aggregateQuerySnapshotEqual
+} from '../src/lite-api/aggregate';
+
 export {
   Firestore as Firestore,
   EmulatorMockTokenOptions,

packages/firestore/src/api.ts

@@ -16,12 +16,12 @@
  */
 
 export {
-  AggregateQuery,
+  AggregateField,
+  AggregateSpec,
+  AggregateSpecData,
   AggregateQuerySnapshot,
-  aggregateQueryEqual,
-  aggregateQuerySnapshotEqual,
-  countQuery,
-  getAggregateFromServerDirect
+  getCountFromServer,
+  aggregateQuerySnapshotEqual
 } from './api/aggregate';
 
 export { FieldPath, documentId } from './api/field_path';

packages/firestore/src/api/aggregate.ts

@@ -14,25 +14,33 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
-import { firestoreClientRunAggregationQuery } from '../core/firestore_client';
-import { AggregateQuery, AggregateQuerySnapshot } from '../lite-api/aggregate';
+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';
 
 export {
-  AggregateQuery,
+  AggregateField,
+  AggregateSpec,
+  AggregateSpecData,
   AggregateQuerySnapshot,
-  aggregateQueryEqual,
-  aggregateQuerySnapshotEqual,
-  countQuery
+  aggregateQuerySnapshotEqual
 } from '../lite-api/aggregate';
 
-export function getAggregateFromServerDirect(
-  query: AggregateQuery
-): Promise<AggregateQuerySnapshot> {
-  const firestore = cast(query.query.firestore, Firestore);
+/**
+ * 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<unknown>
+): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+  const firestore = cast(query.firestore, Firestore);
   const client = ensureFirestoreConfigured(firestore);
-  return firestoreClientRunAggregationQuery(client, query);
+  return firestoreClientRunCountQuery(client, query);
 }

packages/firestore/src/core/firestore_client.ts

@@ -17,14 +17,15 @@
 
 import { GetOptions } from '@firebase/firestore-types';
 
-import { AggregateQuery, AggregateQuerySnapshot } from '../api';
+import { AggregateField, AggregateQuerySnapshot } from '../api/aggregate';
 import { LoadBundleTask } from '../api/bundle';
 import {
   CredentialChangeListener,
   CredentialsProvider
 } from '../api/credentials';
 import { User } from '../auth/user';
-import { getAggregate } from '../lite-api/aggregate';
+import { getCount } from '../lite-api/aggregate';
+import { Query as LiteQuery } from '../lite-api/reference';
 import { LocalStore } from '../local/local_store';
 import {
   localStoreExecuteQuery,
@@ -504,23 +505,25 @@ export function firestoreClientTransaction<T>(
   return deferred.promise;
 }
 
-export function firestoreClientRunAggregationQuery(
+export function firestoreClientRunCountQuery(
   client: FirestoreClient,
-  query: AggregateQuery
-): Promise<AggregateQuerySnapshot> {
-  const deferred = new Deferred<AggregateQuerySnapshot>();
+  query: LiteQuery<unknown>
+): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+  const deferred = new Deferred<
+    AggregateQuerySnapshot<{ count: AggregateField<number> }>
+  >();
   client.asyncQueue.enqueueAndForget(async () => {
     const remoteStore = await getRemoteStore(client);
     if (!canUseNetwork(remoteStore)) {
       deferred.reject(
         new FirestoreError(
           Code.UNAVAILABLE,
-          'Failed to get aggregate result because the client is offline.'
+          'Failed to get count result because the client is offline.'
         )
       );
     } else {
       try {
-        const result = await getAggregate(query);
+        const result = await getCount(query);
         deferred.resolve(result);
       } catch (e) {
         deferred.reject(e as Error);

packages/firestore/src/lite-api/aggregate.ts

@@ -15,6 +15,8 @@
  * 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';
@@ -26,61 +28,87 @@ import { Query, queryEqual } from './reference';
 import { LiteUserDataWriter } from './reference_impl';
 
 /**
- * An `AggregateQuery` computes some aggregation statistics from the result set of
- * a base `Query`.
+ * An `AggregateField`that captures input type T.
  */
-export class AggregateQuery {
-  readonly type = 'AggregateQuery';
-  /**
-   * The query on which you called `countQuery` in order to get this `AggregateQuery`.
-   */
-  readonly query: Query<unknown>;
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export class AggregateField<T> {
+  type = 'AggregateField';
+}
 
-  /** @hideconstructor */
-  constructor(query: Query<unknown>) {
-    this.query = query;
-  }
+/**
+ * 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<number> {
+  return new AggregateField<number>();
 }
 
 /**
- * An `AggregateQuerySnapshot` contains results of a `AggregateQuery`.
+ * The union of all `AggregateField` types that are returned from the factory
+ * functions.
+ */
+type AggregateFieldType = ReturnType<typeof count>;
+
+/**
+ * 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<T extends AggregateSpec> = {
+  [P in keyof T]: T[P] extends AggregateField<infer U> ? U : never;
+};
+
+/**
+ * An `AggregateQuerySnapshot` contains the results of running an aggregate query.
  */
-export class AggregateQuerySnapshot {
+export class AggregateQuerySnapshot<T extends AggregateSpec> {
   readonly type = 'AggregateQuerySnapshot';
-  readonly query: AggregateQuery;
 
   /** @hideconstructor */
-  constructor(query: AggregateQuery, private readonly _count: number) {
-    this.query = query;
-  }
+  constructor(
+    readonly query: Query<unknown>,
+    private readonly _data: AggregateSpecData<T>
+  ) {}
 
   /**
-   * @returns The result of a document count aggregation. Returns null if no count aggregation is
-   * available in the result.
+   * 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.
    */
-  getCount(): number | null {
-    return this._count;
+  data(): AggregateSpecData<T> {
+    return this._data;
   }
 }
 
 /**
- * Creates an `AggregateQuery` counting the number of documents matching this query.
+ * 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 `AggregateQuery` object that can be used to count the number of documents in
- * the result set of this query.
+ * @returns An `AggregateQuerySnapshot` that contains the number of documents.
  */
-export function countQuery(query: Query<unknown>): AggregateQuery {
-  return new AggregateQuery(query);
-}
-
-export function getAggregate(
-  query: AggregateQuery
-): Promise<AggregateQuerySnapshot> {
-  const firestore = cast(query.query.firestore, Firestore);
+export function getCount(
+  query: Query<unknown>
+): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+  const firestore = cast(query.firestore, Firestore);
   const datastore = getDatastore(firestore);
   const userDataWriter = new LiteUserDataWriter(firestore);
-
-  return invokeRunAggregationQueryRpc(datastore, query).then(result => {
+  return invokeRunAggregationQueryRpc(datastore, query._query).then(result => {
     hardAssert(
       result[0] !== undefined,
       'Aggregation fields are missing from result.'
@@ -90,29 +118,36 @@ export function getAggregate(
       .filter(([key, value]) => key === 'count_alias')
       .map(([key, value]) => userDataWriter.convertValue(value as Value));
 
-    const count = counts[0];
+    const countValue = counts[0];
+
     hardAssert(
-      typeof count === 'number',
-      'Count aggeragte field value is not a number: ' + count
+      typeof countValue === 'number',
+      'Count aggregate field value is not a number: ' + countValue
     );
 
-    return Promise.resolve(new AggregateQuerySnapshot(query, count));
+    return Promise.resolve(
+      new AggregateQuerySnapshot<{ count: AggregateField<number> }>(query, {
+        count: countValue
+      })
+    );
   });
 }
 
-export function aggregateQueryEqual(
-  left: AggregateQuery,
-  right: AggregateQuery
-): boolean {
-  return queryEqual(left.query, right.query);
-}
-
-export function aggregateQuerySnapshotEqual(
-  left: AggregateQuerySnapshot,
-  right: AggregateQuerySnapshot
+/**
+ * Compares two `AggregateQuerySnapshot` instances for equality.
+ * Two `AggregateQuerySnapshot` instances are considered "equal" if they have
+ * the same underlying query, the same metadata, and the same data.
+ *
+ * @param left - The `AggregateQuerySnapshot` to compare.
+ * @param right - The `AggregateQuerySnapshot` to compare.
+ *
+ * @returns true if the AggregateQuerySnapshos are equal.
+ */
+export function aggregateQuerySnapshotEqual<T extends AggregateSpec>(
+  left: AggregateQuerySnapshot<T>,
+  right: AggregateQuerySnapshot<T>
 ): boolean {
   return (
-    aggregateQueryEqual(left.query, right.query) &&
-    left.getCount() === right.getCount()
+    queryEqual(left.query, right.query) && deepEqual(left.data(), right.data())
   );
 }

packages/firestore/src/lite-api/reference.ts

@@ -40,14 +40,17 @@ import { FieldValue } from './field_value';
 import { FirestoreDataConverter } from './snapshot';
 import { NestedUpdateFields, Primitive } from './types';
 
+/** Alias dynamic document field value types to any */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type DocumentFieldValue = any;
+
 /**
  * Document data (for use with {@link @firebase/firestore/lite#(setDoc:1)}) consists of fields mapped to
  * values.
  */
 export interface DocumentData {
   /** A mapping between a field and its value. */
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  [field: string]: any;
+  [field: string]: DocumentFieldValue;
 }
 
 /**

packages/firestore/src/remote/connection.ts

@@ -90,7 +90,7 @@ export interface Connection {
    * request message must include the path. If false, then the request message must NOT
    * include the path.
    */
-  get shouldResourcePathBeIncludedInRequest(): boolean;
+  readonly shouldResourcePathBeIncludedInRequest: boolean;
 
   // TODO(mcg): subscribe to connection state changes.
 }