Honour null generic wrapper results and await plain futures

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

@@ -22,6 +22,9 @@ import org.slf4j.LoggerFactory
 import java.lang.reflect.InvocationTargetException
 import java.lang.reflect.Method
 import java.util.*
+import java.util.concurrent.CompletionStage
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.Future
 import java.util.function.Supplier
 import kotlin.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
 import kotlin.reflect.full.valueParameters
@@ -122,12 +125,15 @@ internal class MethodFieldResolver(
             }
         }
 
+        // a returned future is only waited on if the method's declared raw return type is Future, so a returned object that merely implements Future is left alone
+        val returnsFuture = TypeUtils.getRawType(getReturnType(), null) == Future::class.java
+
         return if (numberOfParameters > 0 || isSuspendFunction) {
             // requires arguments and environment or is a suspend function
-            MethodFieldResolverDataFetcher(createSourceResolver(), this.method, args, options, isSuspendFunction)
+            MethodFieldResolverDataFetcher(createSourceResolver(), this.method, args, options, isSuspendFunction, returnsFuture)
         } else {
             // if there are no parameters an optimized version of the data fetcher can be used
-            LightMethodFieldResolverDataFetcher(createSourceResolver(), this.method, options)
+            LightMethodFieldResolverDataFetcher(createSourceResolver(), this.method, options, returnsFuture)
         }
     }
 
