Wait on a future inside a generic wrapper whose transformer returns a future

src/main/kotlin/graphql/kickstart/tools/SchemaParserOptions.kt

@@ -231,6 +231,9 @@ data class SchemaParserOptions internal constructor(
         constructor(type: Class<*>, index: Int) : this(type, index, { x, _ -> x })
         constructor(type: KClass<*>, index: Int) : this(type.java, index, { x, _ -> x })
 
+        internal val transformerReturnsCompletionStage
+            get() = transformer is CompletionStageTransformer<*>
+
         companion object {
             @Suppress("UNCHECKED_CAST")
             @JvmStatic
@@ -289,6 +292,34 @@ data class SchemaParserOptions internal constructor(
             ): GenericWrapper where T : Any {
                 return listCollectionWithTransformer(type.java, index, transformer)
             }
+
+            /**
+             * A wrapper whose [transformer] returns a future of its content, such as Reactor's `Mono::toFuture`. Unlike
+             * [withTransformer], a future declared as the content, as in `Mono<CompletableFuture<Item>>`, is then waited on too.
+             */
+            @JvmStatic
+            fun <T> withCompletionStageTransformer(
+                type: Class<T>,
+                index: Int,
+                transformer: (T) -> CompletionStage<*>?
+            ): GenericWrapper where T : Any {
+                return GenericWrapper(type, index, CompletionStageTransformer(type, transformer))
+            }
+
+            fun <T> withCompletionStageTransformer(
+                type: KClass<T>,
+                index: Int,
+                transformer: (T) -> CompletionStage<*>?
+            ): GenericWrapper where T : Any {
+                return withCompletionStageTransformer(type.java, index, transformer)
+            }
         }
     }
 }
+
+private class CompletionStageTransformer<T : Any>(
+    private val type: Class<T>,
+    private val transformer: (T) -> CompletionStage<*>?
+) : (Any, DataFetchingEnvironment) -> Any? {
+    override fun invoke(value: Any, environment: DataFetchingEnvironment): Any? = transformer(type.cast(value))
+}

src/main/kotlin/graphql/kickstart/tools/resolver/ReturnValueUnwrapper.kt

