Mila/count update api surface

packages/firebase/compat/index.d.ts

@@ -8227,7 +8227,10 @@ 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.

packages/firestore/src/api/aggregate.ts

@@ -14,14 +14,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { Query, SnapshotOptions } from '../api';
+import { Query } from '../api';
 import { firestoreClientRunCountQuery } from '../core/firestore_client';
-import {
-  AggregateField,
-  AggregateQuerySnapshot as AggregateQuerySnapshotLite,
-  AggregateSpec,
-  AggregateSpecData
-} from '../lite-api/aggregate';
+import { AggregateField, AggregateQuerySnapshot } from '../lite-api/aggregate';
 import { cast } from '../util/input_validation';
 
 import { ensureFirestoreConfigured, Firestore } from './database';
@@ -31,29 +26,15 @@ export {
   AggregateSpec,
   AggregateSpecData,
   AggregateQuerySnapshot,
-  aggregateSnapshotEqual,
-  aggregateFieldEqual
+  aggregateSnapshotEqual
 } from '../lite-api/aggregate';
 
 /**
- * A `AggregateQuerySnapshot` contains the result of running an aggregation query.
- * The result data can be extracted with `.data()`.
- */
-class AggregateQuerySnapshot<
-  T extends AggregateSpec
-> extends AggregateQuerySnapshotLite<T> {
-  constructor(query: Query<unknown>, _data: AggregateSpecData<T>) {
-    super(query, _data);
-  }
-  data(options: SnapshotOptions = {}): AggregateSpecData<T> {
-    return super.data();
-  }
-}
-
-/**
- * Executes the query and returns the results as a `QuerySnapshot` from the
+ * 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(

packages/firestore/src/core/firestore_client.ts

@@ -17,18 +17,17 @@
 
 import { GetOptions } from '@firebase/firestore-types';
 
-import {
-  AggregateField,
-  AggregateSpec,
-  AggregateQuerySnapshot
-} from '../api/aggregate';
+import { AggregateField } from '../api/aggregate';
 import { LoadBundleTask } from '../api/bundle';
 import {
   CredentialChangeListener,
   CredentialsProvider
 } from '../api/credentials';
 import { User } from '../auth/user';
-import { getCountFromServer } from '../lite-api/aggregate';
+import {
+  getCount as getCountFromServer,
+  AggregateQuerySnapshot as LiteAggregateQuerySnapshot
+} from '../lite-api/aggregate';
 import { Query as LiteQuery } from '../lite-api/reference';
 import { LocalStore } from '../local/local_store';
 import {
@@ -509,12 +508,12 @@ export function firestoreClientTransaction<T>(
   return deferred.promise;
 }
 
-export function firestoreClientRunCountQuery<T extends AggregateSpec>(
+export function firestoreClientRunCountQuery(
   client: FirestoreClient,
   query: LiteQuery<unknown>
-): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+): Promise<LiteAggregateQuerySnapshot<{ count: AggregateField<number> }>> {
   const deferred = new Deferred<
-    AggregateQuerySnapshot<{ count: AggregateField<number> }>
+    LiteAggregateQuerySnapshot<{ count: AggregateField<number> }>
   >();
   client.asyncQueue.enqueueAndForget(async () => {
     const remoteStore = await getRemoteStore(client);

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

@@ -15,92 +15,77 @@
  * 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 { invokeRunAggregationQueryRpc } from '../remote/datastore';
-import { Value } from '../protos/firestore_proto_api';
 
-import { DocumentFieldValue, Query, queryEqual } from './reference';
-import { Firestore } from './database';
 import { getDatastore } from './components';
+import { Firestore } from './database';
+import { Query, queryEqual } from './reference';
 import { LiteUserDataWriter } from './reference_impl';
 
-import { deepEqual } from '@firebase/util';
-import { FieldPath } from './field_path';
-
 /**
- * An `AggregateField` computes some aggregation statistics from the result set of
- * an aggregation query.
+ * An `AggregateField`that captures input type T.
  */
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
 export class AggregateField<T> {
   type = 'AggregateField';
-  _datum?: T;
+}
 
-  constructor(readonly subType: string) {}
+/**
+ * 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>();
 }
 
 /**
  * The union of all `AggregateField` types that are returned from the factory
  * functions.
  */
-export type AggregateFieldType =
-  | AggregateField<number>
-  | AggregateField<DocumentFieldValue | undefined>
-  | AggregateField<number | undefined>;
+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 type AggregateSpec = { [field: string]: AggregateFieldType };
+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> = {
-  [Property in keyof T]-?: T[Property]['_datum'];
+  [P in keyof T]: T[P] extends AggregateField<infer U> ? U : never;
 };
 
