refactor saveAll flow for publisher

Author: christophstrobl

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/repository/support/SimpleReactiveMongoRepository.java

@@ -116,13 +116,11 @@ public <S extends T> Flux<S> saveAll(Iterable<S> entities) {
 	}
 
 	@Override
-	public <S extends T> Flux<S> saveAll(Publisher<S> entityStream) {
+	public <S extends T> Flux<S> saveAll(Publisher<S> publisher) {
 
-		Assert.notNull(entityStream, "The given Publisher of entities must not be null");
+		Assert.notNull(publisher, "The given Publisher of entities must not be null");
 
-		return Flux.from(entityStream).concatMap(entity -> entityInformation.isNew(entity) ? //
-				mongoOperations.insert(entity, entityInformation.getCollectionName()) : //
-				mongoOperations.save(entity, entityInformation.getCollectionName()));
+		return concatMapSequentially(publisher, this::save);
 	}
 
 	@Override
@@ -498,6 +496,20 @@ static <T> Flux<T> concatMapSequentially(List<T> source,
 		return first.concatWith(theRest);
 	}
 
+	static <T> Flux<T> concatMapSequentially(Publisher<T> publisher,
+			Function<? super T, ? extends Publisher<? extends T>> mapper) {
+
+		return Flux.from(publisher).switchOnFirst(((signal, source) -> {
+
+			if (!signal.hasValue()) {
+				return source.concatMap(mapper);
+			}
+
+			Mono<T> firstCall = Mono.from(mapper.apply(signal.get()));
+			return firstCall.concatWith(source.skip(1).flatMapSequential(mapper));
+		}));
+	}
+
 	private static <E> List<E> toList(Iterable<E> source) {
 		return source instanceof List<E> list ? list : new ArrayList<>(toCollection(source));
 	}

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/SimpleReactiveMongoRepositoryTests.java

@@ -18,18 +18,16 @@
 import static org.assertj.core.api.Assertions.*;
 import static org.springframework.data.domain.ExampleMatcher.*;
 
-import org.junit.jupiter.api.RepeatedTest;
-import org.springframework.data.mongodb.ReactiveMongoTransactionManager;
-import org.springframework.transaction.TransactionDefinition;
-import org.springframework.transaction.reactive.TransactionalOperator;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import java.util.Arrays;
 import java.util.Objects;
+import java.util.stream.Stream;
 
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.springframework.beans.BeansException;
@@ -44,6 +42,7 @@
 import org.springframework.data.domain.Sort;
 import org.springframework.data.domain.Sort.Direction;
 import org.springframework.data.domain.Sort.Order;
+import org.springframework.data.mongodb.ReactiveMongoTransactionManager;
 import org.springframework.data.mongodb.core.ReactiveMongoTemplate;
 import org.springframework.data.mongodb.repository.support.ReactiveMongoRepositoryFactory;
 import org.springframework.data.mongodb.repository.support.SimpleReactiveMongoRepository;
@@ -52,6 +51,8 @@
 import org.springframework.lang.Nullable;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
+import org.springframework.transaction.TransactionDefinition;
+import org.springframework.transaction.reactive.TransactionalOperator;
 import org.springframework.util.ClassUtils;
 
 /**
@@ -337,16 +338,26 @@ void savePublisherOfEntitiesShouldInsertEntity() {
 		assertThat(boyd.getId()).isNotNull();
 	}
 
-	@RepeatedTest(10)
+	@RepeatedTest(10) // GH-4838
 	void transactionalSaveAllForStuffThatIsConsideredAnUpdateOfExistingData() {
 
 		ReactiveMongoTransactionManager txmgr = new ReactiveMongoTransactionManager(template.getMongoDatabaseFactory());
 		TransactionalOperator.create(txmgr, TransactionDefinition.withDefaults()).execute(callback -> {
 			return repository.saveAll(Arrays.asList(oliver, dave, carter, boyd, stefan, leroi, alicia));
-		})
-		.as(StepVerifier::create) //
-			.expectNext(oliver, dave, carter, boyd, stefan, leroi, alicia)
-			.verifyComplete();
+		}).as(StepVerifier::create) //
+				.expectNext(oliver, dave, carter, boyd, stefan, leroi, alicia).verifyComplete();
+	}
+
+	@RepeatedTest(10) // GH-4838
+	void transactionalSaveAllWithPublisherForStuffThatIsConsideredAnUpdateOfExistingData() {
+
+		ReactiveMongoTransactionManager txmgr = new ReactiveMongoTransactionManager(template.getMongoDatabaseFactory());
+		Flux<ReactivePerson> personFlux = Flux.fromStream(Stream.of(oliver, dave, carter, boyd, stefan, leroi, alicia));
+
+		TransactionalOperator.create(txmgr, TransactionDefinition.withDefaults()).execute(callback -> {
+				return repository.saveAll(personFlux);
+			}).as(StepVerifier::create) //
+				.expectNextCount(7).verifyComplete();
 	}
 
 	@Test // GH-3609
@@ -358,7 +369,7 @@ void savePublisherOfImmutableEntitiesShouldInsertEntity() {
 				.consumeNextWith(actual -> {
 					assertThat(actual.id).isNotNull();
 				}) //
-			.verifyComplete();
+				.verifyComplete();
 	}
 
 	@Test // DATAMONGO-1444