Skip to content

Commit 24d3e17

Browse files
wu-huiwu-hui
authored
feat(firestore): global option to turn on implicit orderby (#2337)
1 parent 8c5dd1f commit 24d3e17

7 files changed

Lines changed: 250 additions & 11 deletions

File tree

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/FirestoreOptions.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ public final class FirestoreOptions extends ServiceOptions<Firestore, FirestoreO
6666
private final TransportChannelProvider channelProvider;
6767
private final CredentialsProvider credentialsProvider;
6868
private final String emulatorHost;
69+
private final boolean alwaysUseImplicitOrderBy;
6970
private final transient @Nonnull FirestoreOpenTelemetryOptions openTelemetryOptions;
7071
private final transient @Nonnull com.google.cloud.firestore.telemetry.TraceUtil traceUtil;
7172
private final transient @Nonnull com.google.cloud.firestore.telemetry.MetricsUtil metricsUtil;
@@ -144,6 +145,10 @@ public String getEmulatorHost() {
144145
return emulatorHost;
145146
}
146147

148+
public boolean isAlwaysUseImplicitOrderBy() {
149+
return alwaysUseImplicitOrderBy;
150+
}
151+
147152
@Nonnull
148153
com.google.cloud.firestore.telemetry.TraceUtil getTraceUtil() {
149154
return traceUtil;
@@ -166,6 +171,7 @@ public static class Builder extends ServiceOptions.Builder<Firestore, FirestoreO
166171
@Nullable private TransportChannelProvider channelProvider = null;
167172
@Nullable private CredentialsProvider credentialsProvider = null;
168173
@Nullable private String emulatorHost = null;
174+
private boolean alwaysUseImplicitOrderBy = false;
169175
@Nullable private FirestoreOpenTelemetryOptions openTelemetryOptions = null;
170176

171177
private Builder() {}
@@ -176,6 +182,7 @@ private Builder(FirestoreOptions options) {
176182
this.channelProvider = options.channelProvider;
177183
this.credentialsProvider = options.credentialsProvider;
178184
this.emulatorHost = options.emulatorHost;
185+
this.alwaysUseImplicitOrderBy = options.alwaysUseImplicitOrderBy;
179186
this.openTelemetryOptions = options.openTelemetryOptions;
180187
}
181188

@@ -235,6 +242,26 @@ public Builder setEmulatorHost(@Nonnull String emulatorHost) {
235242
return this;
236243
}
237244

245+
/**
246+
* Sets whether to always include implicit order by clauses in the query request (e.g., for
247+
* inequality queries).
248+
*
249+
* <p>By default, the SDK only sends explicit order by clauses and relies on the backend to
250+
* append implicit ones (unless cursors are used). Firestore Enterprise edition, however, does
251+
* not automatically append these clauses because it does not require an index for every query.
252+
* This option allows users to opt-in to having the SDK always append the implicit order by
253+
* clauses, ensuring behavior consistent with standard edition.
254+
*
255+
* <p>Setting this option to true against Standard Edition is essentially a no-op as Standard
256+
* Edition automatically apply implicit orderby from the backend.
257+
*
258+
* @param alwaysUseImplicitOrderBy Whether to always include implicit order by clauses.
259+
*/
260+
public Builder setAlwaysUseImplicitOrderBy(boolean alwaysUseImplicitOrderBy) {
261+
this.alwaysUseImplicitOrderBy = alwaysUseImplicitOrderBy;
262+
return this;
263+
}
264+
238265
/**
239266
* Sets the database ID to use with this Firestore client.
240267
*
@@ -380,6 +407,7 @@ protected FirestoreOptions(Builder builder) {
380407
: GrpcTransportOptions.setUpCredentialsProvider(this);
381408

382409
this.emulatorHost = builder.emulatorHost;
410+
this.alwaysUseImplicitOrderBy = builder.alwaysUseImplicitOrderBy;
383411
}
384412

385413
private static class FirestoreDefaults implements ServiceDefaults<Firestore, FirestoreOptions> {

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Query.java‎

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1392,15 +1392,31 @@ public Query endAt(@Nonnull DocumentSnapshot snapshot) {
13921392

13931393
/** Build the final Firestore query. */
13941394
StructuredQuery.Builder buildQuery() {
1395-
StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();
1395+
return buildQuery(/* forceImplicitOrderBy= */ false);
1396+
}
1397+
1398+
/**
1399+
* Generates the StructuredQuery for this Query, with an option to force implicit order bys.
1400+
*
1401+
* @param forceImplicitOrderBy Whether to force the inclusion of implicit order by clauses.
1402+
* @return The StructuredQuery model object.
1403+
*/
1404+
StructuredQuery.Builder buildQuery(boolean forceImplicitOrderBy) {
1405+
StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation(forceImplicitOrderBy);
13961406
if (options.getLimitType().equals(LimitType.Last)) {
13971407
structuredQuery.clearOrderBy();
13981408
structuredQuery.clearStartAt();
13991409
structuredQuery.clearEndAt();
14001410

14011411
// Apply client translation for limitToLast.
1402-
if (!options.getFieldOrders().isEmpty()) {
1403-
for (FieldOrder order : options.getFieldOrders()) {
1412+
List<FieldOrder> ordersToFlip = options.getFieldOrders();
1413+
if (forceImplicitOrderBy
1414+
|| rpcContext.getFirestore().getOptions().isAlwaysUseImplicitOrderBy()) {
1415+
ordersToFlip = createImplicitOrderBy();
1416+
}
1417+
1418+
if (!ordersToFlip.isEmpty()) {
1419+
for (FieldOrder order : ordersToFlip) {
14041420
// Flip the orderBy directions since we want the last results
14051421
order =
14061422
new FieldOrder(
@@ -1441,7 +1457,8 @@ StructuredQuery.Builder buildQuery() {
14411457
* representation via {@link BundledQuery.LimitType}.
14421458
*/
14431459
BundledQuery toBundledQuery() {
1444-
StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();
1460+
StructuredQuery.Builder structuredQuery =
1461+
buildWithoutClientTranslation(/* forceImplicitOrderBy= */ false);
14451462

14461463
return BundledQuery.newBuilder()
14471464
.setStructuredQuery(structuredQuery)
@@ -1453,7 +1470,7 @@ BundledQuery toBundledQuery() {
14531470
.build();
14541471
}
14551472

1456-
private StructuredQuery.Builder buildWithoutClientTranslation() {
1473+
private StructuredQuery.Builder buildWithoutClientTranslation(boolean forceImplicitOrderBy) {
14571474
StructuredQuery.Builder structuredQuery = StructuredQuery.newBuilder();
14581475
CollectionSelector.Builder collectionSelector = CollectionSelector.newBuilder();
14591476

@@ -1472,8 +1489,14 @@ private StructuredQuery.Builder buildWithoutClientTranslation() {
14721489
structuredQuery.setWhere(filter.toProto());
14731490
}
14741491

1475-
if (!options.getFieldOrders().isEmpty()) {
1476-
for (FieldOrder order : options.getFieldOrders()) {
1492+
List<FieldOrder> ordersToSerialize = options.getFieldOrders();
1493+
if (forceImplicitOrderBy
1494+
|| rpcContext.getFirestore().getOptions().isAlwaysUseImplicitOrderBy()) {
1495+
ordersToSerialize = createImplicitOrderBy();
1496+
}
1497+
1498+
if (!ordersToSerialize.isEmpty()) {
1499+
for (FieldOrder order : ordersToSerialize) {
14771500
structuredQuery.addOrderBy(order.toProto());
14781501
}
14791502
} else if (LimitType.Last.equals(options.getLimitType())) {

‎google-cloud-firestore/src/main/java/com/google/cloud/firestore/Watch.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@ static Watch forQuery(Query query) {
162162
Target.Builder target = Target.newBuilder();
163163
target.setQuery(
164164
QueryTarget.newBuilder()
165-
.setStructuredQuery(query.buildQuery())
165+
.setStructuredQuery(query.buildQuery(/* forceImplicitOrderBy= */ true))
166166
.setParent(query.options.getParentPath().getName())
167167
.build());
168168
target.setTargetId(WATCH_TARGET_ID);

‎google-cloud-firestore/src/test/java/com/google/cloud/firestore/AggregateQueryTest.java‎

Lines changed: 67 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,21 +22,44 @@
2222
import static com.google.common.truth.Truth.assertThat;
2323
import static java.util.Arrays.asList;
2424
import static java.util.Collections.singletonList;
25+
import static org.mockito.ArgumentMatchers.any;
26+
import static org.mockito.Mockito.doAnswer;
27+
import static org.mockito.Mockito.doReturn;
2528
import static org.mockito.Mockito.mock;
2629

2730
import com.google.cloud.firestore.spi.v1.FirestoreRpc;
31+
import com.google.firestore.v1.RunAggregationQueryRequest;
32+
import com.google.firestore.v1.StructuredQuery;
2833
import java.util.List;
2934
import java.util.Objects;
3035
import org.junit.Test;
3136
import org.junit.runner.RunWith;
32-
import org.mockito.Mock;
37+
import org.mockito.ArgumentCaptor;
38+
import org.mockito.Captor;
39+
import org.mockito.Mockito;
40+
import org.mockito.Spy;
3341
import org.mockito.junit.MockitoJUnitRunner;
3442

3543
@RunWith(MockitoJUnitRunner.class)
3644
public class AggregateQueryTest {
3745

38-
@Mock private Query mockQuery;
39-
@Mock private Query mockQuery2;
46+
private final FirestoreRpc firestoreRpc = Mockito.mock(FirestoreRpc.class);
47+
48+
@Spy
49+
private final FirestoreImpl firestoreMock =
50+
new FirestoreImpl(
51+
FirestoreOptions.newBuilder().setProjectId("test-project").build(), firestoreRpc);
52+
53+
@Captor private ArgumentCaptor<RunAggregationQueryRequest> runQuery;
54+
55+
private Query mockQuery;
56+
private Query mockQuery2;
57+
58+
@org.junit.Before
59+
public void before() {
60+
mockQuery = firestoreMock.collection("coll");
61+
mockQuery2 = firestoreMock.collection("coll2");
62+
}
4063

4164
@Test
4265
public void getQueryShouldReturnTheQuerySpecifiedToTheConstructor() {
@@ -130,4 +153,45 @@ public void toProtoFromProtoRoundTripShouldProduceEqualAggregateQueryObjects() {
130153
assertThat(countQuery2).isEqualTo(countQuery2Recreated);
131154
assertThat(countQuery1).isNotEqualTo(countQuery2);
132155
}
156+
157+
@Test
158+
public void withAlwaysUseImplicitOrderBy() throws Exception {
159+
doAnswer(
160+
invocation -> {
161+
com.google.api.gax.rpc.ResponseObserver<
162+
com.google.firestore.v1.RunAggregationQueryResponse>
163+
observer = invocation.getArgument(1);
164+
observer.onResponse(
165+
com.google.firestore.v1.RunAggregationQueryResponse.newBuilder()
166+
.setResult(
167+
com.google.firestore.v1.AggregationResult.newBuilder()
168+
.putAggregateFields(
169+
"aggregate_0",
170+
com.google.firestore.v1.Value.newBuilder()
171+
.setIntegerValue(1)
172+
.build()))
173+
.build());
174+
observer.onComplete();
175+
return null;
176+
})
177+
.when(firestoreMock)
178+
.streamRequest(runQuery.capture(), any(), any());
179+
180+
doReturn(
181+
FirestoreOptions.newBuilder()
182+
.setProjectId("test-project")
183+
.setAlwaysUseImplicitOrderBy(true)
184+
.build())
185+
.when(firestoreMock)
186+
.getOptions();
187+
188+
mockQuery.whereGreaterThan("a", "b").count().get().get();
189+
190+
RunAggregationQueryRequest queryRequest = runQuery.getValue();
191+
StructuredQuery query = queryRequest.getStructuredAggregationQuery().getStructuredQuery();
192+
193+
assertThat(query.getOrderByCount()).isEqualTo(2);
194+
assertThat(query.getOrderBy(0).getField().getFieldPath()).isEqualTo("a");
195+
assertThat(query.getOrderBy(1).getField().getFieldPath()).isEqualTo("__name__");
196+
}
133197
}

‎google-cloud-firestore/src/test/java/com/google/cloud/firestore/QueryTest.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,59 @@ public void withDocumentIdAndDocumentSnapshotCursor() {
599599
assertEquals(queryRequest, runQuery.getValue());
600600
}
601601

602+
@Test
603+
public void withAlwaysUseImplicitOrderBy() throws Exception {
604+
doAnswer(queryResponse())
605+
.when(firestoreMock)
606+
.streamRequest(runQuery.capture(), streamObserverCapture.capture(), any());
607+
608+
doReturn(
609+
FirestoreOptions.newBuilder()
610+
.setProjectId("test-project")
611+
.setAlwaysUseImplicitOrderBy(true)
612+
.build())
613+
.when(firestoreMock)
614+
.getOptions();
615+
616+
query.whereEqualTo("a", "b").whereGreaterThanOrEqualTo("foo", "bar").get().get();
617+
618+
RunQueryRequest queryRequest =
619+
query(
620+
filter(Operator.EQUAL, "a", "b"),
621+
filter(Operator.GREATER_THAN_OR_EQUAL, "foo", "bar"),
622+
order("foo", Direction.ASCENDING),
623+
order("__name__", StructuredQuery.Direction.ASCENDING));
624+
625+
assertEquals(queryRequest, runQuery.getValue());
626+
}
627+
628+
@Test
629+
public void withAlwaysUseImplicitOrderByAndLimitToLast() throws Exception {
630+
doAnswer(queryResponse())
631+
.when(firestoreMock)
632+
.streamRequest(runQuery.capture(), streamObserverCapture.capture(), any());
633+
634+
doReturn(
635+
FirestoreOptions.newBuilder()
636+
.setProjectId("test-project")
637+
.setAlwaysUseImplicitOrderBy(true)
638+
.build())
639+
.when(firestoreMock)
640+
.getOptions();
641+
642+
query.whereEqualTo("a", "b").whereGreaterThanOrEqualTo("foo", "bar").limitToLast(1).get().get();
643+
644+
RunQueryRequest queryRequest =
645+
query(
646+
filter(Operator.EQUAL, "a", "b"),
647+
filter(Operator.GREATER_THAN_OR_EQUAL, "foo", "bar"),
648+
order("foo", Direction.DESCENDING),
649+
order("__name__", Direction.DESCENDING),
650+
limit(1));
651+
652+
assertEquals(queryRequest, runQuery.getValue());
653+
}
654+
602655
@Test
603656
public void withDocumentReferenceCursor() {
604657
doAnswer(queryResponse())

‎google-cloud-firestore/src/test/java/com/google/cloud/firestore/WatchTest.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,27 @@ public void queryWatchHandlesDocumentChange() throws InterruptedException {
451451
new SnapshotDocument(ChangeType.REMOVED, "coll/doc3", null));
452452
}
453453

454+
@Test
455+
public void queryWatchWithImplicitOrderBy() throws InterruptedException {
456+
listenerRegistration =
457+
firestoreMock
458+
.collection("coll")
459+
.whereGreaterThan("foo", "bar")
460+
.addSnapshotListener((value, error) -> querySnapshots.add(value));
461+
462+
ListenRequest listenRequest = requests.take();
463+
assertEquals(DATABASE_NAME, listenRequest.getDatabase());
464+
assertEquals(TARGET_ID, listenRequest.getAddTarget().getTargetId());
465+
466+
// Verify the query includes the implicit order bys
467+
com.google.firestore.v1.StructuredQuery query =
468+
listenRequest.getAddTarget().getQuery().getStructuredQuery();
469+
470+
assertEquals(2, query.getOrderByCount());
471+
assertEquals("foo", query.getOrderBy(0).getField().getFieldPath());
472+
assertEquals("__name__", query.getOrderBy(1).getField().getFieldPath());
473+
}
474+
454475
@Test
455476
public void queryWatchReconnectsWithResumeToken() throws InterruptedException {
456477
addQueryListener();

‎google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITQueryTest.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1235,4 +1235,54 @@ public void snapshotListenerSortsInvalidUnicodeStringsSameWayAsServer() throws E
12351235
break;
12361236
}
12371237
}
1238+
1239+
@Test
1240+
public void alwaysUseImplicitOrderByReturnsSameResults() throws Exception {
1241+
CollectionReference collection =
1242+
testCollectionWithDocs(
1243+
map(
1244+
"doc01", map("sort", 1),
1245+
"doc02", map("sort", 2),
1246+
"doc03", map("sort", 3),
1247+
"doc04", map("sort", 4),
1248+
"doc05", map("sort", 5),
1249+
"doc06", map("sort", 6),
1250+
"doc07", map("sort", 7),
1251+
"doc08", map("sort", 8),
1252+
"doc09", map("sort", 9),
1253+
"doc10", map("sort", 10)));
1254+
1255+
List<String> expectedOrder =
1256+
Arrays.asList(
1257+
"doc02", "doc03", "doc04", "doc05", "doc06", "doc07", "doc08", "doc09", "doc10");
1258+
1259+
Query originalQuery = firestore.collection(collection.getId()).whereGreaterThan("sort", 1);
1260+
QuerySnapshot originalSnapshot = originalQuery.get().get();
1261+
List<String> originalResult =
1262+
originalSnapshot.getDocuments().stream()
1263+
.map(queryDocumentSnapshot -> queryDocumentSnapshot.getReference().getId())
1264+
.collect(Collectors.toList());
1265+
1266+
if (getFirestoreEdition() == FirestoreEdition.ENTERPRISE) {
1267+
assertThat(originalResult).containsExactlyElementsIn(expectedOrder);
1268+
assertThat(originalResult).isNotEqualTo(expectedOrder);
1269+
} else {
1270+
assertThat(originalResult).isEqualTo(expectedOrder);
1271+
}
1272+
1273+
FirestoreOptions modifiedOptions =
1274+
firestore.getOptions().toBuilder().setAlwaysUseImplicitOrderBy(true).build();
1275+
try (Firestore modifiedFirestore = modifiedOptions.getService()) {
1276+
Query query = modifiedFirestore.collection(collection.getId()).whereGreaterThan("sort", 1);
1277+
1278+
QuerySnapshot snapshot = query.get().get();
1279+
List<String> result =
1280+
snapshot.getDocuments().stream()
1281+
.map(queryDocumentSnapshot -> queryDocumentSnapshot.getReference().getId())
1282+
.collect(Collectors.toList());
1283+
1284+
// since alwaysUseImplicitOrderBy is true, we expect strict ordering even for ENTERPRISE
1285+
assertThat(result).isEqualTo(expectedOrder);
1286+
}
1287+
}
12381288
}

0 commit comments

Comments
 (0)