Skip to content
This repository was archived by the owner on Feb 24, 2026. It is now read-only.

Commit 8c22e61

Browse files
jainsahabjainsahab
authored
feat: Count API (#823)
* Add method in Datastore client to invoke rpc for aggregation query * Creating count aggregation and using it to populate Aggregation proto * Moving aggregation builder method to root level aggregation class * Introducing RecordQuery to represent queries which returns entity records when executed * Updating gitignore with patch extension * Setting up structure of Aggregation query and its builder * Introducing ProtoPreparer to populate the request protos * Delegating responsibility of preparing query proto to QueryPreparer * Populating aggregation query with nested structured query * Delegating responsibility of preparing query proto in GqlQuery to QueryPreparer * Removing RecordQuery from the query hierarchy and making it a standalone interface for now * Populating aggregation query with nested gql query * Removing deprecation warning by using assertThrows instead of ExpectedException rule * Making DatastoreRpc call aggregation query method on client * Creating response transformer to transform aggregation query response into domain objects * Implementing aggregation query executor to execute AggergationQuery * Adding missing assertion statements * Creating RetryExecutor to inject it as a dependency in other components * Making RetryExecutor accept RetrySettings when creating it * Revert "Making RetryExecutor accept RetrySettings when creating it" This reverts commit 1dfafb7. * Revert "Creating RetryExecutor to inject it as a dependency in other components" This reverts commit 8872a55. * Introducing RetryAndTraceDatastoreRpcDecorator to have retry and traceability logic on top of another DatastoreRpc * Extracting out the responsibility of preparing ReadOption in it's own ProtoPreparer * Making QueryExecutor to execute query with provided ReadOptions * Exposing readTime to the user * Ignoring runAggregationQuery method from clirr check * Making readTime final * Allowing namespace to be optional in AggregationQuery * Add capability to fetch aggregation result by passing alias * Implementing User facing datastore.runAggrgation method to run aggregation query * Add integration test for count aggregation * Add transaction Id support in ReadOptionsProtoPreparer * Supporting aggregation query with transactions * Allowing user to create Aggregation directly without involving its builder * Preventing creating duplicated aggregation when creating an aggregation query * Marking RecordQuery implemented method as InternalApi * Writing comments and JavaDoc for aggregation query related class * Adding a default implementation in the public interfaces to avoid compile time failures * covering a scenario to maintain consistent snapshot when executing aggregation query in a transaction * Creating emulator proxy to simulate AggregationQuery response from emulator * Integration test to execute an aggregation query in a read only transaction * Getting rid off limit operation on count aggregation as same behaviour can be achieved by using 'limit' operation on the underlying query * Removing import statement from javadoc and undo changes in .gitignore file * Using Optional instead of returning null from ReadOptionsProtoPreparer * using assertThat from Truth library * fixing unit test * Getting rid off Double braces initialization syntax * Fixing lint * Getting rid off emulator proxy and using easy mock to check the aggregationQuery triggered * Deleting a entity created locally in other test which is causing failure in other test * Deleting all keys in datastore in integration test so that new test can start fresh * Executing two read write transaction simultaneously and verifying their behaviour * Removing tests to verify serializability as it's an underlying implementation detail * Fixing lint * Adding runAggregationQuery method to reflect config so that it's available and accessible in native image through reflection * Fixing equals of CountAggregation * Fixing lint * Adding an integration test of using limit option with aggregation query * Adding BetaApi annotation to public surface to indicate that aggregation query / count is in preview * Fixing lint * Removing unused functiona and fixing javadoc * fixing variable name
1 parent 4641306 commit 8c22e61

38 files changed

Lines changed: 2410 additions & 12 deletions

‎datastore-v1-proto-client/src/main/java/com/google/datastore/v1/client/Datastore.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
import com.google.datastore.v1.ReserveIdsResponse;
2828
import com.google.datastore.v1.RollbackRequest;
2929
import com.google.datastore.v1.RollbackResponse;
30+
import com.google.datastore.v1.RunAggregationQueryRequest;
31+
import com.google.datastore.v1.RunAggregationQueryResponse;
3032
import com.google.datastore.v1.RunQueryRequest;
3133
import com.google.datastore.v1.RunQueryResponse;
3234
import com.google.rpc.Code;
@@ -120,4 +122,13 @@ public RunQueryResponse runQuery(RunQueryRequest request) throws DatastoreExcept
120122
throw invalidResponseException("runQuery", exception);
121123
}
122124
}
125+
126+
public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request)
127+
throws DatastoreException {
128+
try (InputStream is = remoteRpc.call("runAggregationQuery", request)) {
129+
return RunAggregationQueryResponse.parseFrom(is);
130+
} catch (IOException exception) {
131+
throw invalidResponseException("runAggregationQuery", exception);
132+
}
133+
}
123134
}

