DATAMONGO-2255 - Add Coroutines Flow extensions

spring-data-mongodb/src/main/kotlin/org/springframework/data/mongodb/core/ReactiveAggregationOperationExtensions.kt

@@ -15,6 +15,9 @@
  */
 package org.springframework.data.mongodb.core
 
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.reactive.flow.asFlow
 import kotlin.reflect.KClass
 
 /**
@@ -35,3 +38,16 @@ fun <T : Any> ReactiveAggregationOperation.aggregateAndReturn(entityClass: KClas
  */
 inline fun <reified T : Any> ReactiveAggregationOperation.aggregateAndReturn(): ReactiveAggregationOperation.ReactiveAggregation<T> =
         aggregateAndReturn(T::class.java)
+
+/**
+ * Coroutines [Flow] variant of [ReactiveAggregationOperation.TerminatingAggregationOperation.all].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ * @since 2.2
+ */
+@FlowPreview
+fun <T : Any> ReactiveAggregationOperation.TerminatingAggregationOperation<T>.allAsFlow(batchSize: Int = 1): Flow<T> =
+		all().asFlow(batchSize)

spring-data-mongodb/src/main/kotlin/org/springframework/data/mongodb/core/ReactiveFindOperationExtensions.kt

@@ -15,8 +15,12 @@
  */
 package org.springframework.data.mongodb.core
 
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.reactive.awaitFirstOrNull
 import kotlinx.coroutines.reactive.awaitSingle
+import kotlinx.coroutines.reactive.flow.asFlow
+import org.springframework.data.geo.GeoResult
 import kotlin.reflect.KClass
 
 /**
@@ -129,3 +133,52 @@ suspend fun <T : Any> ReactiveFindOperation.TerminatingFind<T>.awaitCount(): Lon
  */
 suspend fun <T : Any> ReactiveFindOperation.TerminatingFind<T>.awaitExists(): Boolean =
 		exists().awaitSingle()
+
+/**
+ * Coroutines [Flow] variant of [ReactiveFindOperation.TerminatingFind.all].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ */
+@FlowPreview
+fun <T : Any> ReactiveFindOperation.TerminatingFind<T>.allAsFlow(batchSize: Int = 1): Flow<T> =
+		all().asFlow(batchSize)
+
+/**
+ * Coroutines [Flow] variant of [ReactiveFindOperation.TerminatingFind.tail].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ */
+@FlowPreview
+fun <T : Any> ReactiveFindOperation.TerminatingFind<T>.tailAsFlow(batchSize: Int = 1): Flow<T> =
+		tail().asFlow(batchSize)
+
+/**
+ * Coroutines [Flow] variant of [ReactiveFindOperation.TerminatingFindNear.all].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ */
+@FlowPreview
+fun <T : Any> ReactiveFindOperation.TerminatingFindNear<T>.allAsFlow(batchSize: Int = 1): Flow<GeoResult<T>> =
+		all().asFlow(batchSize)
+
+/**
+ * Coroutines [Flow] variant of [ReactiveFindOperation.TerminatingDistinct.all].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ * @since 2.2
+ */
+@FlowPreview
+fun <T : Any> ReactiveFindOperation.TerminatingDistinct<T>.allAsFlow(batchSize: Int = 1): Flow<T> =
+		all().asFlow(batchSize)

spring-data-mongodb/src/main/kotlin/org/springframework/data/mongodb/core/ReactiveInsertOperationExtensions.kt

@@ -15,7 +15,10 @@
  */
 package org.springframework.data.mongodb.core
 
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.reactive.awaitSingle
+import kotlinx.coroutines.reactive.flow.asFlow
 import kotlin.reflect.KClass
 
 /**
@@ -45,3 +48,17 @@ inline fun <reified T : Any> ReactiveInsertOperation.insert(): ReactiveInsertOpe
  */
 suspend inline fun <reified T: Any> ReactiveInsertOperation.TerminatingInsert<T>.oneAndAwait(o: T): T =
 		one(o).awaitSingle()
+
+
+/**
+ * Coroutines [Flow] variant of [ReactiveInsertOperation.TerminatingInsert.all].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ * @since 2.2
+ */
+@FlowPreview
+fun <T : Any> ReactiveInsertOperation.TerminatingInsert<T>.allAsFlow(objects: Collection<T>, batchSize: Int = 1): Flow<T> =
+		all(objects).asFlow(batchSize)