-/**
- * Creates and returns an aggregation field that counts the documents in the result set.
- * @returns An `AggregateField` object that includes number of documents.
- */
-export function count(): AggregateField<number> {
-  return new AggregateField<number>('count');
-}
-
-/**
- * Compares two `AggregateField` instances for equality.
- * The two `AggregateField` instances are considered "equal" if and only if
- * they were created by the same factory function (e.g. `count()`, `min()`, and
- * `sum()`) with "equal" arguments.
- */
-export function aggregateFieldEqual(
-  left: AggregateField<unknown>,
-  right: AggregateField<unknown>
-): boolean {
-  return typeof left === typeof right && left.subType === right.subType;
-}
-
 /**
  * An `AggregateQuerySnapshot` contains the results of running an aggregate query.
  */
 export class AggregateQuerySnapshot<T extends AggregateSpec> {
   readonly type = 'AggregateQuerySnapshot';
 
+  /** @hideconstructor */
   constructor(
     readonly query: Query<unknown>,
-    protected readonly _data: AggregateSpecData<T>
+    private readonly _data: AggregateSpecData<T>
   ) {}
 
   /**
    * 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<T> {
     return this._data;
@@ -113,10 +98,11 @@ export class AggregateQuerySnapshot<T extends AggregateSpec> {
  * whatever the server returns. If the server cannot be reached then the
  * returned promise will be rejected.
  *
- * This is a convenience shorthand for:
- * getAggregateFromServer(query, { count: count() }).
+ * @param query - The `Query` to execute.
+ *
+ * @returns An `AggregateQuerySnapshot` that contains the number of documents.
  */
-export function getCountFromServer(
+export function getCount(
   query: Query<unknown>
 ): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
   const firestore = cast(query.firestore, Firestore);
@@ -151,6 +137,11 @@ export function getCountFromServer(
  * 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 aggregateSnapshotEqual<T extends AggregateSpec>(
   left: AggregateQuerySnapshot<T>,

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

@@ -40,6 +40,7 @@ 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;
 

packages/firestore/src/remote/datastore.ts

@@ -18,7 +18,6 @@
 import { CredentialsProvider } from '../api/credentials';
 import { User } from '../auth/user';
 import { Query, queryToTarget } from '../core/query';
-import { AggregateSpec } from '../lite-api/aggregate';
 import { Document } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
@@ -251,10 +250,11 @@ export async function invokeRunAggregationQueryRpc(
   if (!datastoreImpl.connection.shouldResourcePathBeIncludedInRequest) {
     delete request.parent;
   }
+  const EXPECTED_RESPONSE_COUNT = 1;
   const response = await datastoreImpl.invokeStreamingRPC<
     ProtoRunAggregationQueryRequest,
     ProtoRunAggregationQueryResponse
-  >('RunAggregationQuery', parent!, request);
+  >('RunAggregationQuery', parent!, request, EXPECTED_RESPONSE_COUNT);
   return (
     response
       // Omit RunAggregationQueryResponse that only contain readTimes.

packages/firestore/src/remote/serializer.ts

@@ -35,7 +35,6 @@ import {
   Target
 } from '../core/target';
 import { TargetId } from '../core/types';
-import { AggregateSpec } from '../lite-api/aggregate';
 import { Timestamp } from '../lite-api/timestamp';
 import { TargetData, TargetPurpose } from '../local/target_data';
 import { MutableDocument } from '../model/document';
@@ -859,6 +858,7 @@ export function toRunAggregationQueryRequest(
   target: Target
 ): ProtoRunAggregationQueryRequest {
   const queryTarget = toQueryTarget(serializer, target);
+
   return {
     structuredAggregationQuery: {
       aggregations: [

packages/firestore/test/integration/api/aggregation.test.ts

@@ -109,90 +109,4 @@ apiDescribe('Count query', (persistence: boolean) => {
       );
     });
   });
-
-  it('aggregateQuery.query equals to original query', () => {
-    return withEmptyTestCollection(persistence, async coll => {
-      const query_ = query(coll);
-      const aggregateQuery_ = countQuery(query_);
-      expect(aggregateQuery_.query).to.be.equal(query_);
-    });
-  });
-
-  it('aggregateQuerySnapshot.query equals to aggregateQuery', () => {
-    return withEmptyTestCollection(persistence, async coll => {
-      const aggregateQuery_ = countQuery(coll);
-      const snapshot = await getAggregateFromServerDirect(aggregateQuery_);
-      expect(snapshot.query).to.be.equal(aggregateQuery_);
-    });
-  });
-
-  it('aggregate query supports withConverter', () => {
-    const testDocs = {
-      a: { author: 'authorA', title: 'titleA' },
-      b: { author: 'authorB', title: 'titleB' }
-    };
-    return withTestCollection(persistence, testDocs, async coll => {
-      const query_ = query(
-        coll,
-        where('author', '==', 'authorA')
-      ).withConverter(postConverter);
-      const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
-      expect(snapshot.getCount()).to.equal(1);
-    });
-  });
-
-  it('aggregate 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 countQuery_ = countQuery(collectionGroup(db, collectionGroupId));
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
-      expect(snapshot.getCount()).to.equal(2);
-    });
-  });
-
-  it('aggregate query fails if firestore is terminated', () => {
-    return withEmptyTestCollection(persistence, async (coll, firestore) => {
-      await terminate(firestore);
-      const countQuery_ = countQuery(coll);
-      expect(() => getAggregateFromServerDirect(countQuery_)).to.throw(
-        'The client has already been terminated.'
-      );
-    });
-  });
-
-  it("terminate doesn't crash when there is aggregate query in flight", () => {
-    return withEmptyTestCollection(persistence, async (coll, firestore) => {
-      const countQuery_ = countQuery(coll);
-      void getAggregateFromServerDirect(countQuery_);
-      await terminate(firestore);
-    });
-  });
-
-  it('getAggregateFromServerDirect fails if user is offline', () => {
-    return withEmptyTestCollection(
-      persistence,
-      async (collection, firestore) => {
-        await disableNetwork(firestore);
-        const countQuery_ = countQuery(collection);
-        await expect(
-          getAggregateFromServerDirect(countQuery_)
-        ).to.be.eventually.rejectedWith(
-          'Failed to get aggregate result because the client is offline'
-        );
-      }
-    );
-  });
 });