fix: include limit in query for findFirst and findTop support in Fire…

…store (#3387) * fix: include limit in query for findFirst and findTop in Firestore
GoogleCloudPlatform · Nov 20, 2024 · 9fe8e91 · 9fe8e91
1 parent a12638c
commit 9fe8e91
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 0 deletions.
diff --git a/.../java/com/google/cloud/spring/data/firestore/repository/query/PartTreeFirestoreQuery.java b/.../java/com/google/cloud/spring/data/firestore/repository/query/PartTreeFirestoreQuery.java
@@ -156,6 +156,11 @@ public Object execute(Object[] parameters) {
       sort = paramAccessor.getSort();
     }
 
+    if (getQueryMethod().getName().startsWith("findFirst")
+        || getQueryMethod().getName().startsWith("findTop")) {
+      builder.setLimit(Int32Value.newBuilder().setValue(1));
+    }
+
     if (sort == null || sort.isUnsorted()) {
       sort = this.tree.getSort();
     }

diff --git a/...restore/src/test/java/com/google/cloud/spring/data/firestore/entities/UserRepository.java b/...restore/src/test/java/com/google/cloud/spring/data/firestore/entities/UserRepository.java
@@ -28,6 +28,10 @@
 public interface UserRepository extends FirestoreReactiveRepository<User> {
   Flux<User> findBy(Pageable pageable);
 
+  Mono<User> findFirstByAge(Integer age);
+
+  Mono<User> findTopByAge(Integer age);
+
   Flux<User> findByAge(Integer age);
 
   Flux<User> findByAge(Integer age, Sort sort);

diff --git a/...t/java/com/google/cloud/spring/data/firestore/it/FirestoreRepositoryIntegrationTests.java b/...t/java/com/google/cloud/spring/data/firestore/it/FirestoreRepositoryIntegrationTests.java
@@ -343,4 +343,34 @@ void testUpdateTimeNoDocumentChangeDoesNotResultInOptimisticLockingFailure() {
     userRepository.save(user).block();
     // no optimistic locking exception expected
   }
+
+  @Test
+  public void testFirstByAge() {
+    User alice = new User("Alice", 99);
+    User bob = new User("Bob", 99);
+    User zelda = new User("Zelda", 23);
+    this.userRepository
+        .save(alice)
+        .then(this.userRepository.save(bob))
+        .then(this.userRepository.save(zelda))
+        .block();
+
+    Mono<User> testUser = this.userRepository.findFirstByAge(99);
+    assertThat(testUser.block().getName()).isEqualTo("Alice");
+  }
+
+  @Test
+  public void testTopByAge() {
+    User alice = new User("Alice", 99);
+    User bob = new User("Bob", 99);
+    User zelda = new User("Zelda", 23);
+    this.userRepository
+            .save(alice)
+            .then(this.userRepository.save(bob))
+            .then(this.userRepository.save(zelda))
+            .block();
+
+    Mono<User> testUser = this.userRepository.findTopByAge(99);
+    assertThat(testUser.block().getName()).isEqualTo("Alice");
+  }
 }
diff --git a/.../com/google/cloud/spring/data/firestore/repository/query/PartTreeFirestoreQueryTests.java b/.../com/google/cloud/spring/data/firestore/repository/query/PartTreeFirestoreQueryTests.java
@@ -31,6 +31,7 @@
 import com.google.cloud.spring.data.firestore.mapping.FirestoreMappingContext;
 import com.google.firestore.v1.StructuredQuery;
 import com.google.firestore.v1.Value;
+import com.google.protobuf.Int32Value;
 import java.util.function.Consumer;
 import org.junit.jupiter.api.Test;
 import org.mockito.invocation.InvocationOnMock;
@@ -223,6 +224,66 @@ void testPartTreeQueryOrException() {
         .hasMessage("Cloud Firestore doesn't support 'OR' (method name: findByAgeOrName)");
   }
 
+  @Test
+  void testPartTreeQuery_findFirst_limitIsSet() {
+    PartTreeFirestoreQuery partTreeFirestoreQuery =
+        createPartTreeQuery(
+            "findFirstByAge",
+            invocation -> {
+              StructuredQuery.Builder actualBuilder = invocation.getArgument(0);
+              Class clazz = invocation.getArgument(1);
+
+              StructuredQuery.Builder expectedBuilder = StructuredQuery.newBuilder();
+              StructuredQuery.CompositeFilter.Builder compositeFilter =
+                  StructuredQuery.CompositeFilter.newBuilder();
+              compositeFilter.setOp(StructuredQuery.CompositeFilter.Operator.AND);
+              StructuredQuery.Filter.Builder filterAge = StructuredQuery.Filter.newBuilder();
+              filterAge
+                  .getFieldFilterBuilder()
+                  .setField(StructuredQuery.FieldReference.newBuilder().setFieldPath("age").build())
+                  .setOp(StructuredQuery.FieldFilter.Operator.EQUAL)
+                  .setValue(this.classMapper.toFirestoreValue(22));
+              compositeFilter.addFilters(filterAge.build());
+              expectedBuilder.setWhere(
+                  StructuredQuery.Filter.newBuilder().setCompositeFilter(compositeFilter.build()));
+              expectedBuilder.setLimit(Int32Value.of(1));
+              assertThat(actualBuilder.build()).isEqualTo(expectedBuilder.build());
+              assertThat(clazz).isEqualTo(User.class);
+            });
+
+    partTreeFirestoreQuery.execute(new Object[] {22});
+  }
+
+  @Test
+  void testPartTreeQuery_findTop_limitIsSet() {
+    PartTreeFirestoreQuery partTreeFirestoreQuery =
+        createPartTreeQuery(
+            "findTopByAge",
+            invocation -> {
+              StructuredQuery.Builder actualBuilder = invocation.getArgument(0);
+              Class clazz = invocation.getArgument(1);
+
+              StructuredQuery.Builder expectedBuilder = StructuredQuery.newBuilder();
+              StructuredQuery.CompositeFilter.Builder compositeFilter =
+                  StructuredQuery.CompositeFilter.newBuilder();
+              compositeFilter.setOp(StructuredQuery.CompositeFilter.Operator.AND);
+              StructuredQuery.Filter.Builder filterAge = StructuredQuery.Filter.newBuilder();
+              filterAge
+                  .getFieldFilterBuilder()
+                  .setField(StructuredQuery.FieldReference.newBuilder().setFieldPath("age").build())
+                  .setOp(StructuredQuery.FieldFilter.Operator.EQUAL)
+                  .setValue(this.classMapper.toFirestoreValue(22));
+              compositeFilter.addFilters(filterAge.build());
+              expectedBuilder.setWhere(
+                  StructuredQuery.Filter.newBuilder().setCompositeFilter(compositeFilter.build()));
+              expectedBuilder.setLimit(Int32Value.of(1));
+              assertThat(actualBuilder.build()).isEqualTo(expectedBuilder.build());
+              assertThat(clazz).isEqualTo(User.class);
+            });
+
+    partTreeFirestoreQuery.execute(new Object[] {22});
+  }
+
   private PartTreeFirestoreQuery createPartTreeQuery(String methodName) {
     return createPartTreeQuery(methodName, NOOP);
   }