Mila/count implement count query

common/api-review/firestore.api.md

@@ -17,18 +17,17 @@ export type AddPrefixToKeys<Prefix extends string, T extends Record<string, unkn
     [K in keyof T & string as `${Prefix}.${K}`]+?: T[K];
 };
 
-// @public (undocumented)
-export class AggregateQuery<T = DocumentData> {
-    // (undocumented)
-    readonly query: Query<T>;
+// @public
+export class AggregateQuery {
+    readonly query: Query<unknown>;
     // (undocumented)
     readonly type = "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<unknown>): AggregateQuery;
 
 // @public
 export function deleteDoc(reference: DocumentReference<unknown>): Promise<void>;

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

@@ -15,40 +15,95 @@
  * 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<T = DocumentData> {
+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<T>;
+  /**
+   * The query on which you called `countQuery` in order to get this `AggregateQuery`.
+   * Query type is set to unknown to avoid error caused by query type converter.
+   * might change it back to T after testing if the error do exist or not
+   */
+  readonly query: Query<unknown>;
 
   /** @hideconstructor */
-  constructor(query: Query<T>) {
+  constructor(query: Query<unknown>) {
     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<unknown>): AggregateQuery {
   return new AggregateQuery(query);
 }
 
 export function getAggregateFromServerDirect(
   query: AggregateQuery
 ): Promise<AggregateQuerySnapshot> {
-  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(
+      count !== undefined,
+      'count_alias property is missing from result.'
+    );
+    hardAssert(
+      typeof count === 'number',
+      'Count aggeragte field is not a number: ' + count
+    );
+
+    return Promise.resolve(new AggregateQuerySnapshot(query, count));
+  });
 }
 
 export function aggregateQueryEqual(

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<Value>;
+  }
+  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;

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<ProtoValue[]> {
+  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,

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<Req, Resp>(rpcName, url, headers, req).then(
       response => {
         logDebug(LOG_TAG, 'Received: ', response);

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!);