Skip to content

Commit 6da4f4f

Browse files
milaGGLmilaGGL
authored
Mila/count export aggregate to api (#6575)
1 parent a3242ce commit 6da4f4f

File tree

9 files changed

+111
-27
lines changed

9 files changed

+111
-27
lines changed

‎packages/firestore/src/api/aggregate.ts

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,24 @@
1515
* limitations under the License.
1616
*/
1717

18+
import { firestoreClientRunAggregationQuery } from '../core/firestore_client';
19+
import { AggregateQuery, AggregateQuerySnapshot } from '../lite-api/aggregate';
20+
import { cast } from '../util/input_validation';
21+
22+
import { ensureFirestoreConfigured, Firestore } from './database';
23+
1824
export {
1925
AggregateQuery,
2026
AggregateQuerySnapshot,
2127
aggregateQueryEqual,
2228
aggregateQuerySnapshotEqual,
23-
countQuery,
24-
getAggregateFromServerDirect
29+
countQuery
2530
} from '../lite-api/aggregate';
31+
32+
export function getAggregateFromServerDirect(
33+
query: AggregateQuery
34+
): Promise<AggregateQuerySnapshot> {
35+
const firestore = cast(query.query.firestore, Firestore);
36+
const client = ensureFirestoreConfigured(firestore);
37+
return firestoreClientRunAggregationQuery(client, query);
38+
}

‎packages/firestore/src/core/firestore_client.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,14 @@
1717

1818
import { GetOptions } from '@firebase/firestore-types';
1919

20+
import { AggregateQuery, AggregateQuerySnapshot } from '../api';
2021
import { LoadBundleTask } from '../api/bundle';
2122
import {
2223
CredentialChangeListener,
2324
CredentialsProvider
2425
} from '../api/credentials';
2526
import { User } from '../auth/user';
27+
import { getAggregate } from '../lite-api/aggregate';
2628
import { LocalStore } from '../local/local_store';
2729
import {
2830
localStoreExecuteQuery,
@@ -501,6 +503,15 @@ export function firestoreClientTransaction<T>(
501503
return deferred.promise;
502504
}
503505

506+
export function firestoreClientRunAggregationQuery(
507+
client: FirestoreClient,
508+
query: AggregateQuery
509+
): Promise<AggregateQuerySnapshot> {
510+
return client.asyncQueue.enqueue(() => {
511+
return getAggregate(query);
512+
});
513+
}
514+
504515
async function readDocumentFromCache(
505516
localStore: LocalStore,
506517
docKey: DocumentKey,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ export function countQuery(query: Query<unknown>): AggregateQuery {
7373
return new AggregateQuery(query);
7474
}
7575

76-
export function getAggregateFromServerDirect(
76+
export function getAggregate(
7777
query: AggregateQuery
7878
): Promise<AggregateQuerySnapshot> {
7979
const firestore = cast(query.query.firestore, Firestore);

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,10 @@ export class GrpcConnection implements Connection {
8585
// We cache stubs for the most-recently-used token.
8686
private cachedStub: GeneratedGrpcStub | null = null;
8787

88+
get shouldResourcePathBeIncludedInRequest(): boolean {
89+
return true;
90+
}
91+
8892
constructor(protos: grpc.GrpcObject, private databaseInfo: DatabaseInfo) {
8993
// eslint-disable-next-line @typescript-eslint/no-explicit-any
9094
this.firestore = (protos as any)['google']['firestore']['v1'];

‎packages/firestore/src/remote/connection.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,14 @@ export interface Connection {
8484
appCheckToken: Token | null
8585
): Stream<Req, Resp>;
8686

87+
/**
88+
* Returns whether or not the implementation requires that the "path" of the resource
89+
* (a document or a collection) be present in the request message. If true, then the
90+
* request message must include the path. If false, then the request message must NOT
91+
* include the path.
92+
*/
93+
get shouldResourcePathBeIncludedInRequest(): boolean;
94+
8795
// TODO(mcg): subscribe to connection state changes.
8896
}
8997

‎packages/firestore/src/remote/datastore.ts

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -246,12 +246,15 @@ export async function invokeRunAggregationQueryRpc(
246246
datastoreImpl.serializer,
247247
queryToTarget(aggregateQuery.query._query)
248248
);
249+
250+
const parent = request.parent;
251+
if (!datastoreImpl.connection.shouldResourcePathBeIncludedInRequest) {
252+
delete request.parent;
253+
}
249254
const response = await datastoreImpl.invokeStreamingRPC<
250255
ProtoRunAggregationQueryRequest,
251256
ProtoRunAggregationQueryResponse
252-
>('RunAggregationQuery', request.parent!, {
253-
structuredAggregationQuery: request.structuredAggregationQuery
254-
});
257+
>('RunAggregationQuery', parent!, request);
255258
return (
256259
response
257260
// Omit RunAggregationQueryResponse that only contain readTimes.

‎packages/firestore/src/remote/rest_connection.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,10 @@ export abstract class RestConnection implements Connection {
5555
protected readonly baseUrl: string;
5656
private readonly databaseRoot: string;
5757

58+
get shouldResourcePathBeIncludedInRequest(): boolean {
59+
return false;
60+
}
61+
5862
constructor(private readonly databaseInfo: DatabaseInfo) {
5963
this.databaseId = databaseInfo.databaseId;
6064
const proto = databaseInfo.ssl ? 'https' : 'http';

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
/**
2+
* @license
3+
* Copyright 2022 Google LLC
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
import { expect } from 'chai';
19+
20+
import {
21+
countQuery,
22+
getAggregateFromServerDirect,
23+
query
24+
} from '../util/firebase_export';
25+
import { apiDescribe, withTestCollection } from '../util/helpers';
26+
27+
apiDescribe('Aggregation query', (persistence: boolean) => {
28+
it('can run count query getAggregateFromServerDirect', () => {
29+
const testDocs = {
30+
a: { k: 'a', sort: 1 },
31+
b: { k: 'b', sort: 2 },
32+
c: { k: 'c', sort: 2 }
33+
};
34+
return withTestCollection(persistence, testDocs, async coll => {
35+
const query_ = query(coll);
36+
const countQuery_ = countQuery(query_);
37+
const snapshot = await getAggregateFromServerDirect(countQuery_);
38+
expect(snapshot.getCount()).to.equal(3);
39+
});
40+
});
41+
});

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

Lines changed: 21 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import chaiAsPromised from 'chai-as-promised';
2222

2323
import {
2424
countQuery,
25-
getAggregateFromServerDirect,
25+
getAggregate,
2626
aggregateQueryEqual,
2727
aggregateQuerySnapshotEqual
2828
} from '../../src/lite-api/aggregate';
@@ -2050,7 +2050,7 @@ describe('countQuery()', () => {
20502050
return withTestCollection(async coll => {
20512051
const query_ = query(coll);
20522052
const countQuery_ = countQuery(query_);
2053-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2053+
const snapshot = await getAggregate(countQuery_);
20542054
expect(countQuery_.query).to.equal(query_);
20552055
expect(snapshot.query).to.equal(countQuery_);
20562056
expect(snapshot.query.query).to.equal(query_);
@@ -2060,7 +2060,7 @@ describe('countQuery()', () => {
20602060
it('empty test collection count', () => {
20612061
return withTestCollection(async coll => {
20622062
const countQuery_ = countQuery(query(coll));
2063-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2063+
const snapshot = await getAggregate(countQuery_);
20642064
expect(snapshot.getCount()).to.equal(0);
20652065
});
20662066
});
@@ -2073,7 +2073,7 @@ describe('countQuery()', () => {
20732073
];
20742074
return withTestCollectionAndInitialData(testDocs, async collection => {
20752075
const countQuery_ = countQuery(query(collection));
2076-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2076+
const snapshot = await getAggregate(countQuery_);
20772077
expect(snapshot.getCount()).to.equal(3);
20782078
});
20792079
});
@@ -2087,7 +2087,7 @@ describe('countQuery()', () => {
20872087
return withTestCollectionAndInitialData(testDocs, async collection => {
20882088
const query_ = query(collection, where('author', '==', 'authorA'));
20892089
const countQuery_ = countQuery(query_);
2090-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2090+
const snapshot = await getAggregate(countQuery_);
20912091
expect(snapshot.getCount()).to.equal(2);
20922092
});
20932093
});
@@ -2105,7 +2105,7 @@ describe('countQuery()', () => {
21052105
limit(1)
21062106
);
21072107
const countQuery_ = countQuery(query_);
2108-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2108+
const snapshot = await getAggregate(countQuery_);
21092109
expect(snapshot.getCount()).to.equal(1);
21102110
});
21112111
});
@@ -2123,7 +2123,7 @@ describe('countQuery()', () => {
21232123
limit(3)
21242124
);
21252125
const countQuery_ = countQuery(query_);
2126-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2126+
const snapshot = await getAggregate(countQuery_);
21272127
expect(snapshot.getCount()).to.equal(2);
21282128
});
21292129
});
@@ -2138,7 +2138,7 @@ describe('countQuery()', () => {
21382138
return withTestCollectionAndInitialData(testDocs, async collection => {
21392139
const query_ = query(collection, orderBy('title'));
21402140
const countQuery_ = countQuery(query_);
2141-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2141+
const snapshot = await getAggregate(countQuery_);
21422142
expect(snapshot.getCount()).to.equal(3);
21432143
});
21442144
});
@@ -2153,7 +2153,7 @@ describe('countQuery()', () => {
21532153
return withTestCollectionAndInitialData(testDocs, async collection => {
21542154
const query_ = query(collection, orderBy('id'), startAt(2));
21552155
const countQuery_ = countQuery(query_);
2156-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2156+
const snapshot = await getAggregate(countQuery_);
21572157
expect(snapshot.getCount()).to.equal(2);
21582158
});
21592159
});
@@ -2168,7 +2168,7 @@ describe('countQuery()', () => {
21682168
return withTestCollectionAndInitialData(testDocs, async collection => {
21692169
const query_ = query(collection, orderBy('id'), startAfter(2));
21702170
const countQuery_ = countQuery(query_);
2171-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2171+
const snapshot = await getAggregate(countQuery_);
21722172
expect(snapshot.getCount()).to.equal(1);
21732173
});
21742174
});
@@ -2183,7 +2183,7 @@ describe('countQuery()', () => {
21832183
return withTestCollectionAndInitialData(testDocs, async collection => {
21842184
const query_ = query(collection, orderBy('id'), startAt(1), endAt(2));
21852185
const countQuery_ = countQuery(query_);
2186-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2186+
const snapshot = await getAggregate(countQuery_);
21872187
expect(snapshot.getCount()).to.equal(2);
21882188
});
21892189
});
@@ -2198,7 +2198,7 @@ describe('countQuery()', () => {
21982198
return withTestCollectionAndInitialData(testDocs, async collection => {
21992199
const query_ = query(collection, orderBy('id'), startAt(1), endBefore(2));
22002200
const countQuery_ = countQuery(query_);
2201-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2201+
const snapshot = await getAggregate(countQuery_);
22022202
expect(snapshot.getCount()).to.equal(1);
22032203
});
22042204
});
@@ -2215,7 +2215,7 @@ describe('countQuery()', () => {
22152215
where('author', '==', 'authorA')
22162216
).withConverter(postConverter);
22172217
const countQuery_ = countQuery(query_);
2218-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2218+
const snapshot = await getAggregate(countQuery_);
22192219
expect(snapshot.getCount()).to.equal(2);
22202220
});
22212221
});
@@ -2236,7 +2236,7 @@ describe('countQuery()', () => {
22362236
}
22372237
await batch.commit();
22382238
const countQuery_ = countQuery(collectionGroup(db, collectionGroupId));
2239-
const snapshot = await getAggregateFromServerDirect(countQuery_);
2239+
const snapshot = await getAggregate(countQuery_);
22402240
expect(snapshot.getCount()).to.equal(2);
22412241
});
22422242
});
@@ -2283,9 +2283,9 @@ describe('countQuery()', () => {
22832283
const countQuery1A = countQuery(query1);
22842284
const countQuery1B = countQuery(query1);
22852285
const countQuery2 = countQuery(query2);
2286-
const snapshot1A = await getAggregateFromServerDirect(countQuery1A);
2287-
const snapshot1B = await getAggregateFromServerDirect(countQuery1B);
2288-
const snapshot2 = await getAggregateFromServerDirect(countQuery2);
2286+
const snapshot1A = await getAggregate(countQuery1A);
2287+
const snapshot1B = await getAggregate(countQuery1B);
2288+
const snapshot2 = await getAggregate(countQuery2);
22892289
expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot1B)).to.be.true;
22902290
expect(aggregateQuerySnapshotEqual(snapshot1A, snapshot2)).to.be.true;
22912291
});
@@ -2302,8 +2302,8 @@ describe('countQuery()', () => {
23022302
const query2 = query(collection, where('author', '==', 'authorB'));
23032303
const countQuery1 = countQuery(query1);
23042304
const countQuery2 = countQuery(query2);
2305-
const snapshot1 = await getAggregateFromServerDirect(countQuery1);
2306-
const snapshot2 = await getAggregateFromServerDirect(countQuery2);
2305+
const snapshot1 = await getAggregate(countQuery1);
2306+
const snapshot2 = await getAggregate(countQuery2);
23072307
expect(aggregateQuerySnapshotEqual(snapshot1, snapshot2)).to.be.false;
23082308
});
23092309
});
@@ -2312,7 +2312,7 @@ describe('countQuery()', () => {
23122312
return withTestCollection(async collection => {
23132313
await terminate(collection.firestore);
23142314
const countQuery_ = countQuery(query(collection));
2315-
expect(() => getAggregateFromServerDirect(countQuery_)).to.throw(
2315+
expect(() => getAggregate(countQuery_)).to.throw(
23162316
'The client has already been terminated.'
23172317
);
23182318
});
@@ -2326,7 +2326,7 @@ describe('countQuery()', () => {
23262326
];
23272327
return withTestCollectionAndInitialData(testDocs, async collection => {
23282328
const countQuery_ = countQuery(query(collection));
2329-
const promise = getAggregateFromServerDirect(countQuery_);
2329+
const promise = getAggregate(countQuery_);
23302330
await terminate(collection.firestore);
23312331
const snapshot = await promise;
23322332
expect(snapshot.getCount()).to.equal(3);

0 commit comments

Comments
 (0)