Port graphql-java 15 subscription changes to Flow strategy (#742)

In graphql-java 15, changes were made to SubscriptionExecutionStrategy
to allow for DataFetcherResults to be returned, and to instrument
the results of a subscription properly.  This updates the Flow
strategy with the same changes and adds tests for both changes.

See graphql-java/graphql-java#1804

Author: josephlbarnett

graphql-kotlin-schema-generator/src/main/kotlin/com/expediagroup/graphql/execution/FlowSubscriptionExecutionStrategy.kt

@@ -18,21 +18,30 @@ package com.expediagroup.graphql.execution
 
 import graphql.ExecutionResult
 import graphql.ExecutionResultImpl
+import graphql.GraphQLError
+import graphql.execution.AbsoluteGraphQLError
 import graphql.execution.DataFetcherExceptionHandler
+import graphql.execution.DataFetcherResult
 import graphql.execution.ExecutionContext
+import graphql.execution.ExecutionStepInfo
 import graphql.execution.ExecutionStrategy
 import graphql.execution.ExecutionStrategyParameters
 import graphql.execution.FetchedValue
 import graphql.execution.SimpleDataFetcherExceptionHandler
 import graphql.execution.SubscriptionExecutionStrategy
+import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
+import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters
+import graphql.execution.instrumentation.parameters.InstrumentationFieldParameters
 import graphql.execution.reactive.CompletionStageMappingPublisher
+import graphql.schema.GraphQLObjectType
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.future.await
 import kotlinx.coroutines.reactive.asFlow
 import org.reactivestreams.Publisher
 import java.util.Collections
 import java.util.concurrent.CompletableFuture
+import java.util.function.Consumer
 
 /**
  * [SubscriptionExecutionStrategy] replacement that returns an [ExecutionResult]
@@ -54,11 +63,15 @@ class FlowSubscriptionExecutionStrategy(dfe: DataFetcherExceptionHandler) : Exec
         parameters: ExecutionStrategyParameters
     ): CompletableFuture<ExecutionResult> {
 
+        val instrumentation = executionContext.instrumentation
+        val instrumentationParameters = InstrumentationExecutionStrategyParameters(executionContext, parameters)
+        val executionStrategyCtx = instrumentation.beginExecutionStrategy(instrumentationParameters)
+
         val sourceEventStream = createSourceEventStream(executionContext, parameters)
 
         //
         // when the upstream source event stream completes, subscribe to it and wire in our adapter
-        return sourceEventStream.thenApply { sourceFlow ->
+        val overallResult: CompletableFuture<ExecutionResult> = sourceEventStream.thenApply { sourceFlow ->
             if (sourceFlow == null) {
                 ExecutionResultImpl(null, executionContext.errors)
             } else {
@@ -68,6 +81,12 @@ class FlowSubscriptionExecutionStrategy(dfe: DataFetcherExceptionHandler) : Exec
                 ExecutionResultImpl(returnFlow, executionContext.errors)
             }
         }
+
+        // dispatched the subscription query
+        executionStrategyCtx.onDispatched(overallResult)
+        overallResult.whenComplete(executionStrategyCtx::onCompleted)
+
+        return overallResult
     }
 
     /*
@@ -118,15 +137,37 @@ class FlowSubscriptionExecutionStrategy(dfe: DataFetcherExceptionHandler) : Exec
         parameters: ExecutionStrategyParameters,
         eventPayload: Any?
     ): CompletableFuture<ExecutionResult> {
-        val newExecutionContext = executionContext.transform { builder -> builder.root(eventPayload) }
+        val instrumentation = executionContext.instrumentation
 
+        val newExecutionContext = executionContext.transform { builder ->
+            builder
+                .root(eventPayload)
+                .resetErrors()
+        }
         val newParameters = firstFieldOfSubscriptionSelection(parameters)
-        val fetchedValue = FetchedValue.newFetchedValue().fetchedValue(eventPayload)
-            .rawFetchedValue(eventPayload)
-            .localContext(parameters.localContext)
-            .build()
-        return completeField(newExecutionContext, newParameters, fetchedValue).fieldValue
+        val subscribedFieldStepInfo = createSubscribedFieldStepInfo(executionContext, newParameters)
+
+        val i13nFieldParameters = InstrumentationFieldParameters(executionContext, subscribedFieldStepInfo.fieldDefinition, subscribedFieldStepInfo)
+        val subscribedFieldCtx = instrumentation.beginSubscribedFieldEvent(i13nFieldParameters)
+
+        val fetchedValue = unboxPossibleDataFetcherResult(newExecutionContext, parameters, eventPayload)
+
+        val fieldValueInfo = completeField(newExecutionContext, newParameters, fetchedValue)
+        val overallResult = fieldValueInfo
+            .fieldValue
             .thenApply { executionResult -> wrapWithRootFieldName(newParameters, executionResult) }
+
+        // dispatch instrumentation so they can know about each subscription event
+        subscribedFieldCtx.onDispatched(overallResult)
+        overallResult.whenComplete(subscribedFieldCtx::onCompleted)
+
+        // allow them to instrument each ER should they want to
+        val i13ExecutionParameters = InstrumentationExecutionParameters(
+            executionContext.executionInput, executionContext.graphQLSchema, executionContext.instrumentationState)
+
+        return overallResult.thenCompose { executionResult ->
+            instrumentation.instrumentExecutionResult(executionResult, i13ExecutionParameters)
+        }
     }
 
     private fun wrapWithRootFieldName(
@@ -154,4 +195,50 @@ class FlowSubscriptionExecutionStrategy(dfe: DataFetcherExceptionHandler) : Exec
         val fieldPath = parameters.path.segment(ExecutionStrategy.mkNameForPath(firstField.singleField))
         return parameters.transform { builder -> builder.field(firstField).path(fieldPath) }
     }
+
+    private fun createSubscribedFieldStepInfo(
+        executionContext: ExecutionContext,
+        parameters: ExecutionStrategyParameters
+    ): ExecutionStepInfo {
+        val field = parameters.field.singleField
+        val parentType = parameters.executionStepInfo.unwrappedNonNullType as GraphQLObjectType
+        val fieldDef = getFieldDef(executionContext.graphQLSchema, parentType, field)
+        return createExecutionStepInfo(executionContext, parameters, fieldDef, parentType)
+    }
+
+    /**
+     * java->kotlin copy of [ExecutionStrategy.unboxPossibleDataFetcherResult] where it's package-private
+     */
+    private fun unboxPossibleDataFetcherResult(
+        executionContext: ExecutionContext,
+        parameters: ExecutionStrategyParameters,
+        result: Any?
+    ): FetchedValue {
+        return if (result is DataFetcherResult<*>) {
+            if (result.isMapRelativeErrors) {
+                result.errors.stream()
+                    .map { relError: GraphQLError? -> AbsoluteGraphQLError(parameters, relError) }
+                    .forEach { error: AbsoluteGraphQLError? -> executionContext.addError(error) }
+            } else {
+                result.errors.forEach(Consumer { error: GraphQLError? -> executionContext.addError(error) })
+            }
+            var localContext = result.localContext
+            if (localContext == null) {
+                // if the field returns nothing then they get the context of their parent field
+                localContext = parameters.localContext
+            }
+            FetchedValue.newFetchedValue()
+                .fetchedValue(executionContext.valueUnboxer.unbox(result.data))
+                .rawFetchedValue(result.data)
+                .errors(result.errors)
+                .localContext(localContext)
+                .build()
+        } else {
+            FetchedValue.newFetchedValue()
+                .fetchedValue(executionContext.valueUnboxer.unbox(result))
+                .rawFetchedValue(result)
+                .localContext(parameters.localContext)
+                .build()
+        }
+    }
 }

