Mila/count export aggregate to api

packages/firestore/src/api/aggregate.ts

@@ -15,11 +15,24 @@
  * limitations under the License.
  */
 
+import { firestoreClientRunAggregationQuery } from '../core/firestore_client';
+import { AggregateQuery, AggregateQuerySnapshot } from '../lite-api/aggregate';
+import { cast } from '../util/input_validation';
+
+import { ensureFirestoreConfigured, Firestore } from './database';
+
 export {
   AggregateQuery,
   AggregateQuerySnapshot,
   aggregateQueryEqual,
   aggregateQuerySnapshotEqual,
-  countQuery,
-  getAggregateFromServerDirect
+  countQuery
 } from '../lite-api/aggregate';
+
+export function getAggregateFromServerDirect(
+  query: AggregateQuery
+): Promise<AggregateQuerySnapshot> {
+  const firestore = cast(query.query.firestore, Firestore);
+  const client = ensureFirestoreConfigured(firestore);
+  return firestoreClientRunAggregationQuery(client, query);
+}

packages/firestore/src/core/firestore_client.ts

@@ -17,12 +17,14 @@
 
 import { GetOptions } from '@firebase/firestore-types';
 
+import { AggregateQuery, AggregateQuerySnapshot } from '../api';
 import { LoadBundleTask } from '../api/bundle';
 import {
   CredentialChangeListener,
   CredentialsProvider
 } from '../api/credentials';
 import { User } from '../auth/user';