‎datastore-v1-proto-client/src/main/resources/META-INF/native-image/reflect-config.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
{"name":"lookup","parameterTypes":["com.google.datastore.v1.LookupRequest"] },
99
{"name":"reserveIds","parameterTypes":["com.google.datastore.v1.ReserveIdsRequest"] },
1010
{"name":"rollback","parameterTypes":["com.google.datastore.v1.RollbackRequest"] },
11-
{"name":"runQuery","parameterTypes":["com.google.datastore.v1.RunQueryRequest"] }
11+
{"name":"runQuery","parameterTypes":["com.google.datastore.v1.RunQueryRequest"] },
12+
{"name":"runAggregationQuery","parameterTypes":["com.google.datastore.v1.RunAggregationQueryRequest"] }
1213
]
1314
},
1415
{

‎datastore-v1-proto-client/src/test/java/com/google/datastore/v1/client/DatastoreClientTest.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@
3838
import com.google.datastore.v1.ReserveIdsResponse;
3939
import com.google.datastore.v1.RollbackRequest;
4040
import com.google.datastore.v1.RollbackResponse;
41+
import com.google.datastore.v1.RunAggregationQueryRequest;
42+
import com.google.datastore.v1.RunAggregationQueryResponse;
4143
import com.google.datastore.v1.RunQueryRequest;
4244
import com.google.datastore.v1.RunQueryResponse;
4345
import com.google.datastore.v1.client.testing.MockCredential;
@@ -336,6 +338,13 @@ public void runQuery() throws Exception {
336338
expectRpc("runQuery", request.build(), response.build());
337339
}
338340

341+
@Test
342+
public void runAggregationQuery() throws Exception {
343+
RunAggregationQueryRequest.Builder request = RunAggregationQueryRequest.newBuilder();
344+
RunAggregationQueryResponse.Builder response = RunAggregationQueryResponse.newBuilder();
345+
expectRpc("runAggregationQuery", request.build(), response.build());
346+
}
347+
339348
private void expectRpc(String methodName, Message request, Message response) throws Exception {
340349
Datastore datastore = factory.create(options.build());
341350
MockDatastoreFactory mockClient = (MockDatastoreFactory) factory;

‎google-cloud-datastore/clirr-ignored-differences.xml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,19 @@
1111
<method>com.google.datastore.v1.ReserveIdsResponse reserveIds(com.google.datastore.v1.ReserveIdsRequest)</method>
1212
<differenceType>7012</differenceType>
1313
</difference>
14+
<difference>
15+
<className>com/google/cloud/datastore/spi/v1/DatastoreRpc</className>
16+
<method>com.google.datastore.v1.RunAggregationQueryResponse runAggregationQuery(com.google.datastore.v1.RunAggregationQueryRequest)</method>
17+
<differenceType>7012</differenceType>
18+
</difference>
19+
<difference>
20+
<className>com/google/cloud/datastore/Datastore</className>
21+
<method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery, com.google.cloud.datastore.ReadOption[])</method>
22+
<differenceType>7012</differenceType>
23+
</difference>
24+
<difference>
25+
<className>com/google/cloud/datastore/DatastoreReader</className>
26+
<method>com.google.cloud.datastore.AggregationResults runAggregation(com.google.cloud.datastore.AggregationQuery)</method>
27+
<differenceType>7012</differenceType>
28+
</difference>
1429
</differences>

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationQuery.java‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
/*
2+
* Copyright 2022 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.datastore;
17+
18+
import static com.google.cloud.datastore.AggregationQuery.Mode.GQL;
19+
import static com.google.cloud.datastore.AggregationQuery.Mode.STRUCTURED;
20+
import static com.google.common.base.Preconditions.checkArgument;
21+
22+
import com.google.api.core.BetaApi;
23+
import com.google.cloud.datastore.aggregation.Aggregation;
24+
import com.google.cloud.datastore.aggregation.AggregationBuilder;
25+
import java.util.HashSet;
26+
import java.util.Set;
27+
28+
/**
29+
* An implementation of a Google Cloud Datastore Query that returns {@link AggregationResults}, It
30+
* can be constructed by providing a nested query ({@link StructuredQuery} or {@link GqlQuery}) to
31+
* run the aggregations on and a set of {@link Aggregation}.
32+
*
33+
* <p>{@link StructuredQuery} example:
34+
*
35+
* <pre>{@code
36+
* EntityQuery selectAllQuery = Query.newEntityQueryBuilder()
37+
* .setKind("Task")
38+
* .build();
39+
* AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
40+
* .addAggregation(count().as("total_count"))
41+
* .over(selectAllQuery)
42+
* .build();
43+
* AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
44+
* for (AggregationResult aggregationResult : aggregationResults) {
45+
* System.out.println(aggregationResult.get("total_count"));
46+
* }
47+
* }</pre>
48+
*
49+
* <h4>{@link GqlQuery} example:</h4>
50+
*
51+
* <pre>{@code
52+
* GqlQuery<?> selectAllGqlQuery = Query.newGqlQueryBuilder(
53+
* "AGGREGATE COUNT(*) AS total_count, COUNT_UP_TO(100) AS count_upto_100 OVER(SELECT * FROM Task)"
54+
* )
55+
* .setAllowLiteral(true)
56+
* .build();
57+
* AggregationQuery aggregationQuery = Query.newAggregationQueryBuilder()
58+
* .over(selectAllGqlQuery)
59+
* .build();
60+
* AggregationResults aggregationResults = datastore.runAggregation(aggregationQuery);
61+
* for (AggregationResult aggregationResult : aggregationResults) {
62+
* System.out.println(aggregationResult.get("total_count"));
63+
* System.out.println(aggregationResult.get("count_upto_100"));
64+
* }
65+
* }</pre>
66+
*
67+
* @see <a href="https://cloud.google.com/appengine/docs/java/datastore/queries">Datastore
68+
* queries</a>
69+
*/
70+
@BetaApi
71+
public class AggregationQuery extends Query<AggregationResults> {
72+
73+
private Set<Aggregation> aggregations;
74+
private StructuredQuery<?> nestedStructuredQuery;
75+
private final Mode mode;
76+
private GqlQuery<?> nestedGqlQuery;
77+
78+
AggregationQuery(
79+
String namespace, Set<Aggregation> aggregations, StructuredQuery<?> nestedQuery) {
80+
super(namespace);
81+
checkArgument(
82+
!aggregations.isEmpty(),
83+
"At least one aggregation is required for an aggregation query to run");
84+
this.aggregations = aggregations;
85+
this.nestedStructuredQuery = nestedQuery;
86+
this.mode = STRUCTURED;
87+
}
88+
89+
AggregationQuery(String namespace, GqlQuery<?> gqlQuery) {
90+
super(namespace);
91+
this.nestedGqlQuery = gqlQuery;
92+
this.mode = GQL;
93+
}
94+
95+
/** Returns the {@link Aggregation}(s) for this Query. */
96+
public Set<Aggregation> getAggregations() {
97+
return aggregations;
98+
}
99+
100+
/**
101+
* Returns the underlying {@link StructuredQuery for this Query}. Returns null if created with
102+
* {@link GqlQuery}
103+
*/
104+
public StructuredQuery<?> getNestedStructuredQuery() {
105+
return nestedStructuredQuery;
106+
}
107+
108+
/**
109+
* Returns the underlying {@link GqlQuery for this Query}. Returns null if created with {@link
110+
* StructuredQuery}
111+
*/
112+
public GqlQuery<?> getNestedGqlQuery() {
113+
return nestedGqlQuery;
114+
}
115+
116+
/** Returns the {@link Mode} for this query. */
117+
public Mode getMode() {
118+
return mode;
119+
}
120+
121+
public static class Builder {
122+
123+
private String namespace;
124+
private Mode mode;
125+
private final Set<Aggregation> aggregations;
126+
private StructuredQuery<?> nestedStructuredQuery;
127+
private GqlQuery<?> nestedGqlQuery;
128+
129+
public Builder() {
130+
this.aggregations = new HashSet<>();
131+
}
132+
133+
public Builder setNamespace(String namespace) {
134+
this.namespace = namespace;
135+
return this;
136+
}
137+
138+
public Builder addAggregation(AggregationBuilder<?> aggregationBuilder) {
139+
this.aggregations.add(aggregationBuilder.build());
140+
return this;
141+
}
142+
143+
public Builder addAggregation(Aggregation aggregation) {
144+
this.aggregations.add(aggregation);
145+
return this;
146+
}
147+
148+
public Builder over(StructuredQuery<?> nestedQuery) {
149+
this.nestedStructuredQuery = nestedQuery;
150+
this.mode = STRUCTURED;
151+
return this;
152+
}
153+
154+
public Builder over(GqlQuery<?> nestedQuery) {
155+
this.nestedGqlQuery = nestedQuery;
156+
this.mode = GQL;
157+
return this;
158+
}
159+
160+
public AggregationQuery build() {
161+
boolean nestedQueryProvided = nestedGqlQuery != null || nestedStructuredQuery != null;
162+
checkArgument(
163+
nestedQueryProvided, "Nested query is required for an aggregation query to run");
164+
165+
if (mode == GQL) {
166+
return new AggregationQuery(namespace, nestedGqlQuery);
167+
}
168+
return new AggregationQuery(namespace, aggregations, nestedStructuredQuery);
169+
}
170+
}
171+
172+
public enum Mode {
173+
STRUCTURED,
174+
GQL,
175+
}
176+
}

‎google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
/*
2+
* Copyright 2022 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.datastore;
17+
18+
import com.google.api.core.BetaApi;
19+
import com.google.common.base.MoreObjects;
20+
import com.google.common.base.MoreObjects.ToStringHelper;
21+
import java.util.Map;
22+
import java.util.Map.Entry;
23+
import java.util.Objects;
24+
25+
/** Represents a result of an {@link AggregationQuery} query submission. */
26+
@BetaApi
27+
public class AggregationResult {
28+
29+
private final Map<String, LongValue> properties;
30+
31+
public AggregationResult(Map<String, LongValue> properties) {
32+
this.properties = properties;
33+
}
34+
35+
/**
36+
* Returns a result value for the given alias.
37+
*
38+
* @param alias A custom alias provided in the query or an autogenerated alias in the form of
39+
* 'property_\d'
40+
* @return An aggregation result value for the given alias.
41+
*/
42+
public Long get(String alias) {
43+
return properties.get(alias).get();
44+
}
45+
46+
@Override
47+
public boolean equals(Object o) {
48+
if (this == o) {
49+
return true;
50+
}
51+
if (o == null || getClass() != o.getClass()) {
52+
return false;
53+
}
54+
AggregationResult that = (AggregationResult) o;
55+
return properties.equals(that.properties);
56+
}
57+
58+
@Override
59+
public int hashCode() {
60+
return Objects.hash(properties);
61+
}
62+
63+
@Override
64+
public String toString() {
65+
ToStringHelper toStringHelper = MoreObjects.toStringHelper(this);
66+
for (Entry<String, LongValue> entry : properties.entrySet()) {
67+
toStringHelper.add(entry.getKey(), entry.getValue().get());
68+
}
69+
return toStringHelper.toString();
70+
}
71+
}

0 commit comments

Comments
 (0)