Thanks to visit codestin.com
Credit goes to github.com

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

feat: Add long alias support for aggregations. #1267

Merged
ehsannas merged 6 commits into from
Aug 2, 2023
Merged

feat: Add long alias support for aggregations. #1267

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5dcc233
feat: Add long alias support for aggregations.
ehsannas Apr 24, 2023
fce4364
address comments.
ehsannas Jul 27, 2023
279a35f
Better method name and replace hardcoded "count" with "aggregate_0".
ehsannas Jul 28, 2023
daeebf9
Remove duplicate aggregations.
ehsannas Jul 28, 2023
a5febd3
add static import.
ehsannas Jul 28, 2023
ebff5cb
Fix tests. All tests pass.
ehsannas Jul 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 22 additions & 12 deletions google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuery.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,21 +36,16 @@
/** A query that calculates aggregations over an underlying query. */
@InternalExtensionOnly
public class AggregateQuery {

/**
* The "alias" to specify in the {@link RunAggregationQueryRequest} proto when running a count
* query. The actual value is not meaningful, but will be used to get the count out of the {@link
* RunAggregationQueryResponse}.
*/
private static final String ALIAS_COUNT = "count";

@Nonnull private final Query query;

@Nonnull private List<AggregateField> aggregateFieldList;

@Nonnull private Map<String, String> aliasMap;

AggregateQuery(@Nonnull Query query, @Nonnull List<AggregateField> aggregateFields) {
this.query = query;
this.aggregateFieldList = aggregateFields;
this.aliasMap = new HashMap<>();
}

/** Returns the query whose aggregations will be calculated by this object. */
Expand Down Expand Up @@ -114,7 +109,9 @@ long getStartTimeNanos() {

void deliverResult(@Nonnull Map<String, Value> data, Timestamp readTime) {
if (isFutureCompleted.compareAndSet(false, true)) {
future.set(new AggregateQuerySnapshot(AggregateQuery.this, readTime, data));
Map<String, Value> mappedData = new HashMap<>();
data.forEach((serverAlias, value) -> mappedData.put(aliasMap.get(serverAlias), value));
future.set(new AggregateQuerySnapshot(AggregateQuery.this, readTime, mappedData));
}
}

Expand Down Expand Up @@ -202,9 +199,18 @@ RunAggregationQueryRequest toProto(@Nullable final ByteString transactionId) {
request.getStructuredAggregationQueryBuilder();
structuredAggregationQuery.setStructuredQuery(runQueryRequest.getStructuredQuery());

// We use a Set here to automatically remove duplicates.
Set<StructuredAggregationQuery.Aggregation> aggregations = new HashSet<>();
// We use this set to remove duplicate aggregates. e.g. `aggregate(sum("foo"), sum("foo"))`
HashSet<String> uniqueAggregates = new HashSet<>();
List<StructuredAggregationQuery.Aggregation> aggregations = new ArrayList<>();
int aggregationNum = 0;
for (AggregateField aggregateField : aggregateFieldList) {
// `getAlias()` provides a unique representation of an AggregateField.
boolean isNewAggregateField = uniqueAggregates.add(aggregateField.getAlias());
if (!isNewAggregateField) {
// This is a duplicate AggregateField. We don't need to include it in the request.
continue;
}

// If there's a field for this aggregation, build its proto.
StructuredQuery.FieldReference field = null;
if (!aggregateField.getFieldPath().isEmpty()) {
Expand All @@ -224,7 +230,11 @@ RunAggregationQueryRequest toProto(@Nullable final ByteString transactionId) {
} else {
throw new RuntimeException("Unsupported aggregation");
}
aggregation.setAlias(aggregateField.getAlias());
// Map all client-side aliases to a unique short-form alias.
// This avoids issues with client-side aliases that exceed the 1500-byte string size limit.
String serverAlias = "aggregate_" + aggregationNum++;
aliasMap.put(serverAlias, aggregateField.getAlias());
aggregation.setAlias(serverAlias);
aggregations.add(aggregation.build());

@cherylEnkidu cherylEnkidu Apr 28, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will the deduplication still work?

@ehsannas ehsannas Jul 27, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My understanding is:
duplicate aliases are not allowed
duplicate aggregates are allowed

is that right @cherylEnkidu @MarkDuckworth ?

this code uses aggregate_<counter> as alias, so there should not be any duplicate aliases.

also, there's a test named canGetDuplicateAggregations below which currently passes against the emulator.

@ehsannas ehsannas Jul 28, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah... it's coming back to me now... I believe we decided to remove duplicate aggregates even though they are allowed to improve performance.

@ehsannas ehsannas Jul 28, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

Followed what was done in Android. PTAL.

@MarkDuckworth MarkDuckworth Jul 28, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

}
structuredAggregationQuery.addAllAggregations(aggregations);
Expand Down
35 changes: 17 additions & 18 deletions google-cloud-firestore/src/test/java/com/google/cloud/firestore/LocalFirestoreHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,33 +341,32 @@ public static Answer<RunQueryResponse> queryResponseWithDone(
}
}

public static Answer<RunAggregationQueryResponse> aggregationQueryResponse() {
return aggregationQueryResponse(42);
public static Answer<RunAggregationQueryResponse> countQueryResponse() {
return countQueryResponse(42);
}

public static Answer<RunAggregationQueryResponse> aggregationQueryResponse(int count) {
return aggregationQueryResponse(count, null);
public static Answer<RunAggregationQueryResponse> countQueryResponse(int count) {
return countQueryResponse(count, null);
}

public static Answer<RunAggregationQueryResponse> aggregationQueryResponse(
public static Answer<RunAggregationQueryResponse> countQueryResponse(
int count, @Nullable Timestamp readTime) {
return streamingResponse(
new RunAggregationQueryResponse[] {
createRunAggregationQueryResponse(count, readTime),
createCountQueryResponse(count, readTime),
},
/*throwable=*/ null);
}

public static Answer<RunAggregationQueryResponse> aggregationQueryResponse(Throwable throwable) {
public static Answer<RunAggregationQueryResponse> countQueryResponse(Throwable throwable) {
return streamingResponse(new RunAggregationQueryResponse[] {}, throwable);
}

public static Answer<RunAggregationQueryResponse> aggregationQueryResponses(
int count1, int count2) {
return streamingResponse(
new RunAggregationQueryResponse[] {
createRunAggregationQueryResponse(count1, null),
createRunAggregationQueryResponse(count2, null),
createCountQueryResponse(count1, null), createCountQueryResponse(count2, null),
},
/*throwable=*/ null);
}
Expand All @@ -376,17 +375,17 @@ public static Answer<RunAggregationQueryResponse> aggregationQueryResponses(
int count1, Throwable throwable) {
return streamingResponse(
new RunAggregationQueryResponse[] {
createRunAggregationQueryResponse(count1, null),
createCountQueryResponse(count1, null),
},
throwable);
}

private static RunAggregationQueryResponse createRunAggregationQueryResponse(
private static RunAggregationQueryResponse createCountQueryResponse(
int count, @Nullable Timestamp timestamp) {
RunAggregationQueryResponse.Builder builder = RunAggregationQueryResponse.newBuilder();
builder.setResult(
AggregationResult.newBuilder()
.putAggregateFields("count", Value.newBuilder().setIntegerValue(count).build())
.putAggregateFields("aggregate_0", Value.newBuilder().setIntegerValue(count).build())
.build());
if (timestamp != null) {
builder.setReadTime(timestamp.toProto());
Expand Down Expand Up @@ -751,11 +750,11 @@ public static RunQueryRequest query(
return request.build();
}

public static RunAggregationQueryRequest aggregationQuery() {
return aggregationQuery((String) null);
public static RunAggregationQueryRequest countQuery() {
return countQuery((String) null);
}

public static RunAggregationQueryRequest aggregationQuery(@Nullable String transactionId) {
public static RunAggregationQueryRequest countQuery(@Nullable String transactionId) {
RunQueryRequest runQueryRequest = query(TRANSACTION_ID, false);

RunAggregationQueryRequest.Builder request =
Expand All @@ -766,7 +765,7 @@ public static RunAggregationQueryRequest aggregationQuery(@Nullable String trans
.setStructuredQuery(runQueryRequest.getStructuredQuery())
.addAggregations(
Aggregation.newBuilder()
.setAlias("count")
.setAlias("aggregate_0")
.setCount(Aggregation.Count.getDefaultInstance())));

if (transactionId != null) {
Expand All @@ -776,15 +775,15 @@ public static RunAggregationQueryRequest aggregationQuery(@Nullable String trans
return request.build();
}

public static RunAggregationQueryRequest aggregationQuery(RunQueryRequest runQueryRequest) {
public static RunAggregationQueryRequest countQuery(RunQueryRequest runQueryRequest) {
return RunAggregationQueryRequest.newBuilder()
.setParent(runQueryRequest.getParent())
.setStructuredAggregationQuery(
StructuredAggregationQuery.newBuilder()
.setStructuredQuery(runQueryRequest.getStructuredQuery())
.addAggregations(
Aggregation.newBuilder()
.setAlias("count")
.setAlias("aggregate_0")
.setCount(Aggregation.Count.getDefaultInstance())))
.build();
}
Expand Down
45 changes: 22 additions & 23 deletions google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryCountTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,8 @@
package com.google.cloud.firestore;

import static com.google.cloud.firestore.LocalFirestoreHelper.COLLECTION_ID;
import static com.google.cloud.firestore.LocalFirestoreHelper.aggregationQuery;
import static com.google.cloud.firestore.LocalFirestoreHelper.aggregationQueryResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.aggregationQueryResponses;
import static com.google.cloud.firestore.LocalFirestoreHelper.countQueryResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.endAt;
import static com.google.cloud.firestore.LocalFirestoreHelper.limit;
import static com.google.cloud.firestore.LocalFirestoreHelper.order;
Expand Down Expand Up @@ -78,7 +77,7 @@ public void before() {

@Test
public void countShouldBeZeroForEmptyCollection() throws Exception {
doAnswer(aggregationQueryResponse(0))
doAnswer(countQueryResponse(0))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -92,7 +91,7 @@ public void countShouldBeZeroForEmptyCollection() throws Exception {

@Test
public void countShouldBe99ForCollectionWith99Documents() throws Exception {
doAnswer(aggregationQueryResponse(99))
doAnswer(countQueryResponse(99))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -106,7 +105,7 @@ public void countShouldBe99ForCollectionWith99Documents() throws Exception {

@Test
public void countShouldMakeCorrectRequestForACollection() throws Exception {
doAnswer(aggregationQueryResponse(0))
doAnswer(countQueryResponse(0))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -116,12 +115,12 @@ public void countShouldMakeCorrectRequestForACollection() throws Exception {
CollectionReference collection = firestoreMock.collection(COLLECTION_ID);
collection.count().get();

assertThat(runAggregationQuery.getValue()).isEqualTo(aggregationQuery());
assertThat(runAggregationQuery.getValue()).isEqualTo(LocalFirestoreHelper.countQuery());
}

@Test
public void countShouldMakeCorrectRequestForAComplexQuery() throws Exception {
doAnswer(aggregationQueryResponse(0))
doAnswer(countQueryResponse(0))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -132,7 +131,7 @@ public void countShouldMakeCorrectRequestForAComplexQuery() throws Exception {

assertThat(runAggregationQuery.getValue())
.isEqualTo(
aggregationQuery(
LocalFirestoreHelper.countQuery(
query(
limit(42),
order("foo", StructuredQuery.Direction.DESCENDING),
Expand All @@ -142,7 +141,7 @@ public void countShouldMakeCorrectRequestForAComplexQuery() throws Exception {

@Test
public void shouldReturnReadTimeFromResponse() throws Exception {
doAnswer(aggregationQueryResponse(99, Timestamp.ofTimeSecondsAndNanos(123, 456)))
doAnswer(countQueryResponse(99, Timestamp.ofTimeSecondsAndNanos(123, 456)))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand Down Expand Up @@ -185,7 +184,7 @@ public void shouldIgnoreExtraErrors() throws Exception {
@Test
public void shouldPropagateErrors() throws Exception {
Exception exception = new Exception();
doAnswer(aggregationQueryResponse(exception))
doAnswer(countQueryResponse(exception))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -205,7 +204,7 @@ public void aggregateQueryGetQueryShouldReturnCorrectValue() throws Exception {

@Test
public void aggregateQuerySnapshotGetQueryShouldReturnCorrectValue() throws Exception {
doAnswer(aggregationQueryResponse())
doAnswer(countQueryResponse())
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -220,8 +219,8 @@ public void aggregateQuerySnapshotGetQueryShouldReturnCorrectValue() throws Exce

@Test
public void shouldNotRetryIfExceptionIsNotFirestoreException() {
doAnswer(aggregationQueryResponse(new NotFirestoreException()))
.doAnswer(aggregationQueryResponse())
doAnswer(countQueryResponse(new NotFirestoreException()))
.doAnswer(countQueryResponse())
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -235,8 +234,8 @@ public void shouldNotRetryIfExceptionIsNotFirestoreException() {

@Test
public void shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatus() throws Exception {
doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(aggregationQueryResponse(42))
doAnswer(countQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(countQueryResponse(42))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -252,8 +251,8 @@ public void shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatus() thro
@Test
public void shouldNotRetryIfExceptionIsFirestoreExceptionWithNonRetryableStatus() {
doReturn(Duration.ZERO).when(firestoreMock).getTotalRequestTimeout();
doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INVALID_ARGUMENT)))
.doAnswer(aggregationQueryResponse())
doAnswer(countQueryResponse(new FirestoreException("reason", Status.INVALID_ARGUMENT)))
.doAnswer(countQueryResponse())
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -270,8 +269,8 @@ public void shouldNotRetryIfExceptionIsFirestoreExceptionWithNonRetryableStatus(
shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatusWithInfiniteTimeoutWindow()
throws Exception {
doReturn(Duration.ZERO).when(firestoreMock).getTotalRequestTimeout();
doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(aggregationQueryResponse(42))
doAnswer(countQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(countQueryResponse(42))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -288,8 +287,8 @@ public void shouldNotRetryIfExceptionIsFirestoreExceptionWithNonRetryableStatus(
public void shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatusWithinTimeoutWindow()
throws Exception {
doReturn(Duration.ofDays(999)).when(firestoreMock).getTotalRequestTimeout();
doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(aggregationQueryResponse(42))
doAnswer(countQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(countQueryResponse(42))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand All @@ -313,8 +312,8 @@ public void shouldRetryIfExceptionIsFirestoreExceptionWithRetryableStatusWithinT
.when(clockMock)
.nanoTime();
doReturn(Duration.ofSeconds(5)).when(firestoreMock).getTotalRequestTimeout();
doAnswer(aggregationQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(aggregationQueryResponse(42))
doAnswer(countQueryResponse(new FirestoreException("reason", Status.INTERNAL)))
.doAnswer(countQueryResponse(42))
.when(firestoreMock)
.streamRequest(
runAggregationQuery.capture(),
Expand Down
25 changes: 3 additions & 22 deletions google-cloud-firestore/src/test/java/com/google/cloud/firestore/TransactionTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,26 +16,7 @@

package com.google.cloud.firestore;

import static com.google.cloud.firestore.LocalFirestoreHelper.IMMEDIATE_RETRY_SETTINGS;
import static com.google.cloud.firestore.LocalFirestoreHelper.SINGLE_FIELD_PROTO;
import static com.google.cloud.firestore.LocalFirestoreHelper.TRANSACTION_ID;
import static com.google.cloud.firestore.LocalFirestoreHelper.aggregationQuery;
import static com.google.cloud.firestore.LocalFirestoreHelper.aggregationQueryResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.begin;
import static com.google.cloud.firestore.LocalFirestoreHelper.beginResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.commit;
import static com.google.cloud.firestore.LocalFirestoreHelper.commitResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.create;
import static com.google.cloud.firestore.LocalFirestoreHelper.delete;
import static com.google.cloud.firestore.LocalFirestoreHelper.get;
import static com.google.cloud.firestore.LocalFirestoreHelper.getAll;
import static com.google.cloud.firestore.LocalFirestoreHelper.getAllResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.query;
import static com.google.cloud.firestore.LocalFirestoreHelper.queryResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.rollback;
import static com.google.cloud.firestore.LocalFirestoreHelper.rollbackResponse;
import static com.google.cloud.firestore.LocalFirestoreHelper.set;
import static com.google.cloud.firestore.LocalFirestoreHelper.update;
import static com.google.cloud.firestore.LocalFirestoreHelper.*;
import static com.google.common.truth.Truth.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;
Expand Down Expand Up @@ -678,7 +659,7 @@ public void getAggregateQuery() throws Exception {
.when(firestoreMock)
.sendRequest(requestCapture.capture(), Matchers.<UnaryCallable<Message, Message>>any());

doAnswer(aggregationQueryResponse(42))
doAnswer(countQueryResponse(42))
.when(firestoreMock)
.streamRequest(requestCapture.capture(), streamObserverCapture.capture(), Matchers.any());

Expand All @@ -691,7 +672,7 @@ public void getAggregateQuery() throws Exception {
assertEquals(3, requests.size());

assertEquals(begin(), requests.get(0));
assertEquals(aggregationQuery(TRANSACTION_ID), requests.get(1));
assertEquals(countQuery(TRANSACTION_ID), requests.get(1));
assertEquals(commit(TRANSACTION_ID), requests.get(2));
}

Expand Down
Loading