Mila/count implement count query

common/api-review/firestore.api.md

@@ -17,21 +17,23 @@ 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 {
+    getQuery(): Query<unknown>;
+    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)
+    readonly count: number | null;
     getCount(): number | null;
+    getQuery(): AggregateQuery;
     // (undocumented)
     readonly query: AggregateQuery;
     // (undocumented)
@@ -93,8 +95,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>;
@@ -237,7 +239,7 @@ export class GeoPoint {
 }
 
 // @public (undocumented)
-export function getAggregateFromServerDirect(query: AggregateQuery): Promise<AggregateQuerySnapshot>;
+export function getAggregateFromServerDirect(aggregateQuery: AggregateQuery): Promise<AggregateQuerySnapshot>;
 
 // @public
 export function getDoc<T>(reference: DocumentReference<T>): Promise<DocumentSnapshot<T>>;

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

@@ -15,40 +15,115 @@
  * limitations under the License.
  */
 
-import { DocumentData, Query, queryEqual } from './reference';
+import { invokeRunAggregationQueryRpc } from '../remote/datastore';
+import { cast } from '../util/input_validation';
+import { getDatastore } from './components';
+import { Firestore } from './database';
+import { Query, queryEqual } from './reference';
+import { LiteUserDataWriter } from './reference_impl';
 
-export class AggregateQuery<T = DocumentData> {
+/**
+ * A {@code AggregateQuery} computes some aggregation statistics from the result set of a base
+ * {@link Query}.
+ *
+ * <p><b>Subclassing Note</b>: Cloud Firestore classes are not meant to be subclassed except for use
+ * in test mocks. Subclassing is not supported in production code and new SDK releases may break
+ * code that does so.
+ */
+export class AggregateQuery {
   readonly type = 'AggregateQuery';
-  readonly query: Query<T>;
+  /**
+   * The query on which you called {@link 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;
   }
+
+  /** Returns the base {@link Query} for this aggregate query. */
+  getQuery(): Query<unknown> {
+    return this.query;
+  }
 }
 
+/**
+ * A {@code AggregateQuerySnapshot} contains results of a {@link AggregateQuery}.
+ *
+ * <p><b>Subclassing Note</b>: Cloud Firestore classes are not meant to be subclassed except for use
+ * in test mocks. Subclassing is not supported in production code and new SDK releases may break
+ * code that does so.
+ */
 export class AggregateQuerySnapshot {
   readonly type = 'AggregateQuerySnapshot';
   readonly query: AggregateQuery;
+  readonly count: number | null;
 
   /** @hideconstructor */
-  constructor(query: AggregateQuery, readonly _count: number) {
+  constructor(query: AggregateQuery, count: number | null) {
     this.query = query;
+    this.count = count;
   }
 
+  /** @return The original {@link AggregateQuery} this snapshot is a result of. */
+  getQuery(): AggregateQuery {
+    return this.query;
+  }
+
+  /**
+   * @return The result of a document count aggregation. Returns null if no count aggregation is
+   * available in the result.
+   */
   getCount(): number | null {
-    return this._count;
+    return this.count;
   }
 }
 
-export function countQuery(query: Query): AggregateQuery {
+/**
+ * Creates an {@link AggregateQuery} counting the number of documents matching this query.
+ *
+ * @return An {@link 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 {
+  /**
+   * TODO(mila): add the "count" aggregateField to the params after the AggregateQuery is updated.
+   */
   return new AggregateQuery(query);
 }
 
 export function getAggregateFromServerDirect(
-  query: AggregateQuery
+  aggregateQuery: AggregateQuery
 ): Promise<AggregateQuerySnapshot> {
-  return Promise.resolve(new AggregateQuerySnapshot(query, 42));
+  const firestore = cast(aggregateQuery.query.firestore, Firestore);
+  const datastore = getDatastore(firestore);
+  const userDataWriter = new LiteUserDataWriter(firestore);
+
+  return invokeRunAggregationQueryRpc(datastore, aggregateQuery).then(
+    result => {
+      const aggregationFields = new Map();
+      /**
+       * while getting aggregation fields from server direct, it should have only
+       * one RunAggregationQueryResponse returned.
+       * But we used streaming rpc here, so we will have an array of
+       * (one, or possibly more) RunAggregationQueryResponse. For this specific
+       * function, we get the first RunAggregationQueryResponse only.
+       */
+      for (const [key, value] of Object.entries(result[0])) {
+        aggregationFields.set(key, userDataWriter.convertValue(value));
+      }
+      return Promise.resolve(
+        new AggregateQuerySnapshot(
+          aggregateQuery,
+          aggregationFields.get('count_alias')
+        )
+      );
+    }
+  );
 }
 
 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 RunQueryResponses that only contain readTimes.
+      .filter(proto => !!proto.result && !!proto.result.aggregateFields)
+      .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

@@ -82,7 +82,8 @@ import {
   TargetChangeTargetChangeType as ProtoTargetChangeTargetChangeType,
   Timestamp as ProtoTimestamp,
   Write as ProtoWrite,
-  WriteResult as ProtoWriteResult
+  WriteResult as ProtoWriteResult,
+  RunAggregationQueryRequest as ProtoRunAggregationQueryRequest
 } from '../protos/firestore_proto_api';
 import { debugAssert, fail, hardAssert } from '../util/assert';
 import { ByteString } from '../util/byte_string';
@@ -852,6 +853,72 @@ export function toQueryTarget(
   return result;
 }
 
+export function toRunAggregationQueryRequest(
+  serializer: JsonProtoSerializer,
+  target: Target
+): ProtoRunAggregationQueryRequest {
+  // Dissect the path into parent, collectionId, and optional key filter.
+  const result: ProtoRunAggregationQueryRequest = {
+    structuredAggregationQuery: {
+      aggregations: [
+        {
+          count: {
+          },
+          alias:"count_alias"
+        },
+      ],
+      structuredQuery: {}
+    }
+  };
+  const path = target.path;
+  if (target.collectionGroup !== null) {
+    debugAssert(
+      path.length % 2 === 0,
+      'Collection Group queries should be within a document path or root.'
+    );
+    result.parent = toQueryPath(serializer, path);
+    result.structuredAggregationQuery!.structuredQuery!.from = [
+      {
+        collectionId: target.collectionGroup,
+        allDescendants: true
+      }
+    ];
+  } else {
+    debugAssert(
+      path.length % 2 !== 0,
+      'Document queries with filters are not supported.'
+    );
+    result.parent = toQueryPath(serializer, path.popLast());
+    result.structuredAggregationQuery!.structuredQuery!.from = [
+      { collectionId: path.lastSegment() }
+    ];
+  }
+
+  const where = toFilter(target.filters);
+  if (where) {
+    result.structuredAggregationQuery!.structuredQuery!.where = where;
+  }
+
+  const orderBy = toOrder(target.orderBy);
+  if (orderBy) {
+    result.structuredAggregationQuery!.structuredQuery!.orderBy = orderBy;
+  }
+
+  const limit = toInt32Proto(serializer, target.limit);
+  if (limit !== null) {
+    result.structuredAggregationQuery!.structuredQuery!.limit = limit;
+  }
+
+  if (target.startAt) {
+    result.structuredAggregationQuery!.structuredQuery!.startAt = toStartAtCursor(target.startAt);
+  }
+  if (target.endAt) {
+    result.structuredAggregationQuery!.structuredQuery!.endAt = toEndAtCursor(target.endAt);
+  }
+
+  return result;
+}
+
 export function convertQueryTargetToQuery(target: ProtoQueryTarget): Query {
   let path = fromQueryPath(target.parent!);