graphql-kotlin-schema-generator/src/test/kotlin/com/expediagroup/graphql/execution/FlowSubscriptionExecutionStrategyTest.kt

@@ -23,9 +23,13 @@ import com.expediagroup.graphql.hooks.FlowSubscriptionSchemaGeneratorHooks
 import com.expediagroup.graphql.toSchema
 import graphql.ExecutionInput
 import graphql.ExecutionResult
+import graphql.ExecutionResultImpl
 import graphql.GraphQL
 import graphql.GraphQLError
 import graphql.GraphqlErrorBuilder
+import graphql.execution.DataFetcherResult
+import graphql.execution.instrumentation.SimpleInstrumentation
+import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters
 import graphql.schema.GraphQLSchema
 import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.delay
@@ -38,6 +42,7 @@ import kotlinx.coroutines.reactive.asPublisher
 import kotlinx.coroutines.runBlocking
 import org.junit.jupiter.api.Test
 import org.reactivestreams.Publisher
+import java.util.concurrent.CompletableFuture
 import kotlin.test.assertEquals
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
@@ -47,14 +52,17 @@ class FlowSubscriptionExecutionStrategyTest {
 
     private val testSchema: GraphQLSchema = toSchema(
         config = SchemaGeneratorConfig(
-            supportedPackages = listOf("com.expediagroup.graphql.spring.execution"),
+            supportedPackages = listOf("com.expediagroup.graphql.execution"),
             hooks = FlowSubscriptionSchemaGeneratorHooks()
         ),
         queries = listOf(TopLevelObject(BasicQuery())),
         mutations = listOf(TopLevelObject(BasicQuery())),
         subscriptions = listOf(TopLevelObject(FlowSubscription()))
     )
-    private val testGraphQL: GraphQL = GraphQL.newGraphQL(testSchema).subscriptionExecutionStrategy(FlowSubscriptionExecutionStrategy()).build()
+    private val testGraphQL: GraphQL = GraphQL.newGraphQL(testSchema)
+        .subscriptionExecutionStrategy(FlowSubscriptionExecutionStrategy())
+        .instrumentation(TestInstrumentation())
+        .build()
 
     @Test
     fun `verify subscription to flow`() = runBlocking {
@@ -64,6 +72,24 @@ class FlowSubscriptionExecutionStrategyTest {
         val list = mutableListOf<Int>()
         flow.collect {
             list.add(it.getData<Map<String, Int>>().getValue("ticker"))
+            assertEquals(it.extensions["testKey"], "testValue")
+        }
+        assertEquals(5, list.size)
+        for (i in list.indices) {
+            assertEquals(i + 1, list[i])
+        }
+    }
+
+    @Test
+    fun `verify subscription to datafetcher flow`() = runBlocking {
+        val request = ExecutionInput.newExecutionInput().query("subscription { datafetcher }").build()
+        val response = testGraphQL.execute(request)
+        val flow = response.getData<Flow<ExecutionResult>>()
+        val list = mutableListOf<Int>()
+        flow.collect {
+            val intVal = it.getData<Map<String, Int>>().getValue("datafetcher")
+            list.add(intVal)
+            assertEquals(it.extensions["testKey"], "testValue")
         }
         assertEquals(5, list.size)
         for (i in list.indices) {
@@ -163,6 +189,20 @@ class FlowSubscriptionExecutionStrategyTest {
         fun query(): String = "hello"
     }
 
+    class TestInstrumentation : SimpleInstrumentation() {
+        override fun instrumentExecutionResult(
+            executionResult: ExecutionResult,
+            parameters: InstrumentationExecutionParameters?
+        ): CompletableFuture<ExecutionResult> {
+            return CompletableFuture.completedFuture(
+                ExecutionResultImpl.newExecutionResult()
+                    .from(executionResult)
+                    .addExtension("testKey", "testValue")
+                    .build()
+            )
+        }
+    }
+
     class FlowSubscription {
         fun ticker(): Flow<Int> {
             return flow {
@@ -173,6 +213,15 @@ class FlowSubscriptionExecutionStrategyTest {
             }
         }
 
+        fun datafetcher(): Flow<DataFetcherResult<Int>> {
+            return flow {
+                for (i in 1..5) {
+                    delay(100)
+                    emit(DataFetcherResult(i, listOf()))
+                }
+            }
+        }
+
         fun publisherTicker(): Publisher<Int> {
             return flow {
                 for (i in 1..5) {