feat(firestore): global option to turn on implicit orderby

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

@@ -66,6 +66,7 @@ public final class FirestoreOptions extends ServiceOptions<Firestore, FirestoreO
   private final TransportChannelProvider channelProvider;
   private final CredentialsProvider credentialsProvider;
   private final String emulatorHost;
+  private final boolean alwaysUseImplicitOrderBy;
   private final transient @Nonnull FirestoreOpenTelemetryOptions openTelemetryOptions;
   private final transient @Nonnull com.google.cloud.firestore.telemetry.TraceUtil traceUtil;
   private final transient @Nonnull com.google.cloud.firestore.telemetry.MetricsUtil metricsUtil;
@@ -144,6 +145,10 @@ public String getEmulatorHost() {
     return emulatorHost;
   }
 
+  public boolean isAlwaysUseImplicitOrderBy() {
+    return alwaysUseImplicitOrderBy;
+  }
+
   @Nonnull
   com.google.cloud.firestore.telemetry.TraceUtil getTraceUtil() {
     return traceUtil;
@@ -166,6 +171,7 @@ public static class Builder extends ServiceOptions.Builder<Firestore, FirestoreO
     @Nullable private TransportChannelProvider channelProvider = null;
     @Nullable private CredentialsProvider credentialsProvider = null;
     @Nullable private String emulatorHost = null;
+    private boolean alwaysUseImplicitOrderBy = false;
     @Nullable private FirestoreOpenTelemetryOptions openTelemetryOptions = null;
 
     private Builder() {}
@@ -176,6 +182,7 @@ private Builder(FirestoreOptions options) {
       this.channelProvider = options.channelProvider;
       this.credentialsProvider = options.credentialsProvider;
       this.emulatorHost = options.emulatorHost;
+      this.alwaysUseImplicitOrderBy = options.alwaysUseImplicitOrderBy;
       this.openTelemetryOptions = options.openTelemetryOptions;
     }
 
@@ -235,6 +242,22 @@ public Builder setEmulatorHost(@Nonnull String emulatorHost) {
       return this;
     }
 
+    /**
+     * Sets whether to always use implicit order bys (e.g., for inequality queries).
+     *
+     * <p>By default, the SDK only sends explicit order bys and relies on the backend to append
+     * implicit order bys (unless cursors are involved). However, enterprise edition does not
+     * enforce implicit order bys because it does not back all queries with an index. This option
+     * allows enterprise users to opt-in to always forcing the SDK to append the implicit order by
+     * clauses to the request to mimic the standard edition behavior.
+     *
+     * @param alwaysUseImplicitOrderBy Whether to always use implicit order bys.
+     */
+    public Builder setAlwaysUseImplicitOrderBy(boolean alwaysUseImplicitOrderBy) {
+      this.alwaysUseImplicitOrderBy = alwaysUseImplicitOrderBy;
+      return this;
+    }
+
     /**
      * Sets the database ID to use with this Firestore client.
      *
@@ -380,6 +403,7 @@ protected FirestoreOptions(Builder builder) {
             : GrpcTransportOptions.setUpCredentialsProvider(this);
 
     this.emulatorHost = builder.emulatorHost;
+    this.alwaysUseImplicitOrderBy = builder.alwaysUseImplicitOrderBy;
   }
 
   private static class FirestoreDefaults implements ServiceDefaults<Firestore, FirestoreOptions> {

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

@@ -1392,7 +1392,17 @@ public Query endAt(@Nonnull DocumentSnapshot snapshot) {
 
   /** Build the final Firestore query. */
   StructuredQuery.Builder buildQuery() {
-    StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();
+    return buildQuery(/* forceImplicitOrderBy= */ false);
+  }
+
+  /**
+   * Generates the StructuredQuery for this Query, with an option to force implicit order bys.
+   *
+   * @param forceImplicitOrderBy Whether to force the inclusion of implicit order by clauses.
+   * @return The StructuredQuery model object.
+   */
+  StructuredQuery.Builder buildQuery(boolean forceImplicitOrderBy) {
+    StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation(forceImplicitOrderBy);
     if (options.getLimitType().equals(LimitType.Last)) {
       structuredQuery.clearOrderBy();
       structuredQuery.clearStartAt();
@@ -1441,7 +1451,8 @@ StructuredQuery.Builder buildQuery() {
    * representation via {@link BundledQuery.LimitType}.
    */
   BundledQuery toBundledQuery() {
-    StructuredQuery.Builder structuredQuery = buildWithoutClientTranslation();
+    StructuredQuery.Builder structuredQuery =
+        buildWithoutClientTranslation(/* forceImplicitOrderBy= */ false);
 
     return BundledQuery.newBuilder()
         .setStructuredQuery(structuredQuery)
@@ -1453,7 +1464,7 @@ BundledQuery toBundledQuery() {
         .build();
   }
 
-  private StructuredQuery.Builder buildWithoutClientTranslation() {
+  private StructuredQuery.Builder buildWithoutClientTranslation(boolean forceImplicitOrderBy) {
     StructuredQuery.Builder structuredQuery = StructuredQuery.newBuilder();
     CollectionSelector.Builder collectionSelector = CollectionSelector.newBuilder();
 
@@ -1472,8 +1483,14 @@ private StructuredQuery.Builder buildWithoutClientTranslation() {
       structuredQuery.setWhere(filter.toProto());
     }
 
-    if (!options.getFieldOrders().isEmpty()) {
-      for (FieldOrder order : options.getFieldOrders()) {
+    List<FieldOrder> ordersToSerialize = options.getFieldOrders();
+    if (forceImplicitOrderBy
+        || rpcContext.getFirestore().getOptions().isAlwaysUseImplicitOrderBy()) {
+      ordersToSerialize = createImplicitOrderBy();
+    }
+
+    if (!ordersToSerialize.isEmpty()) {
+      for (FieldOrder order : ordersToSerialize) {
         structuredQuery.addOrderBy(order.toProto());
       }
     } else if (LimitType.Last.equals(options.getLimitType())) {

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

@@ -162,7 +162,7 @@ static Watch forQuery(Query query) {
     Target.Builder target = Target.newBuilder();
     target.setQuery(
         QueryTarget.newBuilder()
-            .setStructuredQuery(query.buildQuery())
+            .setStructuredQuery(query.buildQuery(/* forceImplicitOrderBy= */ true))
             .setParent(query.options.getParentPath().getName())
             .build());
     target.setTargetId(WATCH_TARGET_ID);

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

@@ -22,21 +22,44 @@
 import static com.google.common.truth.Truth.assertThat;
 import static java.util.Arrays.asList;
 import static java.util.Collections.singletonList;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 
 import com.google.cloud.firestore.spi.v1.FirestoreRpc;
+import com.google.firestore.v1.RunAggregationQueryRequest;
+import com.google.firestore.v1.StructuredQuery;
 import java.util.List;
 import java.util.Objects;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mockito;
+import org.mockito.Spy;
 import org.mockito.junit.MockitoJUnitRunner;
 
 @RunWith(MockitoJUnitRunner.class)
 public class AggregateQueryTest {
 
-  @Mock private Query mockQuery;
-  @Mock private Query mockQuery2;
+  private final FirestoreRpc firestoreRpc = Mockito.mock(FirestoreRpc.class);
+
+  @Spy
+  private final FirestoreImpl firestoreMock =
+      new FirestoreImpl(
+          FirestoreOptions.newBuilder().setProjectId("test-project").build(), firestoreRpc);
+
+  @Captor private ArgumentCaptor<RunAggregationQueryRequest> runQuery;
+
+  private Query mockQuery;
+  private Query mockQuery2;
+
+  @org.junit.Before
+  public void before() {
+    mockQuery = firestoreMock.collection("coll");
+    mockQuery2 = firestoreMock.collection("coll2");
+  }
 
   @Test
   public void getQueryShouldReturnTheQuerySpecifiedToTheConstructor() {
@@ -130,4 +153,45 @@ public void toProtoFromProtoRoundTripShouldProduceEqualAggregateQueryObjects() {
     assertThat(countQuery2).isEqualTo(countQuery2Recreated);
     assertThat(countQuery1).isNotEqualTo(countQuery2);
   }
+
+  @Test
+  public void withAlwaysUseImplicitOrderBy() throws Exception {
+    doAnswer(
+            invocation -> {
+              com.google.api.gax.rpc.ResponseObserver<
+                      com.google.firestore.v1.RunAggregationQueryResponse>
+                  observer = invocation.getArgument(1);
+              observer.onResponse(
+                  com.google.firestore.v1.RunAggregationQueryResponse.newBuilder()
+                      .setResult(
+                          com.google.firestore.v1.AggregationResult.newBuilder()
+                              .putAggregateFields(
+                                  "aggregate_0",
+                                  com.google.firestore.v1.Value.newBuilder()
+                                      .setIntegerValue(1)
+                                      .build()))
+                      .build());
+              observer.onComplete();
+              return null;
+            })
+        .when(firestoreMock)
+        .streamRequest(runQuery.capture(), any(), any());
+
+    doReturn(
+            FirestoreOptions.newBuilder()
+                .setProjectId("test-project")
+                .setAlwaysUseImplicitOrderBy(true)
+                .build())
+        .when(firestoreMock)
+        .getOptions();
+
+    mockQuery.whereGreaterThan("a", "b").count().get().get();
+
+    RunAggregationQueryRequest queryRequest = runQuery.getValue();
+    StructuredQuery query = queryRequest.getStructuredAggregationQuery().getStructuredQuery();
+
+    assertThat(query.getOrderByCount()).isEqualTo(2);
+    assertThat(query.getOrderBy(0).getField().getFieldPath()).isEqualTo("a");
+    assertThat(query.getOrderBy(1).getField().getFieldPath()).isEqualTo("__name__");
+  }
 }

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

@@ -599,6 +599,32 @@ public void withDocumentIdAndDocumentSnapshotCursor() {
     assertEquals(queryRequest, runQuery.getValue());
   }
 
+  @Test
+  public void withAlwaysUseImplicitOrderBy() throws Exception {
+    doAnswer(queryResponse())
+        .when(firestoreMock)
+        .streamRequest(runQuery.capture(), streamObserverCapture.capture(), any());
+
+    doReturn(
+            FirestoreOptions.newBuilder()
+                .setProjectId("test-project")
+                .setAlwaysUseImplicitOrderBy(true)
+                .build())
+        .when(firestoreMock)
+        .getOptions();
+
+    query.whereEqualTo("a", "b").whereGreaterThanOrEqualTo("foo", "bar").get().get();
+
+    RunQueryRequest queryRequest =
+        query(
+            filter(Operator.EQUAL, "a", "b"),
+            filter(Operator.GREATER_THAN_OR_EQUAL, "foo", "bar"),
+            order("foo", Direction.ASCENDING),
+            order("__name__", StructuredQuery.Direction.ASCENDING));
+
+    assertEquals(queryRequest, runQuery.getValue());
+  }
+
   @Test
   public void withDocumentReferenceCursor() {
     doAnswer(queryResponse())

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

@@ -451,6 +451,27 @@ public void queryWatchHandlesDocumentChange() throws InterruptedException {
         new SnapshotDocument(ChangeType.REMOVED, "coll/doc3", null));
   }
 
+  @Test
+  public void queryWatchWithImplicitOrderBy() throws InterruptedException {
+    listenerRegistration =
+        firestoreMock
+            .collection("coll")
+            .whereGreaterThan("foo", "bar")
+            .addSnapshotListener((value, error) -> querySnapshots.add(value));
+
+    ListenRequest listenRequest = requests.take();
+    assertEquals(DATABASE_NAME, listenRequest.getDatabase());
+    assertEquals(TARGET_ID, listenRequest.getAddTarget().getTargetId());
+
+    // Verify the query includes the implicit order bys
+    com.google.firestore.v1.StructuredQuery query =
+        listenRequest.getAddTarget().getQuery().getStructuredQuery();
+
+    assertEquals(2, query.getOrderByCount());
+    assertEquals("foo", query.getOrderBy(0).getField().getFieldPath());
+    assertEquals("__name__", query.getOrderBy(1).getField().getFieldPath());
+  }
+
   @Test
   public void queryWatchReconnectsWithResumeToken() throws InterruptedException {
     addQueryListener();

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

@@ -1235,4 +1235,54 @@ public void snapshotListenerSortsInvalidUnicodeStringsSameWayAsServer() throws E
         break;
     }
   }
+
+  @Test
+  public void alwaysUseImplicitOrderByReturnsSameResults() throws Exception {
+    CollectionReference collection =
+        testCollectionWithDocs(
+            map(
+                "doc01", map("sort", 1),
+                "doc02", map("sort", 2),
+                "doc03", map("sort", 3),
+                "doc04", map("sort", 4),
+                "doc05", map("sort", 5),
+                "doc06", map("sort", 6),
+                "doc07", map("sort", 7),
+                "doc08", map("sort", 8),
+                "doc09", map("sort", 9),
+                "doc10", map("sort", 10)));
+
+    List<String> expectedOrder =
+        Arrays.asList(
+            "doc02", "doc03", "doc04", "doc05", "doc06", "doc07", "doc08", "doc09", "doc10");
+
+    Query originalQuery = firestore.collection(collection.getId()).whereGreaterThan("sort", 1);
+    QuerySnapshot originalSnapshot = originalQuery.get().get();
+    List<String> originalResult =
+        originalSnapshot.getDocuments().stream()
+            .map(queryDocumentSnapshot -> queryDocumentSnapshot.getReference().getId())
+            .collect(Collectors.toList());
+
+    if (getFirestoreEdition() == FirestoreEdition.ENTERPRISE) {
+      assertThat(originalResult).containsExactlyElementsIn(expectedOrder);
+      assertThat(originalResult).isNotEqualTo(expectedOrder);
+    } else {
+      assertThat(originalResult).isEqualTo(expectedOrder);
+    }
+
+    FirestoreOptions modifiedOptions =
+        firestore.getOptions().toBuilder().setAlwaysUseImplicitOrderBy(true).build();
+    try (Firestore modifiedFirestore = modifiedOptions.getService()) {
+      Query query = modifiedFirestore.collection(collection.getId()).whereGreaterThan("sort", 1);
+
+      QuerySnapshot snapshot = query.get().get();
+      List<String> result =
+          snapshot.getDocuments().stream()
+              .map(queryDocumentSnapshot -> queryDocumentSnapshot.getReference().getId())
+              .collect(Collectors.toList());
+
+      // since alwaysUseImplicitOrderBy is true, we expect strict ordering even for ENTERPRISE
+      assertThat(result).isEqualTo(expectedOrder);
+    }
+  }
 }