spring-data-mongodb/src/main/kotlin/org/springframework/data/mongodb/core/ReactiveMapReduceOperationExtensions.kt

@@ -15,6 +15,9 @@
  */
 package org.springframework.data.mongodb.core
 
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.reactive.flow.asFlow
 import kotlin.reflect.KClass
 
 /**
@@ -54,3 +57,17 @@ fun <T : Any> ReactiveMapReduceOperation.MapReduceWithProjection<*>.asType(resul
  */
 inline fun <reified T : Any> ReactiveMapReduceOperation.MapReduceWithProjection<*>.asType(): ReactiveMapReduceOperation.MapReduceWithQuery<T> =
 		`as`(T::class.java)
+
+
+/**
+ * Coroutines [Flow] variant of [ReactiveMapReduceOperation.TerminatingMapReduce.all].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ * @since 2.2
+ */
+@FlowPreview
+fun <T : Any> ReactiveMapReduceOperation.TerminatingMapReduce<T>.allAsFlow(batchSize: Int = 1): Flow<T> =
+		all().asFlow(batchSize)

spring-data-mongodb/src/main/kotlin/org/springframework/data/mongodb/core/ReactiveRemoveOperationExtensions.kt

@@ -16,7 +16,10 @@
 package org.springframework.data.mongodb.core
 
 import com.mongodb.client.result.DeleteResult
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.reactive.awaitSingle
+import kotlinx.coroutines.reactive.flow.asFlow
 import kotlin.reflect.KClass
 
 /**
@@ -46,3 +49,17 @@ inline fun <reified T : Any> ReactiveRemoveOperation.remove(): ReactiveRemoveOpe
  */
 suspend fun <T : Any> ReactiveRemoveOperation.TerminatingRemove<T>.allAndAwait(): DeleteResult =
 		all().awaitSingle()
+
+
+/**
+ * Coroutines [Flow] variant of [ReactiveRemoveOperation.TerminatingRemove.findAndRemove].
+ *
+ * Backpressure is controlled by [batchSize] parameter that controls the size of in-flight elements
+ * and [org.reactivestreams.Subscription.request] size.
+ *
+ * @author Sebastien Deleuze
+ * @since 2.2
+ */
+@FlowPreview
+fun <T : Any> ReactiveRemoveOperation.TerminatingRemove<T>.findAndRemoveAsFlow(batchSize: Int = 1): Flow<T> =
+		findAndRemove().asFlow(batchSize)

spring-data-mongodb/src/test/kotlin/org/springframework/data/mongodb/core/ReactiveAggregationOperationExtensionsTests.kt

@@ -16,9 +16,15 @@
 package org.springframework.data.mongodb.core
 
 import example.first.First
+import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.runBlocking
+import org.assertj.core.api.Assertions
 import org.junit.Test