@@ -154,11 +160,7 @@ internal class MethodFieldResolver(
     }
 
     override fun scanForMatches(): List<TypeClassMatcher.PotentialMatch> {
-        val returnType = try {
-            method.kotlinFunction?.returnType?.javaType ?: method.genericReturnType
-        } catch (e: InternalError) {
-            method.genericReturnType
-        }
+        val returnType = getReturnType()
         val unwrappedGenericType = genericType.unwrapGenericType(if (search.isSubscription) returnType.asPublisher() else returnType)
         val returnValueMatch = TypeClassMatcher.PotentialMatch.returnValue(field.type, unwrappedGenericType, genericType, SchemaClassScanner.ReturnValueReference(method))
 
@@ -174,6 +176,13 @@ internal class MethodFieldResolver(
             ?.let { TypeUtils.parameterize(Publisher::class.java, it) }
             ?: this
 
+    private fun getReturnType(): JavaType =
+        try {
+            method.kotlinFunction?.returnType?.javaType ?: method.genericReturnType
+        } catch (e: InternalError) {
+            method.genericReturnType
+        }
+
     private fun getIndexOffset(): Int {
         return if (resolverInfo is DataClassTypeResolverInfo && !method.declaringClass.isAssignableFrom(resolverInfo.dataClassType)) {
             1
@@ -203,7 +212,8 @@ internal class MethodFieldResolverDataFetcher(
     private val method: Method,
     private val args: List<ArgumentPlaceholder>,
     private val options: SchemaParserOptions,
-    private val isSuspendFunction: Boolean
+    private val isSuspendFunction: Boolean,
+    private val returnsFuture: Boolean
 ) : DataFetcher<Any> {
 
     override fun get(environment: DataFetchingEnvironment): Any? {
@@ -212,10 +222,10 @@ internal class MethodFieldResolverDataFetcher(
 
         return if (isSuspendFunction) {
             environment.coroutineScope().future(options.coroutineContextProvider.provide()) {
-                invokeSuspend(source, method, args)?.transformWithGenericWrapper(options.genericWrappers) { environment }
+                invokeSuspend(source, method, args)?.transformWithGenericWrapper(options.genericWrappers, returnsFuture) { environment }
             }
         } else {
-            invoke(method, source, args)?.transformWithGenericWrapper(options.genericWrappers) { environment }
+            invoke(method, source, args)?.transformWithGenericWrapper(options.genericWrappers, returnsFuture) { environment }
         }
     }
 
@@ -234,13 +244,14 @@ internal class MethodFieldResolverDataFetcher(
 internal class LightMethodFieldResolverDataFetcher(
     private val sourceResolver: SourceResolver,
     private val method: Method,
-    private val options: SchemaParserOptions
+    private val options: SchemaParserOptions,
+    private val returnsFuture: Boolean
 ) : LightDataFetcher<Any?> {
 
     override fun get(fieldDefinition: GraphQLFieldDefinition?, sourceObject: Any?, environmentSupplier: Supplier<DataFetchingEnvironment>): Any? {
         val source = sourceResolver.resolve(null, sourceObject)
 
-        return invoke(method, source, emptyArray())?.transformWithGenericWrapper(options.genericWrappers, environmentSupplier)
+        return invoke(method, source, emptyArray())?.transformWithGenericWrapper(options.genericWrappers, returnsFuture, environmentSupplier)
     }
 
     override fun get(environment: DataFetchingEnvironment): Any? {
@@ -258,14 +269,34 @@ internal class LightMethodFieldResolverDataFetcher(
 
 private fun Any.transformWithGenericWrapper(
     genericWrappers: List<GenericWrapper>,
+    returnsFuture: Boolean,
     environmentSupplier: Supplier<DataFetchingEnvironment>
-): Any {
-    return genericWrappers
+): Any? {
+    val genericWrapper = genericWrappers
         .asSequence()
         .filter { it.type.isInstance(this) }
         .sortedWith(CompareGenericWrappers)
         .firstOrNull()
-        ?.transformer?.invoke(this, environmentSupplier.get()) ?: this
+        ?: return this
+
+    val value = genericWrapper.transformer(this, environmentSupplier.get())
+    return if (returnsFuture) value.awaitFuture() else value
+}
+
+// graphql-java only awaits completion stages, so any other future is waited on by blocking the fetching thread
+private fun Any?.awaitFuture(): Any? {
+    if (this !is Future<*> || this is CompletionStage<*>) {
+        return this
+    }
+
+    return try {
+        get()
+    } catch (e: ExecutionException) {
+        throw e.cause ?: e
+    } catch (e: InterruptedException) {
+        Thread.currentThread().interrupt()
+        throw e
+    }
 }
 
 private class CompareGenericWrappers {

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

@@ -239,6 +239,51 @@ class MethodFieldResolverTest {
         assertEquals(exceptionWhileDataFetching.exception.message, "Whoops")
     }
 
+    @Test
+    fun `should return null when generic wrapper transformer returns null`() {
+        val schema = SchemaParser.newParser()
+            .schemaString(
+                """
+                type Query {
+                    name: String
+                    names: [String!]
+                    nameWithArgument(input: String): String
+                }
+                """)
+            .resolvers(object : GraphQLQueryResolver {
+                fun name(): Maybe<String> = Maybe.None
+                fun names(): Maybe<List<String>> = Maybe.None
+                fun nameWithArgument(input: String?): Maybe<String> = Maybe.None
+            })
+            .options(SchemaParserOptions.newOptions()
+                .genericWrappers(SchemaParserOptions.GenericWrapper.withTransformer(Maybe::class, 0, { maybe -> (maybe as? Maybe.Some)?.value }))
+                .build())
+            .build()
+            .makeExecutableSchema()
+
+        val gql = GraphQL.newGraphQL(schema).build()
+        val result = gql.execute(ExecutionInput.newExecutionInput().query(
+            """
+            query {
+                name
+                names
+                nameWithArgument(input: "test")
+            }
+            """))
+
+        assertEquals(result.errors, listOf())
+        assertEquals(result.getData(), mapOf(
+            "name" to null,
+            "names" to null,
+            "nameWithArgument" to null
+        ))
+    }
+
+    sealed class Maybe<out T> {
+        data class Some<T>(val value: T) : Maybe<T>()
+        object None : Maybe<Nothing>()
+    }
+
     /**
      * Custom Scalar Class type that doesn't work with Jackson serialization/deserialization
      */

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

@@ -1,5 +1,6 @@
 package graphql.kickstart.tools
 
+import graphql.ExceptionWhileDataFetching
 import graphql.GraphQL
 import graphql.execution.AsyncExecutionStrategy
 import graphql.kickstart.tools.SchemaParser.Companion.newParser
@@ -9,6 +10,7 @@ import org.junit.Test
 import java.util.*
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.Future
+import java.util.concurrent.FutureTask
 
 //import io.reactivex.Single;
 //import io.reactivex.internal.operators.single.SingleJust;
@@ -39,6 +41,73 @@ class ReactiveTest {
         }
     }
 
+    @Test
+    fun `future that is not a completion stage succeeds`() {
+        val schema = newParser().file("Reactive.graphqls")
+            .resolvers(object : GraphQLQueryResolver {
+                fun organization(organizationId: Int): Future<Organization> =
+                    FutureTask { Organization(User(organizationId.toLong())) }.also { Thread(it).start() }
+            })
+            .build()
+            .makeExecutableSchema()
+
+        val gql = GraphQL.newGraphQL(schema).build()
+
+        val data = assertNoGraphQlErrors(gql) {
+            "query { organization(organizationId: 1) { user { id } } }"
+        }
+
+        assertEquals(data, mapOf("organization" to mapOf("user" to mapOf("id" to "1"))))
+    }
+
+    @Test
+    fun `future that is not a completion stage rethrows its failure`() {
+        val schema = newParser().file("Reactive.graphqls")
+            .resolvers(object : GraphQLQueryResolver {
+                fun organization(organizationId: Int): Future<Organization> =
+                    FutureTask<Organization> { throw IllegalStateException("Whoops") }.also { it.run() }
+            })
+            .build()
+            .makeExecutableSchema()
+
+        val gql = GraphQL.newGraphQL(schema).build()
+        val result = gql.execute("query { organization(organizationId: 1) { user { id } } }")
+
+        assertEquals(result.errors.size, 1)
+        val exceptionWhileDataFetching = result.errors[0] as ExceptionWhileDataFetching
+        assertEquals(exceptionWhileDataFetching.exception.message, "Whoops")
+    }
+
+    @Test
+    fun `future is only waited on if the resolver declares a future`() {
+        val schema = newParser()
+            .schemaString(
+                """
+                type Query {
+                    job: Job
+                    futureJob: Job
+                }
+
+                type Job {
+                    id: ID
+                }
+                """)
+            .resolvers(object : GraphQLQueryResolver {
+                fun job(): Job = Job(1)
+                fun futureJob(): Future<Job> = FutureTask { Job(2) }.also { it.run() }
+            })
+            .build()
+            .makeExecutableSchema()
+
+        val gql = GraphQL.newGraphQL(schema).build()
+
+        val data = assertNoGraphQlErrors(gql) {
+            "query { job { id } futureJob { id } }"
+        }
+
+        assertEquals(data, mapOf("job" to mapOf("id" to "1"), "futureJob" to mapOf("id" to "2")))
+    }
+
     private class Query : GraphQLQueryResolver {
         //        Single<Optional<Organization>> organization(int organizationid) {
         //            return Single.just(Optional.empty()); //CompletableFuture.completedFuture(null);
@@ -48,12 +117,11 @@ class ReactiveTest {
         }
     }
 
-    private class Organization {
-        private val user: User? = null
-    }
+    private class Organization(private val user: User? = null)
 
-    private class User {
-        private val id: Long? = null
+    private class User(private val id: Long? = null) {
         private val name: String? = null
     }
+
+    private class Job(private val id: Long) : Future<String> by CompletableFuture.completedFuture("done")
 }