Skip to content

Commit e35db6f

Browse files
authored
Mila/count (#6597)
1 parent 5b696de commit e35db6f

19 files changed

+1058
-3
lines changed

.changeset/hot-insects-wink.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@firebase/firestore': minor
3+
---
4+
5+
Implement count query for internal use.
+39
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
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 { Query } from '../api';
19+
import { firestoreClientRunCountQuery } from '../core/firestore_client';
20+
import { AggregateField, AggregateQuerySnapshot } from '../lite-api/aggregate';
21+
import { cast } from '../util/input_validation';
22+
23+
import { ensureFirestoreConfigured, Firestore } from './database';
24+
25+
/**
26+
* Executes the query and returns the results as a `AggregateQuerySnapshot` from the
27+
* server. Returns an error if the network is not available.
28+
*
29+
* @param query - The `Query` to execute.
30+
*
31+
* @returns A `Promise` that will be resolved with the results of the query.
32+
*/
33+
export function getCountFromServer(
34+
query: Query<unknown>
35+
): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
36+
const firestore = cast(query.firestore, Firestore);
37+
const client = ensureFirestoreConfigured(firestore);
38+
return firestoreClientRunCountQuery(client, query);
39+
}

packages/firestore/src/core/firestore_client.ts

+35
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,12 @@ import {
2323
CredentialsProvider
2424
} from '../api/credentials';
2525
import { User } from '../auth/user';
26+
import {
27+
AggregateField,
28+
AggregateQuerySnapshot,
29+
getCount
30+
} from '../lite-api/aggregate';
31+
import { Query as LiteQuery } from '../lite-api/reference';
2632
import { LocalStore } from '../local/local_store';
2733
import {
2834
localStoreExecuteQuery,
@@ -38,6 +44,7 @@ import { toByteStreamReader } from '../platform/byte_stream_reader';
3844
import { newSerializer, newTextEncoder } from '../platform/serializer';
3945
import { Datastore } from '../remote/datastore';
4046
import {
47+
canUseNetwork,
4148
RemoteStore,
4249
remoteStoreDisableNetwork,
4350
remoteStoreEnableNetwork,
@@ -501,6 +508,34 @@ export function firestoreClientTransaction<T>(
501508
return deferred.promise;
502509
}
503510

511+
export function firestoreClientRunCountQuery(
512+
client: FirestoreClient,
513+
query: LiteQuery<unknown>
514+
): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
515+
const deferred = new Deferred<
516+
AggregateQuerySnapshot<{ count: AggregateField<number> }>
517+
>();
518+
client.asyncQueue.enqueueAndForget(async () => {
519+
try {
520+
const remoteStore = await getRemoteStore(client);
521+
if (!canUseNetwork(remoteStore)) {
522+
deferred.reject(
523+
new FirestoreError(
524+
Code.UNAVAILABLE,
525+
'Failed to get count result because the client is offline.'
526+
)
527+
);
528+
} else {
529+
const result = await getCount(query);
530+
deferred.resolve(result);
531+
}
532+
} catch (e) {
533+
deferred.reject(e as Error);
534+
}
535+
});
536+
return deferred.promise;
537+
}
538+
504539
async function readDocumentFromCache(
505540
localStore: LocalStore,
506541
docKey: DocumentKey,
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
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 { deepEqual } from '@firebase/util';
19+
20+
import { Value } from '../protos/firestore_proto_api';
21+
import { invokeRunAggregationQueryRpc } from '../remote/datastore';
22+
import { hardAssert } from '../util/assert';
23+
import { cast } from '../util/input_validation';
24+
25+
import { getDatastore } from './components';
26+
import { Firestore } from './database';
27+
import { Query, queryEqual } from './reference';
28+
import { LiteUserDataWriter } from './reference_impl';
29+
30+
/**
31+
* An `AggregateField`that captures input type T.
32+
*/
33+
// eslint-disable-next-line @typescript-eslint/no-unused-vars
34+
export class AggregateField<T> {
35+
type = 'AggregateField';
36+
}
37+
38+
/**
39+
* Creates and returns an aggregation field that counts the documents in the result set.
40+
* @returns An `AggregateField` object with number input type.
41+
*/
42+
export function count(): AggregateField<number> {
43+
return new AggregateField<number>();
44+
}
45+
46+
/**
47+
* The union of all `AggregateField` types that are returned from the factory
48+
* functions.
49+
*/
50+
export type AggregateFieldType = ReturnType<typeof count>;
51+
52+
/**
53+
* A type whose values are all `AggregateField` objects.
54+
* This is used as an argument to the "getter" functions, and the snapshot will
55+
* map the same names to the corresponding values.
56+
*/
57+
export interface AggregateSpec {
58+
[field: string]: AggregateFieldType;
59+
}
60+
61+
/**
62+
* A type whose keys are taken from an `AggregateSpec` type, and whose values
63+
* are the result of the aggregation performed by the corresponding
64+
* `AggregateField` from the input `AggregateSpec`.
65+
*/
66+
export type AggregateSpecData<T extends AggregateSpec> = {
67+
[P in keyof T]: T[P] extends AggregateField<infer U> ? U : never;
68+
};
69+
70+
/**
71+
* An `AggregateQuerySnapshot` contains the results of running an aggregate query.
72+
*/
73+
export class AggregateQuerySnapshot<T extends AggregateSpec> {
74+
readonly type = 'AggregateQuerySnapshot';
75+
76+
/** @hideconstructor */
77+
constructor(
78+
readonly query: Query<unknown>,
79+
private readonly _data: AggregateSpecData<T>
80+
) {}
81+
82+
/**
83+
* The results of the requested aggregations. The keys of the returned object
84+
* will be the same as those of the `AggregateSpec` object specified to the
85+
* aggregation method, and the values will be the corresponding aggregation
86+
* result.
87+
*
88+
* @returns The aggregation statistics result of running a query.
89+
*/
90+
data(): AggregateSpecData<T> {
91+
return this._data;
92+
}
93+
}
94+
95+
/**
96+
* Counts the number of documents in the result set of the given query, ignoring
97+
* any locally-cached data and any locally-pending writes and simply surfacing
98+
* whatever the server returns. If the server cannot be reached then the
99+
* returned promise will be rejected.
100+
*
101+
* @param query - The `Query` to execute.
102+
*
103+
* @returns An `AggregateQuerySnapshot` that contains the number of documents.
104+
*/
105+
export function getCount(
106+
query: Query<unknown>
107+
): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
108+
const firestore = cast(query.firestore, Firestore);
109+
const datastore = getDatastore(firestore);
110+
const userDataWriter = new LiteUserDataWriter(firestore);
111+
return invokeRunAggregationQueryRpc(datastore, query._query).then(result => {
112+
hardAssert(
113+
result[0] !== undefined,
114+
'Aggregation fields are missing from result.'
115+
);
116+
117+
const counts = Object.entries(result[0])
118+
.filter(([key, value]) => key === 'count_alias')
119+
.map(([key, value]) => userDataWriter.convertValue(value as Value));
120+
121+
const countValue = counts[0];
122+
123+
hardAssert(
124+
typeof countValue === 'number',
125+
'Count aggregate field value is not a number: ' + countValue
126+
);
127+
128+
return Promise.resolve(
129+
new AggregateQuerySnapshot<{ count: AggregateField<number> }>(query, {
130+
count: countValue
131+
})
132+
);
133+
});
134+
}
135+
136+
/**
137+
* Compares two `AggregateQuerySnapshot` instances for equality.
138+
* Two `AggregateQuerySnapshot` instances are considered "equal" if they have
139+
* the same underlying query, and the same data.
140+
*
141+
* @param left - The `AggregateQuerySnapshot` to compare.
142+
* @param right - The `AggregateQuerySnapshot` to compare.
143+
*
144+
* @returns true if the AggregateQuerySnapshots are equal.
145+
*/
146+
export function aggregateQuerySnapshotEqual<T extends AggregateSpec>(
147+
left: AggregateQuerySnapshot<T>,
148+
right: AggregateQuerySnapshot<T>
149+
): boolean {
150+
return (
151+
queryEqual(left.query, right.query) && deepEqual(left.data(), right.data())
152+
);
153+
}

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

+6
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,12 @@ 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+
// Both `invokeRPC()` and `invokeStreamingRPC()` ignore their `path` arguments, and expect
90+
// the "path" to be part of the given `request`.
91+
return true;
92+
}
93+
8894
constructor(protos: grpc.GrpcObject, private databaseInfo: DatabaseInfo) {
8995
// eslint-disable-next-line @typescript-eslint/no-explicit-any
9096
this.firestore = (protos as any)['google']['firestore']['v1'];

packages/firestore/src/protos/firestore_proto_api.ts

+30
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,32 @@ export declare namespace firestoreV1ApiClientInterfaces {
334334
readTime?: string;
335335
skippedResults?: number;
336336
}
337+
interface RunAggregationQueryRequest {
338+
parent?: string;
339+
structuredAggregationQuery?: StructuredAggregationQuery;
340+
transaction?: string;
341+
newTransaction?: TransactionOptions;
342+
readTime?: string;
343+
}
344+
interface RunAggregationQueryResponse {
345+
result?: AggregationResult;
346+
transaction?: string;
347+
readTime?: string;
348+
}
349+
interface AggregationResult {
350+
aggregateFields?: ApiClientObjectMap<Value>;
351+
}
352+
interface StructuredAggregationQuery {
353+
structuredQuery?: StructuredQuery;
354+
aggregations?: Aggregation[];
355+
}
356+
interface Aggregation {
357+
count?: Count;
358+
alias?: string;
359+
}
360+
interface Count {
361+
upTo?: number;
362+
}
337363
interface Status {
338364
code?: number;
339365
message?: string;
@@ -479,6 +505,10 @@ export declare type RunQueryRequest =
479505
firestoreV1ApiClientInterfaces.RunQueryRequest;
480506
export declare type RunQueryResponse =
481507
firestoreV1ApiClientInterfaces.RunQueryResponse;
508+
export declare type RunAggregationQueryRequest =
509+
firestoreV1ApiClientInterfaces.RunAggregationQueryRequest;
510+
export declare type RunAggregationQueryResponse =
511+
firestoreV1ApiClientInterfaces.RunAggregationQueryResponse;
482512
export declare type Status = firestoreV1ApiClientInterfaces.Status;
483513
export declare type StructuredQuery =
484514
firestoreV1ApiClientInterfaces.StructuredQuery;
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
// Copyright 2022 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
syntax = "proto3";
16+
17+
package google.firestore.v1;
18+
19+
import "google/firestore/v1/document.proto";
20+
21+
option csharp_namespace = "Google.Cloud.Firestore.V1";
22+
option go_package = "google.golang.org/genproto/googleapis/firestore/v1;firestore";
23+
option java_multiple_files = true;
24+
option java_outer_classname = "AggregationResultProto";
25+
option java_package = "com.google.firestore.v1";
26+
option objc_class_prefix = "GCFS";
27+
option php_namespace = "Google\\Cloud\\Firestore\\V1";
28+
option ruby_package = "Google::Cloud::Firestore::V1";
29+
30+
// The result of a single bucket from a Firestore aggregation query.
31+
//
32+
// The keys of `aggregate_fields` are the same for all results in an aggregation
33+
// query, unlike document queries which can have different fields present for
34+
// each result.
35+
message AggregationResult {
36+
// The result of the aggregation functions, ex: `COUNT(*) AS total_docs`.
37+
//
38+
// The key is the [alias][google.firestore.v1.StructuredAggregationQuery.Aggregation.alias]
39+
// assigned to the aggregation function on input and the size of this map
40+
// equals the number of aggregation functions in the query.
41+
map<string, Value> aggregate_fields = 2;
42+
}

0 commit comments

Comments
 (0)