Add support for Kotlin coroutine channels for subscription queries

Author: varahash

pom.xml

@@ -39,6 +39,11 @@
             <artifactId>kotlinx-coroutines-jdk8</artifactId>
             <version>1.0.0</version>
         </dependency>
+        <dependency>
+            <groupId>org.jetbrains.kotlinx</groupId>
+            <artifactId>kotlinx-coroutines-reactive</artifactId>
+            <version>1.0.0</version>
+        </dependency>
         <dependency>
             <groupId>com.graphql-java</groupId>
             <artifactId>graphql-java</artifactId>
@@ -120,6 +125,12 @@
             <version>2.1</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.reactivestreams</groupId>
+            <artifactId>reactive-streams-tck</artifactId>
+            <version>1.0.2</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

src/main/kotlin/com/coxautodev/graphql/tools/FieldResolverScanner.kt

@@ -10,6 +10,7 @@ import org.slf4j.LoggerFactory
 import java.lang.reflect.Modifier
 import java.lang.reflect.ParameterizedType
 import kotlin.reflect.full.valueParameters
+import kotlin.reflect.jvm.javaType
 import kotlin.reflect.jvm.kotlinFunction
 
 /**
@@ -114,9 +115,11 @@ internal class FieldResolverScanner(val options: SchemaParserOptions) {
             true
         }
 
-        val correctParameterCount = method.parameterCount == requiredCount ||
-                (method.parameterCount == (requiredCount + 1) && allowedLastArgumentTypes.contains(method.parameterTypes.last())) ||
-                (method.kotlinFunction?.run { isSuspend && valueParameters.size == requiredCount } == true)
+        val methodParameterCount = method.kotlinFunction?.valueParameters?.size ?: method.parameterCount
+        val methodLastParameter = method.kotlinFunction?.valueParameters?.lastOrNull()?.type?.javaType ?: method.parameterTypes.lastOrNull()
+
+        val correctParameterCount = methodParameterCount == requiredCount ||
+                (methodParameterCount == (requiredCount + 1) && allowedLastArgumentTypes.contains(methodLastParameter))
         return correctParameterCount && appropriateFirstParameter
     }
 

src/main/kotlin/com/coxautodev/graphql/tools/MethodFieldResolver.kt

@@ -155,31 +155,26 @@ open class MethodFieldResolverDataFetcher(private val sourceResolver: SourceReso
         val source = sourceResolver(environment)
         val args = this.args.map { it(environment) }.toTypedArray()
 
-        val result = if (isSuspendFunction) {
+        return if (isSuspendFunction) {
             GlobalScope.future(options.coroutineContext) {
                 suspendCoroutineUninterceptedOrReturn<Any?> { continuation ->
-                    methodAccess.invoke(source, methodIndex, *args + continuation)
+                    methodAccess.invoke(source, methodIndex, *args + continuation)?.transformWithGenericWrapper(environment)
                 }
             }
         } else {
-            methodAccess.invoke(source, methodIndex, *args)
-        }
-        return if (result == null) {
-            null
-        } else {
-            val wrapper = options.genericWrappers
-                    .asSequence()
-                    .filter { it.type.isInstance(result) }
-                    .sortedWith(CompareGenericWrappers)
-                    .firstOrNull()
-            if (wrapper == null) {
-                result
-            } else {
-                wrapper.transformer.invoke(result, environment)
-            }
+            methodAccess.invoke(source, methodIndex, *args)?.transformWithGenericWrapper(environment)
         }
     }
 
+    private fun Any.transformWithGenericWrapper(environment: DataFetchingEnvironment): Any? {
+        return options.genericWrappers
+                .asSequence()
+                .filter { it.type.isInstance(this) }
+                .sortedWith(CompareGenericWrappers)
+                .firstOrNull()
+                ?.transformer?.invoke(this, environment) ?: this
+    }
+
     /**
      * Function that return the object used to fetch the data
      * It can be a DataFetcher or an entity

src/main/kotlin/com/coxautodev/graphql/tools/SchemaParserBuilder.kt

@@ -9,6 +9,9 @@ import graphql.parser.Parser
 import graphql.schema.DataFetchingEnvironment
 import graphql.schema.GraphQLScalarType
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.channels.ReceiveChannel
+import kotlinx.coroutines.reactive.publish
 import org.antlr.v4.runtime.RecognitionException
 import org.antlr.v4.runtime.misc.ParseCancellationException
 import org.reactivestreams.Publisher
@@ -331,20 +334,31 @@ data class SchemaParserOptions internal constructor(
         }
 
         fun build(): SchemaParserOptions {
+            val coroutineContext = coroutineContext ?: Dispatchers.Default
             val wrappers = if (useDefaultGenericWrappers) {
                 genericWrappers + listOf(
                         GenericWrapper(Future::class, 0),
                         GenericWrapper(CompletableFuture::class, 0),
                         GenericWrapper(CompletionStage::class, 0),
-                        GenericWrapper(Publisher::class, 0)
+                        GenericWrapper(Publisher::class, 0),
+                        GenericWrapper.withTransformer(ReceiveChannel::class, 0, { receiveChannel ->
+                            GlobalScope.publish(coroutineContext) {
+                                try {
+                                    for (item in receiveChannel) {
+                                        send(item)
+                                    }
+                                } finally {
+                                    receiveChannel.cancel()
+                                }
+                            }
+                        })
                 )
             } else {
                 genericWrappers
             }
 
             return SchemaParserOptions(contextClass, wrappers, allowUnimplementedResolvers, objectMapperProvider,
-                    proxyHandlers, preferGraphQLResolver, introspectionEnabled,
-                    coroutineContext ?: Dispatchers.Default)
+                    proxyHandlers, preferGraphQLResolver, introspectionEnabled, coroutineContext)
         }
     }
 

src/test/groovy/com/coxautodev/graphql/tools/EndToEndSpec.groovy

@@ -7,6 +7,7 @@ import graphql.execution.batched.BatchedExecutionStrategy
 import graphql.schema.GraphQLSchema
 import org.reactivestreams.Publisher
 import org.reactivestreams.Subscriber
+import org.reactivestreams.tck.TestEnvironment
 import spock.lang.Shared
 import spock.lang.Specification
 
@@ -563,18 +564,56 @@ class EndToEndSpec extends Specification {
 
     def "generated schema supports Kotlin suspend functions"() {
         when:
-        def data = Utils.assertNoGraphQlErrors(gql) {
-            '''
-                {
-                    coroutineItems {
-                        id
-                        name
+            def data = Utils.assertNoGraphQlErrors(gql) {
+                '''
+                    {
+                        coroutineItems {
+                            id
+                            name
+                        }
                     }
-                }
+                    '''
+            }
+
+        then:
+            data.coroutineItems == [[id:0, name:"item1"], [id:1, name:"item2"]]
+    }
+
+    def "generated schema supports Kotlin coroutine channels for the subscription query"() {
+        when:
+            def newItem = new Item(1, "item", Type.TYPE_1, UUID.randomUUID(), [])
+            def data = Utils.assertNoGraphQlErrors(gql, [:], new OnItemCreatedContext(newItem)) {
                 '''
-        }
+                    subscription {
+                        onItemCreatedCoroutineChannel {
+                            id
+                        }
+                    } 
+                    '''
+            }
+            def subscriber = new TestEnvironment().newManualSubscriber(data as Publisher<ExecutionResult>)
+
+        then:
+            subscriber.requestNextElement().data.get("onItemCreatedCoroutineChannel").id == 1
+            subscriber.expectCompletion()
+    }
+
+    def "generated schema supports Kotlin coroutine channels with suspend function for the subscription query"() {
+        when:
+            def newItem = new Item(1, "item", Type.TYPE_1, UUID.randomUUID(), [])
+            def data = Utils.assertNoGraphQlErrors(gql, [:], new OnItemCreatedContext(newItem)) {
+                '''
+                        subscription {
+                            onItemCreatedCoroutineChannelAndSuspendFunction {
+                                id
+                            }
+                        } 
+                        '''
+            }
+            def subscriber = new TestEnvironment().newManualSubscriber(data as Publisher<ExecutionResult>)
 
         then:
-        data.coroutineItems == [[id:0, name:"item1"], [id:1, name:"item2"]]
+            subscriber.requestNextElement().data.get("onItemCreatedCoroutineChannelAndSuspendFunction").id == 1
+            subscriber.expectCompletion()
     }
 }

src/test/kotlin/com/coxautodev/graphql/tools/EndToEndSpec.kt

@@ -7,7 +7,9 @@ import graphql.language.StringValue
 import graphql.schema.Coercing
 import graphql.schema.DataFetchingEnvironment
 import graphql.schema.GraphQLScalarType
-import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.ReceiveChannel
 import org.reactivestreams.Publisher
 import java.util.Optional
 import java.util.UUID
@@ -107,6 +109,8 @@ type Mutation {
 
 type Subscription {
     onItemCreated: Item!
+    onItemCreatedCoroutineChannel: Item!
+    onItemCreatedCoroutineChannelAndSuspendFunction: Item!
 }
 
 input ItemSearchInput {
@@ -306,6 +310,20 @@ class Subscription : GraphQLSubscriptionResolver {
             subscriber.onNext(env.getContext<OnItemCreatedContext>().newItem)
 //            subscriber.onComplete()
         }
+
+    fun onItemCreatedCoroutineChannel(env: DataFetchingEnvironment): ReceiveChannel<Item> {
+        val channel = Channel<Item>(1)
+        channel.offer(env.getContext<OnItemCreatedContext>().newItem)
+        return channel
+    }
+
+    suspend fun onItemCreatedCoroutineChannelAndSuspendFunction(env: DataFetchingEnvironment): ReceiveChannel<Item> {
+        return coroutineScope {
+            val channel = Channel<Item>(1)
+            channel.offer(env.getContext<OnItemCreatedContext>().newItem)
+            channel
+        }
+    }
 }
 
 class ItemResolver : GraphQLResolver<Item> {

src/test/kotlin/com/coxautodev/graphql/tools/MethodFieldResolverDataFetcherTest.kt

@@ -11,8 +11,12 @@ import graphql.schema.DataFetchingEnvironment
 import graphql.schema.DataFetchingEnvironmentBuilder
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.ReceiveChannel
 import org.junit.Assert
 import org.junit.Test
+import org.reactivestreams.Publisher
+import org.reactivestreams.tck.TestEnvironment
 import java.util.concurrent.CompletableFuture
 import kotlin.coroutines.coroutineContext
 
@@ -39,6 +43,53 @@ class MethodFieldResolverDataFetcherTest {
         Assert.assertTrue(future.get())
     }
 
+    @Test
+    fun `canceling subscription Publisher also cancels underlying Kotlin coroutine channel`() {
+        // setup
+        val channel = Channel<String>(10)
+        channel.offer("A")
+        channel.offer("B")
+
+        val resolver = createFetcher("onDataNameChanged", object : GraphQLResolver<DataClass> {
+            fun onDataNameChanged(date: DataClass): ReceiveChannel<String> {
+                return channel
+            }
+        })
+
+        // expect
+        @Suppress("UNCHECKED_CAST")
+        val publisher = resolver.get(createEnvironment(DataClass())) as Publisher<String>
+        val subscriber = TestEnvironment().newManualSubscriber(publisher)
+
+        Assert.assertEquals("A", subscriber.requestNextElement())
+
+        subscriber.cancel()
+        Thread.sleep(100)
+        Assert.assertTrue(channel.isClosedForReceive)
+    }
+
+    @Test
+    fun `canceling underlying Kotlin coroutine channel also cancels subscription Publisher`() {
+        // setup
+        val channel = Channel<String>(10)
+        channel.offer("A")
+        channel.close(IllegalStateException("Channel error"))
+
+        val resolver = createFetcher("onDataNameChanged", object : GraphQLResolver<DataClass> {
+            fun onDataNameChanged(date: DataClass): ReceiveChannel<String> {
+                return channel
+            }
+        })
+
+        // expect
+        @Suppress("UNCHECKED_CAST")
+        val publisher = resolver.get(createEnvironment(DataClass())) as Publisher<String>
+        val subscriber = TestEnvironment().newManualSubscriber(publisher)
+
+        Assert.assertEquals("A", subscriber.requestNextElement())
+        subscriber.expectErrorWithMessage(IllegalStateException::class.java, "Channel error")
+    }
+
     private fun createFetcher(
             methodName: String,
             resolver: GraphQLResolver<*>,