@@ -23,8 +23,14 @@ internal class ReturnValueUnwrapper private constructor(
 ) {
 
     fun unwrap(value: Any, environmentSupplier: Supplier<DataFetchingEnvironment>): Any? {
-        val unwrapped = awaitable.requireOutermost(value.transformWithGenericWrapper(genericWrappers, environmentSupplier))
-        return if (isOutermostFuture) unwrapped.then(nested, environmentSupplier) else unwrapped.thenUnlessNextAwaits(nested, environmentSupplier)
+        val wrapper = genericWrappers.wrapperOf(value)
+        val transformed = if (wrapper == null) value else wrapper.transformer(value, environmentSupplier.get())
+        val unwrapped = awaitable.requireOutermost(transformed)
+        return if (isOutermostFuture || wrapper?.transformerReturnsCompletionStage == true) {
+            unwrapped.then(nested, environmentSupplier)
+        } else {
+            unwrapped.thenUnlessNextAwaits(nested, environmentSupplier)
+        }
     }
 
     suspend fun unwrapSuspending(value: Any, environmentSupplier: Supplier<DataFetchingEnvironment>): Any? {
@@ -69,7 +75,11 @@ private class Await(private val wrapper: GenericWrapper, private val awaitable:
 private class Peel(private val wrapper: GenericWrapper, private val awaitable: Awaitable, private val next: Step?) : Step {
     override fun unwrap(value: Any?, environmentSupplier: Supplier<DataFetchingEnvironment>): Any? {
         val transformed = value?.let { wrapper.transformer(it, environmentSupplier.get()) }
-        return if (next is Await) next.unwrap(transformed, environmentSupplier) else awaitable.requireNested(transformed).then(next, environmentSupplier)
+        return if (next is Await && !wrapper.transformerReturnsCompletionStage) {
+            next.unwrap(transformed, environmentSupplier)
+        } else {
+            awaitable.requireNested(transformed).then(next, environmentSupplier)
+        }
     }
 }
 
@@ -133,24 +143,16 @@ private fun List<Any?>.awaitAll(): Any {
 private fun Any?.asCompletionStage(): CompletionStage<Any?> =
     this as? CompletionStage<Any?> ?: CompletableFuture.completedFuture(this)
 
-private fun Any.transformWithGenericWrapper(
-    genericWrappers: List<GenericWrapper>,
-    environmentSupplier: Supplier<DataFetchingEnvironment>
-): Any? {
+private fun List<GenericWrapper>.wrapperOf(value: Any): GenericWrapper? {
     // the JIT doesn't optimize away the sequence below on this call path, so values that aren't wrappers, which is
     // most of them, return before it's allocated
-    if (genericWrappers.none { it.type.isInstance(this) }) {
-        return this
+    if (none { it.type.isInstance(value) }) {
+        return null
     }
-    val genericWrapper = genericWrappers
-        .asSequence()
-        .filter { it.type.isInstance(this) }
+    return asSequence()
+        .filter { it.type.isInstance(value) }
         .sortedWith(CompareGenericWrappers)
         .firstOrNull()
-        ?: return this
-
-    // returned as is, even if null, so a transformer can map a value to null
-    return genericWrapper.transformer(this, environmentSupplier.get())
 }
 
 // graphql-java only waits on a CompletionStage returned by a data fetcher, so any other future, or a future

src/test/kotlin/graphql/kickstart/tools/NestedGenericWrapperTest.kt

@@ -43,6 +43,7 @@ import org.junit.Test
 import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.CompletableFuture.completedFuture
+import java.util.concurrent.CompletionStage
 import java.util.concurrent.Future
 import java.util.concurrent.FutureTask
 
@@ -171,6 +172,19 @@ class NestedGenericWrapperTest {
     @Test
     fun `hands a future graphql-java can't wait on to its own wrapper`() = assertResolves("listOfWrappedPending", ab)
 
+    @Test
+    fun `waits on a future in a wrapper whose transformer returns a future`() = assertResolves("promisedFuture", mapOf("data" to "a"))
+
+    @Test
+    fun `waits on a scalar future in a wrapper whose transformer returns a future`() = assertResolves("promisedString", "a", selection = "")
+
+    @Test
+    fun `waits on futures in wrappers whose transformer returns a future in a list`() = assertResolves("listOfPromisedFutures", ab)
+
+    @Test
+    fun `resolves a union member that implements Future in a list of wrappers whose transformer returns a future`() =
+        assertResolves("listOfPromisedItems", ab, selection = "{ ... on Payload { data } ... on FutureItem { data } }")
+
     @Test
     fun `unwraps the wrappers a schemaWrapper introduces`() = assertResolves("batchOfPayloads", ab)
 
@@ -368,6 +382,10 @@ class NestedGenericWrapperTest {
                 plainList: [Payload!]!
                 futureOfList: [Payload!]!
                 wrappedFuture: Payload!
+                promisedFuture: Payload!
+                promisedString: String
+                listOfPromisedFutures: [Payload!]!
+                listOfPromisedItems: [Item!]!
                 ${if (scalars.isEmpty()) "" else "json: [JSON] futureJson: JSON jsonFutures: [JSON] wrappedJsonOfAny: [JSON]"}
             }
 
@@ -386,7 +404,7 @@ class NestedGenericWrapperTest {
         .resolvers(query)
         .dictionary(FutureItem::class)
         .scalars(*scalars)
-        .options(SchemaParserOptions.newOptions().genericWrappers(tryWrapper, pageWrapper, batchWrapper, boxWrapper, specialBoxWrapper, laterWrapper, pendingWrapper).build())
+        .options(SchemaParserOptions.newOptions().genericWrappers(tryWrapper, pageWrapper, batchWrapper, boxWrapper, specialBoxWrapper, laterWrapper, pendingWrapper, promisedWrapper).build())
         .build()
         .makeExecutableSchema()
 
@@ -496,6 +514,10 @@ class NestedGenericWrapperTest {
         fun wrappedJsonOfAny(): List<Try<Any>> = listOf(Try.success(FutureLeaf("f")))
         fun batchOfPayloads(): Batch<Payload> = Batch(listOfFutures())
         fun listOfBoxes(): List<Box<Payload>> = listOf(Box(Payload("a")), SpecialBox(Payload("b")))
+        fun promisedFuture(): Promised<CompletableFuture<Payload>> = Promised(CompletableFuture.supplyAsync { completedFuture(Payload("a")) })
+        fun promisedString(): Promised<CompletionStage<String>> = Promised(CompletableFuture.supplyAsync { completedFuture("a") })
+        fun listOfPromisedFutures(): List<Promised<CompletableFuture<Payload>>> = listOfFutures().map { Promised(CompletableFuture.supplyAsync { it }) }
+        fun listOfPromisedItems(): List<Promised<Any>> = listOf(Promised(completedFuture(Payload("a"))), Promised(CompletableFuture.supplyAsync { FutureItem("b") }))
     }
 
     class Batch<T>(val items: List<CompletableFuture<T>>)
@@ -508,6 +530,8 @@ class NestedGenericWrapperTest {
 
     class Pending<T>(val value: T) : Future<T> by completedFuture(value)
 
+    class Promised<T>(val future: CompletableFuture<T>)
+
     companion object {
         private val tryWrapper = GenericWrapper.withTransformer(Try::class, 0, { wrapper -> wrapper.get() })
         private val pageWrapper = GenericWrapper.listCollectionWithTransformer(Page::class, 0, { page -> page.content })
@@ -518,6 +542,7 @@ class NestedGenericWrapperTest {
         private val specialBoxWrapper = GenericWrapper.withTransformer(SpecialBox::class, 0, { _ -> throw IllegalStateException("picked by the value's class") })
         private val laterWrapper = GenericWrapper.withTransformer(Later::class, 0, { later -> FutureTask { later.value }.apply { run() } })
         private val pendingWrapper = GenericWrapper.withTransformer(Pending::class, 0, { pending -> completedFuture(pending.value) })
+        private val promisedWrapper = GenericWrapper.withCompletionStageTransformer(Promised::class, 0, Promised<*>::future)
         private val payloadType = GraphQLObjectType.newObject()
             .name("Payload")
             .field(GraphQLFieldDefinition.newFieldDefinition().name("data").type(GraphQLNonNull.nonNull(Scalars.GraphQLString)))