Firestore find by (#1836)

findById and findAllById in FirestoreTemplate ; fixes #1778

Author: dmitry-s

spring-cloud-gcp-data-firestore/src/main/java/org/springframework/cloud/gcp/data/firestore/FirestoreTemplate.java

@@ -24,6 +24,7 @@
 import com.google.firestore.v1.DeleteDocumentRequest;
 import com.google.firestore.v1.Document;
 import com.google.firestore.v1.FirestoreGrpc.FirestoreStub;
+import com.google.firestore.v1.GetDocumentRequest;
 import com.google.firestore.v1.RunQueryRequest;
 import com.google.firestore.v1.RunQueryResponse;
 import com.google.firestore.v1.StructuredQuery;
@@ -74,6 +75,20 @@ public FirestoreTemplate(FirestoreStub firestore, String parent) {
 		this.databasePath = parent.substring(0, StringUtils.ordinalIndexOf(parent, "/", 4));
 	}
 
+	public <T> Mono<T> findById(Publisher idPublisher, Class<T> aClass) {
+		return findAllById(idPublisher, aClass).next();
+	}
+
+	public <T> Flux<T> findAllById(Publisher idPublisher, Class<T> aClass) {
+		return ((Flux<String>) Flux.from(idPublisher)).flatMap(id -> {
+			FirestorePersistentEntity<?> persistentEntity = this.mappingContext.getPersistentEntity(aClass);
+			GetDocumentRequest request = GetDocumentRequest.newBuilder()
+					.setName(this.parent + "/" + persistentEntity.collectionName() + "/" + id).build();
+			return ObservableReactiveUtil.<Document>unaryCall(obs -> this.firestore.getDocument(request, obs));
+		}).onErrorMap(throwable -> new FirestoreDataException("Unable to find an entry by id", throwable))
+				.map(document -> PublicClassMapper.convertToCustomClass(document, aClass));
+	}
+
 	public <T> Mono<T> save(T entity) {
 		return Mono.defer(() -> {
 			FirestorePersistentEntity<?> persistentEntity = this.mappingContext.getPersistentEntity(entity.getClass());

spring-cloud-gcp-data-firestore/src/test/java/org/springframework/cloud/gcp/data/firestore/FirestoreIntegrationTests.java

@@ -31,6 +31,7 @@
 import org.junit.Test;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.data.annotation.Id;
 
@@ -85,6 +86,10 @@ public void writeReadDeleteTest() {
 		this.firestoreTemplate.save(alice).block();
 		this.firestoreTemplate.save(bob).block();
 
+		assertThat(this.firestoreTemplate.findById(Mono.just("Bob"), User.class).block()).isEqualTo(bob);
+		assertThat(this.firestoreTemplate.findAllById(Flux.just("Bob", "Alice"), User.class).collectList().block())
+				.containsExactly(bob, alice);
+
 		List<User> usersBeforeDelete = this.firestoreTemplate.findAll(User.class).collectList().block();
 
 		assertThat(this.firestoreTemplate.deleteAll(User.class).block()).isEqualTo(2);
@@ -104,6 +109,7 @@ public void saveAllTest() throws InterruptedException {
 		this.firestoreTemplate.saveAll(users).collectList().block();
 
 		assertThat(this.firestoreTemplate.findAll(User.class).collectList().block().size()).isEqualTo(2);
+		assertThat(this.firestoreTemplate.deleteAll(User.class).block()).isEqualTo(2);
 	}
 }
 

spring-cloud-gcp-data-firestore/src/test/java/org/springframework/cloud/gcp/data/firestore/FirestoreTemplateTests.java

@@ -24,6 +24,7 @@
 import com.google.firestore.v1.DeleteDocumentRequest;
 import com.google.firestore.v1.Document;
 import com.google.firestore.v1.FirestoreGrpc.FirestoreStub;
+import com.google.firestore.v1.GetDocumentRequest;
 import com.google.firestore.v1.RunQueryRequest;
 import com.google.firestore.v1.RunQueryResponse;
 import com.google.firestore.v1.StructuredQuery;
@@ -32,6 +33,8 @@
 import io.grpc.stub.StreamObserver;
 import org.junit.Before;
 import org.junit.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import org.springframework.data.annotation.Id;
@@ -104,6 +107,82 @@ public void findAllTest() {
 		verify(this.firestoreStub, times(1)).runQuery(any(), any());
 	}
 
+	@Test
+	public void findByIdTest() {
+		doAnswer(invocation -> {
+			StreamObserver<Document> streamObserver = invocation.getArgument(1);
+			streamObserver.onNext(buildDocument("e1", 100L));
+
+			streamObserver.onCompleted();
+			return null;
+		}).when(this.firestoreStub).getDocument(any(), any());
+
+		StepVerifier.create(this.firestoreTemplate.findById(Mono.just("e1"), TestEntity.class))
+				.expectNext(new TestEntity("e1", 100L))
+				.verifyComplete();
+
+		GetDocumentRequest request = GetDocumentRequest.newBuilder()
+				.setName(this.parent + "/testEntities/" + "e1")
+				.build();
+
+		verify(this.firestoreStub, times(1)).getDocument(eq(request), any());
+	}
+
+	@Test
+	public void findByIdErrorTest() {
+		doAnswer(invocation -> {
+			StreamObserver<Document> streamObserver = invocation.getArgument(1);
+			streamObserver.onError(new RuntimeException("Firestore error"));
+			return null;
+		}).when(this.firestoreStub).getDocument(any(), any());
+
+		StepVerifier.create(this.firestoreTemplate.findById(Mono.just("e1"), TestEntity.class))
+				.expectErrorMatches(e ->
+						e instanceof FirestoreDataException
+						&& e.getMessage().contains("Firestore error")
+						&& e.getMessage().contains("Unable to find an entry by id"))
+				.verify();
+
+		GetDocumentRequest request = GetDocumentRequest.newBuilder()
+				.setName(this.parent + "/testEntities/" + "e1")
+				.build();
+
+		verify(this.firestoreStub, times(1)).getDocument(eq(request), any());
+	}
+
+	@Test
+	public void findAllByIdTest() {
+		GetDocumentRequest request1 = GetDocumentRequest.newBuilder()
+				.setName(this.parent + "/testEntities/" + "e1")
+				.build();
+
+		GetDocumentRequest request2 = GetDocumentRequest.newBuilder()
+				.setName(this.parent + "/testEntities/" + "e2")
+				.build();
+
+		doAnswer(invocation -> {
+			StreamObserver<Document> streamObserver = invocation.getArgument(1);
+			streamObserver.onNext(buildDocument("e1", 100L));
+
+			streamObserver.onCompleted();
+			return null;
+		}).when(this.firestoreStub).getDocument(eq(request1), any());
+
+		doAnswer(invocation -> {
+			StreamObserver<Document> streamObserver = invocation.getArgument(1);
+			streamObserver.onNext(buildDocument("e2", 200L));
+
+			streamObserver.onCompleted();
+			return null;
+		}).when(this.firestoreStub).getDocument(eq(request2), any());
+
+		StepVerifier.create(this.firestoreTemplate.findAllById(Flux.just("e1", "e2"), TestEntity.class))
+				.expectNext(new TestEntity("e1", 100L), new TestEntity("e2", 200L))
+				.verifyComplete();
+
+		verify(this.firestoreStub, times(1)).getDocument(eq(request1), any());
+	}
+
 	@Test
 	public void deleteAllTest() {
 		mockRunQueryMethod();