update the firestore-client with datastore from client

Author: milaGGL

packages/firestore/src/api/aggregate.ts

@@ -21,6 +21,7 @@ import { AggregateField, AggregateQuerySnapshot } from '../lite-api/aggregate';
 import { cast } from '../util/input_validation';
 
 import { ensureFirestoreConfigured, Firestore } from './database';
+import { ExpUserDataWriter } from './reference_impl';
 
 export {
   AggregateField,
@@ -45,5 +46,6 @@ export function getCountFromServer(
 ): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
   const firestore = cast(query.firestore, Firestore);
   const client = ensureFirestoreConfigured(firestore);
-  return firestoreClientRunCountQuery(client, query);
+  const userDataWriter = new ExpUserDataWriter(firestore);
+  return firestoreClientRunCountQuery(client, query, userDataWriter);
 }

packages/firestore/src/core/count_query_runner.ts

@@ -0,0 +1,67 @@
+/**
+ * @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 { AbstractUserDataWriter, Query } from '../api';
+import { AggregateField, AggregateQuerySnapshot } from '../lite-api/aggregate';
+import { Value } from '../protos/firestore_proto_api';
+import { Datastore, invokeRunAggregationQueryRpc } from '../remote/datastore';
+import { hardAssert } from '../util/assert';
+
+/**
+ * CountQueryRunner encapsulates the logic needed to run the count aggregation
+ * queries.
+ */
+export class CountQueryRunner<T> {
+  constructor(
+    private readonly query: Query<unknown>,
+    private readonly datastore: Datastore,
+    private readonly userDataWriter: AbstractUserDataWriter
+  ) {}
+
+  run(): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+    return invokeRunAggregationQueryRpc(this.datastore, this.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]) =>
+            this.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> }>(
+            this.query,
+            {
+              count: countValue
+            }
+          )
+        );
+      }
+    );
+  }
+}

packages/firestore/src/core/firestore_client.ts

@@ -16,18 +16,17 @@
  */
 
 import { GetOptions } from '@firebase/firestore-types';
-
+import { AbstractUserDataWriter } 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 {
-  AggregateField,
-  AggregateQuerySnapshot,
-  getCount
-} from '../lite-api/aggregate';
 import { Query as LiteQuery } from '../lite-api/reference';
 import { LocalStore } from '../local/local_store';
 import {
@@ -68,6 +67,7 @@ import {
   OfflineComponentProvider,
   OnlineComponentProvider
 } from './component_provider';
+import { CountQueryRunner } from './count_query_runner';
 import { DatabaseId, DatabaseInfo } from './database_info';
 import {
   addSnapshotsInSyncListener,
@@ -510,7 +510,8 @@ export function firestoreClientTransaction<T>(
 
 export function firestoreClientRunCountQuery(
   client: FirestoreClient,
-  query: LiteQuery<unknown>
+  query: LiteQuery<unknown>,
+  userDataWriter: AbstractUserDataWriter
 ): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
   const deferred = new Deferred<
     AggregateQuerySnapshot<{ count: AggregateField<number> }>
@@ -526,7 +527,12 @@ export function firestoreClientRunCountQuery(
           )
         );
       } else {
-        const result = await getCount(query);
+        const datastore = await getDatastore(client);
+        const result = new CountQueryRunner(
+          query,
+          datastore,
+          userDataWriter
+        ).run();
         deferred.resolve(result);
       }
     } catch (e) {

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

@@ -16,10 +16,7 @@
  */
 
 import { deepEqual } from '@firebase/util';
-
-import { Value } from '../protos/firestore_proto_api';
-import { invokeRunAggregationQueryRpc } from '../remote/datastore';
-import { hardAssert } from '../util/assert';
+import { CountQueryRunner } from '../core/count_query_runner';
 import { cast } from '../util/input_validation';
 
 import { getDatastore } from './components';
@@ -108,29 +105,7 @@ export function getCount(
   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
-      })
-    );
-  });
+  return new CountQueryRunner(query, datastore, userDataWriter).run();
 }
 
 /**