+import reactor.core.publisher.Flux
 
 /**
  * @author Mark Paluch
@@ -42,4 +48,20 @@ class ReactiveAggregationOperationExtensionsTests {
 		verify { operation.aggregateAndReturn(First::class.java) }
 	}
 
+	@Test
+	@FlowPreview
+	fun terminatingAggregationOperationAllAsFlow() {
+
+		val spec = mockk<ReactiveAggregationOperation.TerminatingAggregationOperation<String>>()
+		every { spec.all() } returns Flux.just("foo", "bar", "baz")
+
+		runBlocking {
+			Assertions.assertThat(spec.allAsFlow().toList()).contains("foo", "bar", "baz")
+		}
+
+		verify {
+			spec.all()
+		}
+	}
+
 }

spring-data-mongodb/src/test/kotlin/org/springframework/data/mongodb/core/ReactiveFindOperationExtensionsTests.kt

@@ -19,10 +19,15 @@ import example.first.First
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
 import org.assertj.core.api.Assertions.assertThat
 import org.assertj.core.api.Assertions.assertThatExceptionOfType
 import org.junit.Test
+import org.springframework.data.geo.Distance
+import org.springframework.data.geo.GeoResult
+import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
 
 /**
@@ -228,4 +233,71 @@ class ReactiveFindOperationExtensionsTests {
 			find.exists()
 		}
 	}
+
+	@Test
+	@FlowPreview
+	fun terminatingFindAllAsFlow() {
+
+		val spec = mockk<ReactiveFindOperation.TerminatingFind<String>>()
+		every { spec.all() } returns Flux.just("foo", "bar", "baz")
+
+		runBlocking {
+			assertThat(spec.allAsFlow().toList()).contains("foo", "bar", "baz")
+		}
+
+		verify {
+			spec.all()
+		}
+	}
+
+	@Test
+	@FlowPreview
+	fun terminatingFindTailAsFlow() {
+
+		val spec = mockk<ReactiveFindOperation.TerminatingFind<String>>()
+		every { spec.tail() } returns Flux.just("foo", "bar", "baz")
+
+		runBlocking {
+			assertThat(spec.tailAsFlow().toList()).contains("foo", "bar", "baz")
+		}
+
+		verify {
+			spec.tail()
+		}
+	}
+
+	@Test
+	@FlowPreview
+	fun terminatingFindNearAllAsFlow() {
+
+		val spec = mockk<ReactiveFindOperation.TerminatingFindNear<String>>()
+		val foo = GeoResult("foo", Distance(0.0))
+		val bar = GeoResult("bar", Distance(0.0))
+		val baz = GeoResult("baz", Distance(0.0))
+		every { spec.all() } returns Flux.just(foo, bar, baz)
+
+		runBlocking {
+			assertThat(spec.allAsFlow().toList()).contains(foo, bar, baz)
+		}
+
+		verify {
+			spec.all()
+		}
+	}
+
+	@Test
+	@FlowPreview
+	fun terminatingDistinctAllAsFlow() {
+
+		val spec = mockk<ReactiveFindOperation.TerminatingDistinct<String>>()
+		every { spec.all() } returns Flux.just("foo", "bar", "baz")
+
+		runBlocking {
+			assertThat(spec.allAsFlow().toList()).contains("foo", "bar", "baz")
+		}
+
+		verify {
+			spec.all()
+		}
+	}
 }

spring-data-mongodb/src/test/kotlin/org/springframework/data/mongodb/core/ReactiveInsertOperationExtensionsTests.kt

@@ -19,9 +19,12 @@ import example.first.First
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Test
+import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
 
 /**
@@ -47,17 +50,34 @@ class ReactiveInsertOperationExtensionsTests {
 	}
 
 	@Test // DATAMONGO-2209
-	fun terminatingFindAwaitOne() {
+	fun terminatingInsertOneAndAwait() {
 
-		val find = mockk<ReactiveInsertOperation.TerminatingInsert<String>>()
-		every { find.one("foo") } returns Mono.just("foo")
+		val insert = mockk<ReactiveInsertOperation.TerminatingInsert<String>>()
+		every { insert.one("foo") } returns Mono.just("foo")
 
 		runBlocking {
-			assertThat(find.oneAndAwait("foo")).isEqualTo("foo")
+			assertThat(insert.oneAndAwait("foo")).isEqualTo("foo")
 		}
 
 		verify {
-			find.one("foo")
+			insert.one("foo")
+		}
+	}
+
+	@Test
+	@FlowPreview
+	fun terminatingInsertAllAsFlow() {
+
+		val insert = mockk<ReactiveInsertOperation.TerminatingInsert<String>>()
+		val list = listOf("foo", "bar")
+		every { insert.all(any()) } returns Flux.fromIterable(list)
+
+		runBlocking {
+			assertThat(insert.allAsFlow(list).toList()).containsAll(list)
+		}
+
+		verify {
+			insert.all(list)
 		}
 	}
 }

spring-data-mongodb/src/test/kotlin/org/springframework/data/mongodb/core/ReactiveMapReduceOperationExtensionsTests.kt

@@ -16,9 +16,15 @@
 package org.springframework.data.mongodb.core
 
 import example.first.First
+import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
+import kotlinx.coroutines.FlowPreview
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.runBlocking
+import org.assertj.core.api.Assertions
 import org.junit.Test
+import reactor.core.publisher.Flux
 
 /**
  * @author Christoph Strobl
@@ -57,4 +63,20 @@ class ReactiveMapReduceOperationExtensionsTests {
 		operationWithProjection.asType<User>()
 		verify { operationWithProjection.`as`(User::class.java) }
 	}
+
+	@Test
+	@FlowPreview
+	fun terminatingMapReduceAllAsFlow() {
+
+		val spec = mockk<ReactiveMapReduceOperation.TerminatingMapReduce<String>>()
+		every { spec.all() } returns Flux.just("foo", "bar", "baz")
+
+		runBlocking {
+			Assertions.assertThat(spec.allAsFlow().toList()).contains("foo", "bar", "baz")
+		}
+
+		verify {
+			spec.all()
+		}
+	}
 }