Add Apollo Subscription Hooks (#544)

* Add Apollo Subscription Lifecycle Events (WIP):

* switch lifecycle events to monos

* adjust tests

* update tests and add state variable for onConnect

* add MySubscriptionHooks bean

* fix lint issues

* Update SubscriptionAutoConfiguration.kt

Change apolloLifeCycleEvents to apolloSubscriptionHooks

* add documentation and tests

* reorder test

* get rid of extra assertion in test

* alter tests to reflect that flux won't play twice

* fix lint

* fix lint

* default to using empty params in onConnect rather than empty mono

* switch some tests to using StepVerifier

Author: de55

docs/spring-server/subscriptions.md

@@ -7,15 +7,27 @@ title: Subscriptions
 To see more details of how to implement subscriptions in your schema, see [executing subscriptions](../execution/subscriptions).
 
 ## `graphql-ws` subprotocol
+###Overview
 We have implemented subscriptions in Spring WebSockets following the [`graphql-ws` subprotocol](https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md) defined by Apollo. This requires that your client send and parse messages in a specific format.
 
 You can see more details in the file [ApolloSubscriptionProtocolHandler](https://github.com/ExpediaGroup/graphql-kotlin/blob/master/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/spring/execution/ApolloSubscriptionProtocolHandler.kt).
 
 If you would like to implement your own subscription handler, you can provide a primary spring bean for `HandlerMapping` that overrides the [default one](https://github.com/ExpediaGroup/graphql-kotlin/blob/master/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/spring/SubscriptionAutoConfiguration.kt) which sets the url for subscriptions to the Apollo subscription handler.
 
+### Subscription Hooks
+In line with the protocol, we have implemented hooks to execute functions at different stages of the connection lifecycle:
+- onConnect
+- onOperation
+- onOperationComplete
+- onDisconnect
+
+You can see more details in the file [ApolloSubscriptionHooks](https://github.com/ExpediaGroup/graphql-kotlin/blob/master/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/spring/execution/ApolloSubscriptionHooks.kt).
+
+If you would like to implement your own subscription hooks, you can provide a primary spring bean for `ApolloSubscriptionHooks` that overrides the [default one](https://github.com/ExpediaGroup/graphql-kotlin/blob/master/graphql-kotlin-spring-server/src/main/kotlin/com/expediagroup/graphql/spring/SubscriptionAutoConfiguration.kt) which do not perform any actions.
+
+
 ## Example
 You can see an example implementation of a `Subscription` in the [example app](https://github.com/ExpediaGroup/graphql-kotlin/blob/master/examples/spring/src/main/kotlin/com/expediagroup/graphql/sample/subscriptions/SimpleSubscription.kt).
 
 
 
-

examples/spring/src/main/kotlin/com/expediagroup/graphql/examples/Application.kt

@@ -21,7 +21,9 @@ import com.expediagroup.graphql.examples.datafetchers.CustomDataFetcherFactoryPr
 import com.expediagroup.graphql.examples.datafetchers.SpringDataFetcherFactory
 import com.expediagroup.graphql.examples.directives.CustomDirectiveWiringFactory
 import com.expediagroup.graphql.examples.exceptions.CustomDataFetcherExceptionHandler
+import com.expediagroup.graphql.examples.execution.MySubscriptionHooks
 import com.expediagroup.graphql.examples.extension.CustomSchemaGeneratorHooks
+import com.expediagroup.graphql.spring.execution.ApolloSubscriptionHooks
 import com.fasterxml.jackson.databind.ObjectMapper
 import graphql.execution.DataFetcherExceptionHandler
 import org.springframework.boot.autoconfigure.SpringBootApplication
@@ -43,6 +45,9 @@ class Application {
 
     @Bean
     fun dataFetcherExceptionHandler(): DataFetcherExceptionHandler = CustomDataFetcherExceptionHandler()
+
+    @Bean
+    fun apolloSubscriptionHooks(): ApolloSubscriptionHooks = MySubscriptionHooks()
 }
 
 fun main(args: Array<String>) {

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

@@ -23,4 +23,4 @@ import org.springframework.http.server.reactive.ServerHttpResponse
 /**
  * Simple [GraphQLContext] that holds extra value.
  */
-class MyGraphQLContext(val myCustomValue: String, val request: ServerHttpRequest, val response: ServerHttpResponse)
+class MyGraphQLContext(val myCustomValue: String, val request: ServerHttpRequest, val response: ServerHttpResponse, var subscriptionValue: String? = null)

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

@@ -12,7 +12,8 @@ import org.springframework.stereotype.Component
 class MyGraphQLContextFactory : GraphQLContextFactory<MyGraphQLContext> {
 
     override suspend fun generateContext(request: ServerHttpRequest, response: ServerHttpResponse): MyGraphQLContext = MyGraphQLContext(
-            myCustomValue = request.headers.getFirst("MyHeader") ?: "defaultContext",
-            request = request,
-            response = response)
+        myCustomValue = request.headers.getFirst("MyHeader") ?: "defaultContext",
+        request = request,
+        response = response,
+        subscriptionValue = null)
 }

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

@@ -0,0 +1,23 @@
+package com.expediagroup.graphql.examples.execution
+
+import com.expediagroup.graphql.examples.context.MyGraphQLContext
+import com.expediagroup.graphql.spring.execution.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.
+ */
+open class MySubscriptionHooks :
+    ApolloSubscriptionHooks {
+    override fun onConnect(
+        connectionParams: Map<String, String>,
+        session: WebSocketSession,
+        graphQLContext: Any?
+    ): Mono<Unit> = mono {
+        val bearer = connectionParams["Authorization"] ?: "none"
+        val context = graphQLContext as? MyGraphQLContext
+        context?.subscriptionValue = bearer
+    }
+}

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

@@ -16,7 +16,9 @@
 
 package com.expediagroup.graphql.examples.subscriptions
 
+import com.expediagroup.graphql.annotations.GraphQLContext
 import com.expediagroup.graphql.annotations.GraphQLDescription
+import com.expediagroup.graphql.examples.context.MyGraphQLContext
 import com.expediagroup.graphql.spring.operations.Subscription
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.reactive.asPublisher
@@ -59,4 +61,8 @@ class SimpleSubscription : Subscription {
 
     @GraphQLDescription("Returns list of values")
     fun flow(): Publisher<Int> = flowOf(1, 2, 4).asPublisher()
+
+    @GraphQLDescription("Returns a value from the subscription context")
+    fun subscriptionContext(@GraphQLContext myGraphQLContext: MyGraphQLContext): Publisher<String> =
+        flowOf(myGraphQLContext.subscriptionValue ?: "", "value 2", "value3").asPublisher()
 }

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

@@ -16,8 +16,10 @@
 
 package com.expediagroup.graphql.spring
 
+import com.expediagroup.graphql.spring.execution.ApolloSubscriptionHooks
 import com.expediagroup.graphql.spring.execution.ApolloSubscriptionProtocolHandler
 import com.expediagroup.graphql.spring.execution.SimpleSubscriptionHandler
+import com.expediagroup.graphql.spring.execution.SimpleSubscriptionHooks
 import com.expediagroup.graphql.spring.execution.SubscriptionHandler
 import com.expediagroup.graphql.spring.execution.SubscriptionWebSocketHandler
 import com.expediagroup.graphql.spring.operations.Subscription
@@ -57,8 +59,16 @@ class SubscriptionAutoConfiguration {
     fun websocketHandlerAdapter(): WebSocketHandlerAdapter = WebSocketHandlerAdapter()
 
     @Bean
-    fun apolloSubscriptionProtocolHandler(config: GraphQLConfigurationProperties, handler: SubscriptionHandler, objectMapper: ObjectMapper) =
-        ApolloSubscriptionProtocolHandler(config, handler, objectMapper)
+    @ConditionalOnMissingBean
+    fun apolloSubscriptionHooks(): ApolloSubscriptionHooks = SimpleSubscriptionHooks()
+
+    @Bean
+    fun apolloSubscriptionProtocolHandler(
+        config: GraphQLConfigurationProperties,
+        handler: SubscriptionHandler,
+        objectMapper: ObjectMapper,
+        apolloSubscriptionHooks: ApolloSubscriptionHooks
+    ) = ApolloSubscriptionProtocolHandler(config, handler, objectMapper, apolloSubscriptionHooks)
 
     @Bean
     fun subscriptionWebSocketHandler(handler: ApolloSubscriptionProtocolHandler, objectMapper: ObjectMapper) =

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

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2019 Expedia, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.expediagroup.graphql.spring.execution
+
+import com.expediagroup.graphql.spring.model.SubscriptionOperationMessage
+import kotlinx.coroutines.reactor.mono
+import org.springframework.web.reactive.socket.WebSocketSession
+import reactor.core.publisher.Mono
+
+/**
+ * Implementation of Apollo Subscription Server Lifecycle Events
+ * https://www.apollographql.com/docs/graphql-subscriptions/lifecycle-events/
+ */
+interface ApolloSubscriptionHooks {
+    /**
+     * Allows validation of connectionParams prior to starting the connection.
+     * You can reject the connection by throwing an exception
+     */
+    fun onConnect(connectionParams: Map<String, String>, session: WebSocketSession, graphQLContext: Any?): Mono<Unit> = mono {}
+
+    /**
+     * Called when the client executes a GraphQL operation
+     */
+    fun onOperation(
+        operationMessage: SubscriptionOperationMessage,
+        session: WebSocketSession,
+        graphQLContext: Any?
+    ): Mono<Unit> = mono { }
+
+    /**
+     * Called when client's unsubscribes
+     */
+    fun onOperationComplete(session: WebSocketSession): Mono<Unit> = mono { }
+
+    /**
+     * Called when the client disconnects
+     */
+    fun onDisconnect(session: WebSocketSession, graphQLContext: Any?): Mono<Unit> = mono { }
+}
+
+/**
+ * Default implementation of Apollo Subscription Lifecycle Events.
+ */
+open class SimpleSubscriptionHooks : ApolloSubscriptionHooks

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

@@ -34,6 +34,7 @@ import com.fasterxml.jackson.module.kotlin.readValue
 import org.slf4j.LoggerFactory
 import org.springframework.web.reactive.socket.WebSocketSession
 import reactor.core.publisher.Flux
+import reactor.core.publisher.Mono
 import java.time.Duration
 
 /**
@@ -43,7 +44,8 @@ import java.time.Duration
 class ApolloSubscriptionProtocolHandler(
     private val config: GraphQLConfigurationProperties,
     private val subscriptionHandler: SubscriptionHandler,
-    private val objectMapper: ObjectMapper
+    private val objectMapper: ObjectMapper,
+    private val subscriptionHooks: ApolloSubscriptionHooks
 ) {
     private val sessionState = ApolloSubscriptionSessionState()
     private val logger = LoggerFactory.getLogger(ApolloSubscriptionProtocolHandler::class.java)
@@ -53,37 +55,39 @@ class ApolloSubscriptionProtocolHandler(
 
     @Suppress("Detekt.TooGenericExceptionCaught")
     fun handle(payload: String, session: WebSocketSession): Flux<SubscriptionOperationMessage> {
-        try {
-            val operationMessage: SubscriptionOperationMessage = objectMapper.readValue(payload)
-
-            logger.debug("GraphQL subscription client message, sessionId=${session.id} operationMessage=$operationMessage")
-
-            when (operationMessage.type) {
-                GQL_CONNECTION_INIT.type -> {
-                    val ackowledgeMessageFlux = Flux.just(acknowledgeMessage)
-                    val keepAliveFlux = getKeepAliveFlux(session)
-                    return ackowledgeMessageFlux.concatWith(keepAliveFlux)
-                }
-                GQL_START.type -> return startSubscription(operationMessage, session)
-                GQL_STOP.type -> return sessionState.stopOperation(session, operationMessage)
-                GQL_CONNECTION_TERMINATE.type -> {
-                    sessionState.terminateSession(session)
-                    return Flux.empty()
-                }
-                else -> {
-                    logger.error("Unknown subscription operation $operationMessage")
-                    sessionState.stopOperation(session, operationMessage)
-                    return Flux.just(SubscriptionOperationMessage(type = GQL_CONNECTION_ERROR.type, id = operationMessage.id))
+        val operationMessage = convertToMessageOrNull(payload) ?: return Flux.just(basicConnectionErrorMessage)
+        logger.debug("GraphQL subscription client message, sessionId=${session.id} operationMessage=$operationMessage")
+        return Mono.subscriberContext().flatMapMany<SubscriptionOperationMessage> { reactorContext ->
+            try {
+                val graphQLContext = reactorContext.getOrDefault<Any>(GRAPHQL_CONTEXT_KEY, null)
+                when (operationMessage.type) {
+                    GQL_CONNECTION_INIT.type -> onInit(operationMessage, session, graphQLContext)
+                    GQL_START.type -> onStart(operationMessage, session, graphQLContext)
+                    GQL_STOP.type -> onStop(operationMessage, session)
+                    GQL_CONNECTION_TERMINATE.type -> onDisconnect(session, graphQLContext)
+                    else -> {
+                        logger.error("Unknown subscription operation $operationMessage")
+                        sessionState.stopOperation(session, operationMessage)
+                        Flux.just(SubscriptionOperationMessage(type = GQL_CONNECTION_ERROR.type, id = operationMessage.id))
+                    }
                 }
+            } catch (exception: Exception) {
+                logger.error("Error parsing the subscription message", exception)
+                Flux.just(basicConnectionErrorMessage)
             }
+        }
+    }
+    @Suppress("Detekt.TooGenericExceptionCaught")
+    private fun convertToMessageOrNull(payload: String): SubscriptionOperationMessage? {
+        return try {
+            objectMapper.readValue(payload)
         } catch (exception: Exception) {
             logger.error("Error parsing the subscription message", exception)
-            return Flux.just(basicConnectionErrorMessage)
+            null
         }
     }
-
     /**
-     * If the keep alive configuraation is set, send a message back to client at every interval until the session is terminated.
+     * If the keep alive configuration is set, send a message back to client at every interval until the session is terminated.
      * Otherwise just return empty flux to append to the acknowledge message.
      */
     private fun getKeepAliveFlux(session: WebSocketSession): Flux<SubscriptionOperationMessage> {
@@ -98,7 +102,10 @@ class ApolloSubscriptionProtocolHandler(
     }
 
     @Suppress("Detekt.TooGenericExceptionCaught")
-    private fun startSubscription(operationMessage: SubscriptionOperationMessage, session: WebSocketSession): Flux<SubscriptionOperationMessage> {
+    private fun startSubscription(
+        operationMessage: SubscriptionOperationMessage,
+        session: WebSocketSession
+    ): Flux<SubscriptionOperationMessage> {
         if (operationMessage.id == null) {
             logger.error("GraphQL subscription operation id is required")
             return Flux.just(basicConnectionErrorMessage)
@@ -135,4 +142,40 @@ class ApolloSubscriptionProtocolHandler(
             return Flux.just(SubscriptionOperationMessage(type = GQL_CONNECTION_ERROR.type, id = operationMessage.id))
         }
     }
+
+    private fun onInit(operationMessage: SubscriptionOperationMessage, session: WebSocketSession, graphQLContext: Any?): Flux<SubscriptionOperationMessage> {
+        val connectionParams = operationMessage.payload as? Map<String, String> ?: emptyMap()
+        val onConnect = subscriptionHooks.onConnect(connectionParams, session, graphQLContext)
+        sessionState.saveOnConnectHook(session, onConnect)
+        val acknowledgeMessageFlux = Flux.just(acknowledgeMessage)
+        val keepAliveFlux = getKeepAliveFlux(session)
+        return acknowledgeMessageFlux.concatWith(keepAliveFlux)
+    }
+
+    private fun onStart(
+        operationMessage: SubscriptionOperationMessage,
+        session: WebSocketSession,
+        graphQLContext: Any?
+    ): Flux<SubscriptionOperationMessage> {
+        val onConnect = sessionState.onConnect(session) ?: subscriptionHooks.onConnect(emptyMap(), session, graphQLContext)
+        return onConnect.flatMap { subscriptionHooks.onOperation(operationMessage, session, graphQLContext) }
+            .flatMapMany { startSubscription(operationMessage, session) }
+    }
+
+    private fun onStop(
+        operationMessage: SubscriptionOperationMessage,
+        session: WebSocketSession
+    ): Flux<SubscriptionOperationMessage> = subscriptionHooks.onOperationComplete(session)
+        .flatMapMany {
+            sessionState.stopOperation(session, operationMessage)
+        }
+
+    private fun onDisconnect(
+        session: WebSocketSession,
+        graphQLContext: Any?
+    ): Flux<SubscriptionOperationMessage> = subscriptionHooks.onDisconnect(session, graphQLContext)
+        .flatMapMany {
+            sessionState.terminateSession(session)
+            Flux.empty<SubscriptionOperationMessage>()
+        }
 }