Mila/count (#6597)

.changeset/hot-insects-wink.md

@@ -0,0 +1,5 @@
+---
+'@firebase/firestore': minor
+---
+
+Implement count query for internal use.

packages/firestore/src/api/aggregate.ts

@@ -0,0 +1,39 @@
+/**
+ * @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 { 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';
+
+/**
+ * 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 firestoreClientRunCountQuery(client, query);
+}

packages/firestore/src/core/firestore_client.ts

@@ -23,6 +23,12 @@ import {
   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 {
   localStoreExecuteQuery,
@@ -38,6 +44,7 @@ import { toByteStreamReader } from '../platform/byte_stream_reader';
 import { newSerializer, newTextEncoder } from '../platform/serializer';
 import { Datastore } from '../remote/datastore';
 import {
+  canUseNetwork,
   RemoteStore,
   remoteStoreDisableNetwork,
   remoteStoreEnableNetwork,
@@ -501,6 +508,34 @@ export function firestoreClientTransaction<T>(
   return deferred.promise;
 }
 
+export function firestoreClientRunCountQuery(
+  client: FirestoreClient,
+  query: LiteQuery<unknown>
+): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+  const deferred = new Deferred<
+    AggregateQuerySnapshot<{ count: AggregateField<number> }>
+  >();
+  client.asyncQueue.enqueueAndForget(async () => {
+    try {
+      const remoteStore = await getRemoteStore(client);
+      if (!canUseNetwork(remoteStore)) {
+        deferred.reject(
+          new FirestoreError(
+            Code.UNAVAILABLE,
+            'Failed to get count result because the client is offline.'
+          )
+        );
+      } else {
+        const result = await getCount(query);
+        deferred.resolve(result);
+      }
+    } catch (e) {
+      deferred.reject(e as Error);
+    }
+  });
+  return deferred.promise;
+}
+
 async function readDocumentFromCache(
   localStore: LocalStore,
   docKey: DocumentKey,

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

@@ -0,0 +1,153 @@
+/**
+ * @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 { 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 { getDatastore } from './components';
+import { Firestore } from './database';
+import { Query, queryEqual } from './reference';
+import { LiteUserDataWriter } from './reference_impl';
+
+/**
+ * An `AggregateField`that captures input type T.
+ */
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export class AggregateField<T> {
+  type = 'AggregateField';
+}
+
+/**
+ * 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 = 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<T extends AggregateSpec> {
+  readonly type = 'AggregateQuerySnapshot';
+
+  /** @hideconstructor */
+  constructor(
+    readonly query: Query<unknown>,
+    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;
+  }
+}
+
+/**
+ * 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 `AggregateQuerySnapshot` that contains the number of documents.
+ */
+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._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 countValue = counts[0];
+
+    hardAssert(
+      typeof countValue === 'number',
+      'Count aggregate field value is not a number: ' + countValue
+    );
+
+    return Promise.resolve(
+      new AggregateQuerySnapshot<{ count: AggregateField<number> }>(query, {
+        count: countValue
+      })
+    );
+  });
+}
+
+/**
+ * Compares two `AggregateQuerySnapshot` instances for equality.
+ * Two `AggregateQuerySnapshot` instances are considered "equal" if they have
+ * the same underlying query, and the same data.
+ *
+ * @param left - The `AggregateQuerySnapshot` to compare.
+ * @param right - The `AggregateQuerySnapshot` to compare.
+ *
+ * @returns true if the AggregateQuerySnapshots are equal.
+ */
+export function aggregateQuerySnapshotEqual<T extends AggregateSpec>(
+  left: AggregateQuerySnapshot<T>,
+  right: AggregateQuerySnapshot<T>
+): boolean {
+  return (
+    queryEqual(left.query, right.query) && deepEqual(left.data(), right.data())
+  );
+}

packages/firestore/src/platform/node/grpc_connection.ts

@@ -85,6 +85,12 @@ export class GrpcConnection implements Connection {
   // We cache stubs for the most-recently-used token.
   private cachedStub: GeneratedGrpcStub | null = null;
 
+  get shouldResourcePathBeIncludedInRequest(): boolean {
+    // Both `invokeRPC()` and `invokeStreamingRPC()` ignore their `path` arguments, and expect
+    // the "path" to be part of the given `request`.
+    return true;
+  }
+
   constructor(protos: grpc.GrpcObject, private databaseInfo: DatabaseInfo) {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     this.firestore = (protos as any)['google']['firestore']['v1'];

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/protos/google/firestore/v1/aggregation_result.proto

@@ -0,0 +1,42 @@
+// 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.
+
+syntax = "proto3";
+
+package google.firestore.v1;
+
+import "google/firestore/v1/document.proto";
+
+option csharp_namespace = "Google.Cloud.Firestore.V1";
+option go_package = "google.golang.org/genproto/googleapis/firestore/v1;firestore";
+option java_multiple_files = true;
+option java_outer_classname = "AggregationResultProto";
+option java_package = "com.google.firestore.v1";
+option objc_class_prefix = "GCFS";
+option php_namespace = "Google\\Cloud\\Firestore\\V1";
+option ruby_package = "Google::Cloud::Firestore::V1";
+
+// The result of a single bucket from a Firestore aggregation query.
+//
+// The keys of `aggregate_fields` are the same for all results in an aggregation
+// query, unlike document queries which can have different fields present for
+// each result.
+message AggregationResult {
+  // The result of the aggregation functions, ex: `COUNT(*) AS total_docs`.
+  //
+  // The key is the [alias][google.firestore.v1.StructuredAggregationQuery.Aggregation.alias]
+  // assigned to the aggregation function on input and the size of this map
+  // equals the number of aggregation functions in the query.
+  map<string, Value> aggregate_fields = 2;
+}