+import { getAggregate } from '../lite-api/aggregate';
 import { LocalStore } from '../local/local_store';
 import {
   localStoreExecuteQuery,
@@ -501,6 +503,15 @@ export function firestoreClientTransaction<T>(
   return deferred.promise;
 }
 
+export function firestoreClientRunAggregationQuery(
+  client: FirestoreClient,
+  query: AggregateQuery
+): Promise<AggregateQuerySnapshot> {
+  return client.asyncQueue.enqueue(() => {
+    return getAggregate(query);
+  });
+}
+
 async function readDocumentFromCache(
   localStore: LocalStore,
   docKey: DocumentKey,

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

@@ -73,7 +73,7 @@ export function countQuery(query: Query<unknown>): AggregateQuery {
   return new AggregateQuery(query);
 }
 
-export function getAggregateFromServerDirect(
+export function getAggregate(
   query: AggregateQuery
 ): Promise<AggregateQuerySnapshot> {
   const firestore = cast(query.query.firestore, Firestore);

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

@@ -85,6 +85,10 @@ export class GrpcConnection implements Connection {
   // We cache stubs for the most-recently-used token.
   private cachedStub: GeneratedGrpcStub | null = null;
 
+  get shouldResourcePathBeIncludedInRequest(): boolean {
+    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/remote/connection.ts

@@ -84,6 +84,14 @@ export interface Connection {
     appCheckToken: Token | null
   ): Stream<Req, Resp>;
 
+  /**
+   * Returns whether or not the implementation requires that the "path" of the resource
+   * (a document or a collection) be present in the request message. If true, then the
+   * request message must include the path. If false, then the request message must NOT
+   * include the path.
+   */
+  get shouldResourcePathBeIncludedInRequest(): boolean;
+
   // TODO(mcg): subscribe to connection state changes.
 }
 

packages/firestore/src/remote/datastore.ts

@@ -246,12 +246,15 @@ export async function invokeRunAggregationQueryRpc(
     datastoreImpl.serializer,
     queryToTarget(aggregateQuery.query._query)
   );
+
+  const parent = request.parent;
+  if (!datastoreImpl.connection.shouldResourcePathBeIncludedInRequest) {
+    delete request.parent;
+  }
   const response = await datastoreImpl.invokeStreamingRPC<
     ProtoRunAggregationQueryRequest,
     ProtoRunAggregationQueryResponse
-  >('RunAggregationQuery', request.parent!, {
-    structuredAggregationQuery: request.structuredAggregationQuery
-  });
+  >('RunAggregationQuery', parent!, request);
   return (
     response
       // Omit RunAggregationQueryResponse that only contain readTimes.

packages/firestore/src/remote/rest_connection.ts

@@ -55,6 +55,10 @@ export abstract class RestConnection implements Connection {
   protected readonly baseUrl: string;
   private readonly databaseRoot: string;
 
+  get shouldResourcePathBeIncludedInRequest(): boolean {
+    return false;
+  }
+
   constructor(private readonly databaseInfo: DatabaseInfo) {
     this.databaseId = databaseInfo.databaseId;
     const proto = databaseInfo.ssl ? 'https' : 'http';

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

@@ -0,0 +1,41 @@
+/**
+ * @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 { expect } from 'chai';
+
+import {
+  countQuery,
+  getAggregateFromServerDirect,
+  query
+} from '../util/firebase_export';
+import { apiDescribe, withTestCollection } from '../util/helpers';
+
+apiDescribe('Aggregation query', (persistence: boolean) => {
+  it('can run count query getAggregateFromServerDirect', () => {
+    const testDocs = {
+      a: { k: 'a', sort: 1 },
+      b: { k: 'b', sort: 2 },
+      c: { k: 'c', sort: 2 }
+    };
+    return withTestCollection(persistence, testDocs, async coll => {
+      const query_ = query(coll);
+      const countQuery_ = countQuery(query_);
+      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      expect(snapshot.getCount()).to.equal(3);
+    });
+  });
+});

packages/firestore/test/lite/integration.test.ts

@@ -22,7 +22,7 @@ import chaiAsPromised from 'chai-as-promised';
 
 import {
   countQuery,
-  getAggregateFromServerDirect,
+  getAggregate,
   aggregateQueryEqual,
   aggregateQuerySnapshotEqual
 } from '../../src/lite-api/aggregate';
@@ -2050,7 +2050,7 @@ describe('countQuery()', () => {
     return withTestCollection(async coll => {
       const query_ = query(coll);
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(countQuery_.query).to.equal(query_);
       expect(snapshot.query).to.equal(countQuery_);
       expect(snapshot.query.query).to.equal(query_);
@@ -2060,7 +2060,7 @@ describe('countQuery()', () => {
   it('empty test collection count', () => {
     return withTestCollection(async coll => {
       const countQuery_ = countQuery(query(coll));
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(0);
     });
   });
@@ -2073,7 +2073,7 @@ describe('countQuery()', () => {
     ];
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const countQuery_ = countQuery(query(collection));
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(3);
     });
   });
@@ -2087,7 +2087,7 @@ describe('countQuery()', () => {
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const query_ = query(collection, where('author', '==', 'authorA'));
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(2);
     });
   });
@@ -2105,7 +2105,7 @@ describe('countQuery()', () => {
         limit(1)
       );
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(1);
     });
   });
@@ -2123,7 +2123,7 @@ describe('countQuery()', () => {
         limit(3)
       );
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(2);
     });
   });
@@ -2138,7 +2138,7 @@ describe('countQuery()', () => {
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const query_ = query(collection, orderBy('title'));
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(3);
     });
   });
@@ -2153,7 +2153,7 @@ describe('countQuery()', () => {
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const query_ = query(collection, orderBy('id'), startAt(2));
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(2);
     });
   });
@@ -2168,7 +2168,7 @@ describe('countQuery()', () => {
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const query_ = query(collection, orderBy('id'), startAfter(2));
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(1);
     });
   });
@@ -2183,7 +2183,7 @@ describe('countQuery()', () => {
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const query_ = query(collection, orderBy('id'), startAt(1), endAt(2));
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(2);
     });
   });
@@ -2198,7 +2198,7 @@ describe('countQuery()', () => {
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const query_ = query(collection, orderBy('id'), startAt(1), endBefore(2));
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(1);
     });
   });
@@ -2215,7 +2215,7 @@ describe('countQuery()', () => {
         where('author', '==', 'authorA')
       ).withConverter(postConverter);
       const countQuery_ = countQuery(query_);
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(2);
     });
   });
@@ -2236,7 +2236,7 @@ describe('countQuery()', () => {
       }
       await batch.commit();
       const countQuery_ = countQuery(collectionGroup(db, collectionGroupId));
-      const snapshot = await getAggregateFromServerDirect(countQuery_);
+      const snapshot = await getAggregate(countQuery_);
       expect(snapshot.getCount()).to.equal(2);
     });
   });
@@ -2283,9 +2283,9 @@ describe('countQuery()', () => {
       const countQuery1A = countQuery(query1);
       const countQuery1B = countQuery(query1);
       const countQuery2 = countQuery(query2);
-      const snapshot1A = await getAggregateFromServerDirect(countQuery1A);
-      const snapshot1B = await getAggregateFromServerDirect(countQuery1B);
-      const snapshot2 = await getAggregateFromServerDirect(countQuery2);
+      const snapshot1A = await getAggregate(countQuery1A);
+      const snapshot1B = await getAggregate(countQuery1B);
+      const snapshot2 = await getAggregate(countQuery2);
       expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot1B)).to.be.true;
       expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot2)).to.be.true;
     });
@@ -2302,8 +2302,8 @@ describe('countQuery()', () => {
       const query2 = query(collection, where('author', '==', 'authorB'));
       const countQuery1 = countQuery(query1);
       const countQuery2 = countQuery(query2);
-      const snapshot1 = await getAggregateFromServerDirect(countQuery1);
-      const snapshot2 = await getAggregateFromServerDirect(countQuery2);
+      const snapshot1 = await getAggregate(countQuery1);
+      const snapshot2 = await getAggregate(countQuery2);
       expect(aggregateQuerySnapshotEqual(snapshot1, snapshot2)).to.be.false;
     });
   });
@@ -2312,7 +2312,7 @@ describe('countQuery()', () => {
     return withTestCollection(async collection => {
       await terminate(collection.firestore);
       const countQuery_ = countQuery(query(collection));
-      expect(() => getAggregateFromServerDirect(countQuery_)).to.throw(
+      expect(() => getAggregate(countQuery_)).to.throw(
         'The client has already been terminated.'
       );
     });
@@ -2326,7 +2326,7 @@ describe('countQuery()', () => {
     ];
     return withTestCollectionAndInitialData(testDocs, async collection => {
       const countQuery_ = countQuery(query(collection));
-      const promise = getAggregateFromServerDirect(countQuery_);
+      const promise = getAggregate(countQuery_);
       await terminate(collection.firestore);
       const snapshot = await promise;
       expect(snapshot.getCount()).to.equal(3);