Generate context with suspend and nullable in subscriptions

examples/build.gradle.kts

@@ -40,6 +40,7 @@ subprojects {
 
     val kotlinVersion: String by project
     val junitVersion: String by project
+    val kotlinCoroutinesVersion: String by project
 
     val detektVersion: String by project
     val ktlintVersion: String by project
@@ -50,6 +51,8 @@ subprojects {
 
     dependencies {
         implementation(kotlin("stdlib", kotlinVersion))
+        implementation(kotlin("reflect", kotlinVersion))
+        implementation("org.jetbrains.kotlinx:kotlinx-coroutines-jdk8:$kotlinCoroutinesVersion")
         testImplementation(kotlin("test-junit5", kotlinVersion))
         testImplementation("org.junit.jupiter:junit-jupiter-api:$junitVersion")
         testImplementation("org.junit.jupiter:junit-jupiter-engine:$junitVersion")

examples/server/ktor-server/src/main/kotlin/com/expediagroup/graphql/examples/server/ktor/KtorGraphQLContextFactory.kt

@@ -25,7 +25,7 @@ import io.ktor.request.ApplicationRequest
  */
 class KtorGraphQLContextFactory : GraphQLContextFactory<AuthorizedContext, ApplicationRequest> {
 
-    override fun generateContext(request: ApplicationRequest): AuthorizedContext {
+    override suspend fun generateContext(request: ApplicationRequest): AuthorizedContext {
         val loggedInUser = User(
             email = "fake@site.com",
             firstName = "Someone",

examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContext.kt

@@ -24,14 +24,14 @@ import org.springframework.web.reactive.socket.WebSocketSession
  * Simple [GraphQLContext] that holds extra value and the [ServerRequest]
  */
 class MyGraphQLContext(
-    val myCustomValue: String,
-    val request: ServerRequest
+    val request: ServerRequest,
+    val myCustomValue: String
 ) : GraphQLContext
 
 /**
  * Simple [GraphQLContext] that holds extra value and the [WebSocketSession]
  */
 class MySubscriptionGraphQLContext(
     val request: WebSocketSession,
-    var subscriptionValue: String? = null
+    var auth: String? = null
 ) : GraphQLContext

examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/context/MyGraphQLContextFactory.kt

@@ -29,17 +29,20 @@ import org.springframework.web.reactive.socket.WebSocketSession
 @Component
 class MyGraphQLContextFactory : SpringGraphQLContextFactory<MyGraphQLContext>() {
 
-    override fun generateContext(request: ServerRequest): MyGraphQLContext = MyGraphQLContext(
-        myCustomValue = request.headers().firstHeader("MyHeader") ?: "defaultContext",
-        request = request
+    override suspend fun generateContext(request: ServerRequest): MyGraphQLContext = MyGraphQLContext(
+        request = request,
+        myCustomValue = request.headers().firstHeader("MyHeader") ?: "defaultContext"
     )
 }
 
+/**
+ * [GraphQLContextFactory] that generates [MySubscriptionGraphQLContext] that will be available when processing subscription operations.
+ */
 @Component
 class MySubscriptionGraphQLContextFactory : SpringSubscriptionGraphQLContextFactory<MySubscriptionGraphQLContext>() {
 
-    override fun generateContext(request: WebSocketSession): MySubscriptionGraphQLContext = MySubscriptionGraphQLContext(
+    override suspend fun generateContext(request: WebSocketSession): MySubscriptionGraphQLContext = MySubscriptionGraphQLContext(
         request = request,
-        subscriptionValue = null
+        auth = null
     )
 }

examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/execution/MySubscriptionHooks.kt

@@ -19,9 +19,7 @@ package com.expediagroup.graphql.examples.server.spring.execution
 import com.expediagroup.graphql.examples.server.spring.context.MySubscriptionGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.spring.subscriptions.ApolloSubscriptionHooks
-import kotlinx.coroutines.reactor.mono
 import org.springframework.web.reactive.socket.WebSocketSession
-import reactor.core.publisher.Mono
 
 /**
  * A simple implementation of Apollo Subscription Lifecycle Events.
@@ -31,12 +29,11 @@ class MySubscriptionHooks : ApolloSubscriptionHooks {
     override fun onConnect(
         connectionParams: Map<String, String>,
         session: WebSocketSession,
-        graphQLContext: GraphQLContext
-    ): Mono<GraphQLContext> = mono {
-        if (graphQLContext is MySubscriptionGraphQLContext) {
-            val bearer = connectionParams["Authorization"] ?: "none"
-            graphQLContext.subscriptionValue = bearer
+        graphQLContext: GraphQLContext?
+    ): GraphQLContext? {
+        if (graphQLContext != null && graphQLContext is MySubscriptionGraphQLContext) {
+            graphQLContext.auth = connectionParams["Authorization"]
         }
-        graphQLContext
+        return graphQLContext
     }
 }

examples/server/spring-server/src/main/kotlin/com/expediagroup/graphql/examples/server/spring/subscriptions/SimpleSubscription.kt

@@ -82,5 +82,5 @@ class SimpleSubscription : Subscription {
 
     @GraphQLDescription("Returns a value from the subscription context")
     fun subscriptionContext(myGraphQLContext: MySubscriptionGraphQLContext): Flux<String> =
-        Flux.just(myGraphQLContext.subscriptionValue ?: "", "value 2", "value3")
+        Flux.just(myGraphQLContext.auth ?: "no-auth")
 }

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

@@ -59,7 +59,7 @@ open class FunctionDataFetcher(
      * Invoke a suspend function or blocking function, passing in the [target] if not null or default to using the source from the environment.
      */
     override fun get(environment: DataFetchingEnvironment): Any? {
-        val instance = target ?: environment.getSource<Any?>()
+        val instance: Any? = target ?: environment.getSource<Any?>()
         val instanceParameter = fn.instanceParameter
 
         return if (instance != null && instanceParameter != null) {

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

@@ -16,19 +16,8 @@
 
 package com.expediagroup.graphql.generator.execution
 
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.ConcurrentMap
-
 /**
  * Marker interface to indicate that the implementing class should be considered
  * as the GraphQL context. This means the implementing class will not appear in the schema.
  */
 interface GraphQLContext
-
-/**
- * Default [GraphQLContext] that can be used if there is none provided. Exposes generic concurrent hash map
- * that can be populated with custom data.
- */
-class DefaultGraphQLContext : GraphQLContext {
-    val contents: ConcurrentMap<Any, Any> = ConcurrentHashMap()
-}

servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/execution/GraphQLContextFactory.kt

@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.execution
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 
 /**
@@ -28,12 +27,5 @@ interface GraphQLContextFactory<out Context : GraphQLContext, Request> {
      * Generate GraphQL context based on the incoming request and the corresponding response.
      * If no context should be generated and used in the request, return null.
      */
-    fun generateContext(request: Request): Context?
-}
-
-/**
- * Default context factory that generates GraphQL context with empty concurrent map that can store any elements.
- */
-class DefaultGraphQLContextFactory<T> : GraphQLContextFactory<DefaultGraphQLContext, T> {
-    override fun generateContext(request: T) = DefaultGraphQLContext()
+    suspend fun generateContext(request: Request): Context?
 }

servers/graphql-kotlin-server/src/main/kotlin/com/expediagroup/graphql/server/extensions/requestExtensions.kt

@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.extensions
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.types.GraphQLRequest
 import graphql.ExecutionInput
 import org.dataloader.DataLoaderRegistry
@@ -29,6 +28,6 @@ fun GraphQLRequest.toExecutionInput(graphQLContext: Any? = null, dataLoaderRegis
         .query(this.query)
         .operationName(this.operationName)
         .variables(this.variables ?: emptyMap())
-        .context(graphQLContext ?: DefaultGraphQLContext())
+        .context(graphQLContext)
         .dataLoaderRegistry(dataLoaderRegistry ?: DataLoaderRegistry())
         .build()

servers/graphql-kotlin-server/src/test/kotlin/com/expediagroup/graphql/server/extensions/RequestExtensionsKtTest.kt

@@ -16,14 +16,13 @@
 
 package com.expediagroup.graphql.server.extensions
 
-import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.types.GraphQLRequest
 import io.mockk.mockk
 import org.dataloader.DataLoaderRegistry
 import org.junit.jupiter.api.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
-import kotlin.test.assertTrue
+import kotlin.test.assertNull
 
 class RequestExtensionsKtTest {
 
@@ -32,7 +31,7 @@ class RequestExtensionsKtTest {
         val request = GraphQLRequest(query = "query { whatever }")
         val executionInput = request.toExecutionInput()
         assertEquals(request.query, executionInput.query)
-        assertTrue(executionInput.context is GraphQLContext)
+        assertNull(executionInput.context)
         assertNotNull(executionInput.dataLoaderRegistry)
     }
 

servers/graphql-kotlin-spring-server/build.gradle.kts

@@ -35,7 +35,7 @@ tasks {
                 limit {
                     counter = "BRANCH"
                     value = "COVEREDRATIO"
-                    minimum = "0.74".toBigDecimal()
+                    minimum = "0.70".toBigDecimal()
                 }
             }
         }

servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/SubscriptionAutoConfiguration.kt

@@ -67,7 +67,7 @@ class SubscriptionAutoConfiguration {
 
     @Bean
     @ConditionalOnMissingBean
-    fun springSubscriptionGraphQLContextFactory(): SpringSubscriptionGraphQLContextFactory<*> = DefaultSpringSubscriptionGraphQLContextFactory
+    fun springSubscriptionGraphQLContextFactory(): SpringSubscriptionGraphQLContextFactory<*> = DefaultSpringSubscriptionGraphQLContextFactory()
 
     @Bean
     fun apolloSubscriptionProtocolHandler(

servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/execution/SpringGraphQLContextFactory.kt

@@ -16,7 +16,6 @@
 
 package com.expediagroup.graphql.server.spring.execution
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.execution.GraphQLContextFactory
 import org.springframework.web.reactive.function.server.ServerRequest
@@ -27,8 +26,8 @@ import org.springframework.web.reactive.function.server.ServerRequest
 abstract class SpringGraphQLContextFactory<out T : GraphQLContext> : GraphQLContextFactory<T, ServerRequest>
 
 /**
- * Basic implementation of [SpringGraphQLContextFactory] that just returns a [DefaultGraphQLContext]
+ * Basic implementation of [SpringGraphQLContextFactory] that just returns null
  */
-class DefaultSpringGraphQLContextFactory : SpringGraphQLContextFactory<DefaultGraphQLContext>() {
-    override fun generateContext(request: ServerRequest) = DefaultGraphQLContext()
+class DefaultSpringGraphQLContextFactory : SpringGraphQLContextFactory<GraphQLContext>() {
+    override suspend fun generateContext(request: ServerRequest): GraphQLContext? = null
 }

servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionHooks.kt

@@ -17,7 +17,6 @@ package com.expediagroup.graphql.server.spring.subscriptions
 
 import com.expediagroup.graphql.generator.execution.GraphQLContext
 import org.springframework.web.reactive.socket.WebSocketSession
-import reactor.core.publisher.Mono
 
 /**
  * Implementation of Apollo Subscription Server Lifecycle Events
@@ -33,8 +32,8 @@ interface ApolloSubscriptionHooks {
     fun onConnect(
         connectionParams: Map<String, String>,
         session: WebSocketSession,
-        graphQLContext: GraphQLContext
-    ): Mono<GraphQLContext> = Mono.just(graphQLContext)
+        graphQLContext: GraphQLContext?
+    ): GraphQLContext? = graphQLContext
 
     /**
      * Called when the client executes a GraphQL operation.
@@ -43,7 +42,7 @@ interface ApolloSubscriptionHooks {
     fun onOperation(
         operationMessage: SubscriptionOperationMessage,
         session: WebSocketSession,
-        graphQLContext: GraphQLContext
+        graphQLContext: GraphQLContext?
     ): Unit = Unit
 
     /**

servers/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/server/spring/subscriptions/ApolloSubscriptionProtocolHandler.kt

@@ -16,8 +16,6 @@
 
 package com.expediagroup.graphql.server.spring.subscriptions
 
-import com.expediagroup.graphql.generator.execution.DefaultGraphQLContext
-import com.expediagroup.graphql.generator.execution.GraphQLContext
 import com.expediagroup.graphql.server.spring.GraphQLConfigurationProperties
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ClientMessages.GQL_CONNECTION_INIT
 import com.expediagroup.graphql.server.spring.subscriptions.SubscriptionOperationMessage.ClientMessages.GQL_CONNECTION_TERMINATE
@@ -33,6 +31,7 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.kotlin.convertValue
 import com.fasterxml.jackson.module.kotlin.readValue
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
 import org.slf4j.LoggerFactory
 import org.springframework.web.reactive.socket.WebSocketSession
 import reactor.core.publisher.Flux
@@ -66,7 +65,7 @@ class ApolloSubscriptionProtocolHandler(
         return try {
             when (operationMessage.type) {
                 GQL_CONNECTION_INIT.type -> onInit(operationMessage, session)
-                GQL_START.type -> onStart(operationMessage, session)
+                GQL_START.type -> startSubscription(operationMessage, session)
                 GQL_STOP.type -> onStop(operationMessage, session)
                 GQL_CONNECTION_TERMINATE.type -> onDisconnect(session)
                 else -> onUnknownOperation(operationMessage, session)
@@ -104,17 +103,18 @@ class ApolloSubscriptionProtocolHandler(
     @Suppress("Detekt.TooGenericExceptionCaught")
     private fun startSubscription(
         operationMessage: SubscriptionOperationMessage,
-        session: WebSocketSession,
-        context: GraphQLContext
+        session: WebSocketSession
     ): Flux<SubscriptionOperationMessage> {
+        val context = sessionState.getContext(session)
+
         subscriptionHooks.onOperation(operationMessage, session, context)
 
         if (operationMessage.id == null) {
             logger.error("GraphQL subscription operation id is required")
             return Flux.just(basicConnectionErrorMessage)
         }
 
-        if (sessionState.operationExists(session, operationMessage)) {
+        if (sessionState.doesOperationExist(session, operationMessage)) {
             logger.info("Already subscribed to operation ${operationMessage.id} for session ${session.id}")
             return Flux.empty()
         }
@@ -149,13 +149,23 @@ class ApolloSubscriptionProtocolHandler(
     }
 
     private fun onInit(operationMessage: SubscriptionOperationMessage, session: WebSocketSession): Flux<SubscriptionOperationMessage> {
-        val connectionParams = getConnectionParams(operationMessage.payload)
-        val graphQLContext = contextFactory.generateContext(session) ?: DefaultGraphQLContext()
-        val onConnect = subscriptionHooks.onConnect(connectionParams, session, graphQLContext)
-        sessionState.saveContext(session, onConnect)
-        val acknowledgeMessage = Flux.just(acknowledgeMessage)
+        saveContext(operationMessage, session)
+        val acknowledgeMessage = Mono.just(acknowledgeMessage)
         val keepAliveFlux = getKeepAliveFlux(session)
         return acknowledgeMessage.concatWith(keepAliveFlux)
+            .onErrorReturn(getConnectionErrorMessage(operationMessage))
+    }
+
+    /**
+     * Generate the context and save it for all future messages.
+     */
+    private fun saveContext(operationMessage: SubscriptionOperationMessage, session: WebSocketSession) {
+        runBlocking {
+            val connectionParams = getConnectionParams(operationMessage.payload)
+            val context = contextFactory.generateContext(session)
+            val onConnect = subscriptionHooks.onConnect(connectionParams, session, context)
+            sessionState.saveContext(session, onConnect)
+        }
     }
 
     /**
@@ -172,25 +182,6 @@ class ApolloSubscriptionProtocolHandler(
         return emptyMap()
     }
 
-    /**
-     * Called when the client sends the start message.
-     * It triggers the specific hooks first, runs the operation, and appends it with a complete message.
-     */
-    private fun onStart(
-        operationMessage: SubscriptionOperationMessage,
-        session: WebSocketSession
-    ): Flux<SubscriptionOperationMessage> {
-        val context = sessionState.getContext(session)
-
-        // If we do not have a context, that means the init message was never sent
-        return if (context != null) {
-            context.flatMapMany { startSubscription(operationMessage, session, it) }
-        } else {
-            val message = getConnectionErrorMessage(operationMessage)
-            Flux.just(message)
-        }
-    }
-
     /**
      * Called with the publisher has completed on its own.
      */
@@ -216,7 +207,7 @@ class ApolloSubscriptionProtocolHandler(
     private fun onDisconnect(session: WebSocketSession): Flux<SubscriptionOperationMessage> {
         subscriptionHooks.onDisconnect(session)
         sessionState.terminateSession(session)
-        return Flux.empty<SubscriptionOperationMessage>()
+        return Flux.empty()
     }
 
     private fun onUnknownOperation(operationMessage: SubscriptionOperationMessage, session: WebSocketSession): Flux<SubscriptionOperationMessage> {