Added websockets support for ktor (#44)

martinbonnin · web-flow · commit b8c177b46a6a · 2024-11-26T10:05:04.000+01:00
* Added websockets support for ktor

* update docs

* add changelog entry
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 * Add `Ftv1Instrumentation` and `ApolloReportingInstrumentation` for respectively federated and monograph operation and fields insights.
 * Publish to GCS
+* Add `apolloSubscriptionModule` for subscription support with Ktor.
 
 # Version 0.1.0
 _2024-10-23_
diff --git a/Writerside/topics/ktor.md b/Writerside/topics/ktor.md
@@ -3,6 +3,8 @@
 
 > See [sample-ktor](https://github.com/apollographql/apollo-kotlin-execution/tree/main/sample-ktor) for a project using the Ktor integration. 
 
+## Add apollo-execution-ktor to your project
+
 To use the Ktor integration, add `apollo-execution-ktor` to your dependencies and a Ktor engine:
 
 ```kotlin
@@ -15,28 +17,23 @@ dependencies {
 }
 ```
 
-`apollo-execution-ktor` provides an `apolloModule(ExecutableSchema)` function that adds a `/graphql` route to your application:
-
-```kotlin
-embeddedServer(Netty, port = 8080) {
-  // /graphql route
-  apolloModule(ServiceExecutableSchemaBuilder().build())
-}.start(wait = true)
-```
-
-> `ServiceExecutableSchemaBuilder` is generated by the KSP processor. See ["Getting started"](getting-started.md) for more details.
+`apollo-execution-ktor` provides 3 modules:
 
-You can also opt in the Apollo Sandbox route by using `apolloSandboxModule()`
+- `apolloModule(ExecutableSchema)` adds the main `/graphql` route for queries/mutations.
+- `apolloSubscriptionModule(ExecutableSchema)` adds the `/subscription` route for subscriptions.
+- `apolloSandboxModule(ExecutableSchema)` adds the `/sandbox/index.html` for the online IDE
 
 ```kotlin
 embeddedServer(Netty, port = 8080) {
+  val executableSchema = ServiceExecutableSchemaBuilder().build()
+  // /graphql route
+  apolloModule(executableSchema)
+  // /subscription route
+  apolloSubscriptionModule(executableSchema)
   // /sandbox/index.html route
   apolloSandboxModule()
-}
+}.start(wait = true)
 ```
 
-`apolloSandboxModule()` adds a `sandbox/index.html` route to your application.
-
-Open [`http://localhost:8080/sandbox/index.html`](http://localhost:8080/sandbox/index.html) and try out your API in the [Apollo sandbox](https://www.apollographql.com/docs/graphos/explorer/sandbox/)
+> `ServiceExecutableSchemaBuilder` is generated by the KSP processor. See ["Getting started"](getting-started.md) for more details.
 
-[![Apollo Sandbox](sandbox.png)](http://localhost:8080/sandbox/index.html)
diff --git a/apollo-execution-ktor/api/apollo-execution-ktor.api b/apollo-execution-ktor/api/apollo-execution-ktor.api
@@ -3,6 +3,8 @@ public final class com/apollographql/execution/ktor/MainKt {
 	public static synthetic fun apolloModule$default (Lio/ktor/server/application/Application;Lcom/apollographql/execution/ExecutableSchema;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static final fun apolloSandboxModule (Lio/ktor/server/application/Application;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
 	public static synthetic fun apolloSandboxModule$default (Lio/ktor/server/application/Application;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;ILjava/lang/Object;)V
+	public static final fun apolloSubscriptionModule (Lio/ktor/server/application/Application;Lcom/apollographql/execution/ExecutableSchema;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V
+	public static synthetic fun apolloSubscriptionModule$default (Lio/ktor/server/application/Application;Lcom/apollographql/execution/ExecutableSchema;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V
 	public static final fun parseAsGraphQLRequest (Lio/ktor/server/request/ApplicationRequest;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun respondGraphQL (Lio/ktor/server/application/ApplicationCall;Lcom/apollographql/execution/ExecutableSchema;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static synthetic fun respondGraphQL$default (Lio/ktor/server/application/ApplicationCall;Lcom/apollographql/execution/ExecutableSchema;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
diff --git a/apollo-execution-ktor/build.gradle.kts b/apollo-execution-ktor/build.gradle.kts
@@ -16,6 +16,7 @@ kotlin {
         implementation(libs.atomicfu)
         api(libs.coroutines)
         api(libs.ktor.server.core)
+        implementation(libs.ktor.server.websockets)
       }
     }
     getByName("jvmTest") {
diff --git a/apollo-execution-ktor/src/commonMain/kotlin/com/apollographql/execution/ktor/main.kt b/apollo-execution-ktor/src/commonMain/kotlin/com/apollographql/execution/ktor/main.kt
@@ -2,14 +2,24 @@ package com.apollographql.execution.ktor
 
 import com.apollographql.apollo.api.ExecutionContext
 import com.apollographql.execution.*
+import com.apollographql.execution.websocket.ConnectionInitAck
+import com.apollographql.execution.websocket.SubscriptionWebSocketHandler
+import com.apollographql.execution.websocket.WebSocketBinaryMessage
+import com.apollographql.execution.websocket.WebSocketHandler
+import com.apollographql.execution.websocket.WebSocketTextMessage
 import io.ktor.http.*
 import io.ktor.http.content.*
 import io.ktor.server.application.*
 import io.ktor.server.request.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
 import io.ktor.server.util.*
+import io.ktor.server.websocket.WebSockets
+import io.ktor.server.websocket.webSocket
 import io.ktor.utils.io.*
+import io.ktor.websocket.Frame
+import io.ktor.websocket.readText
+import kotlinx.coroutines.coroutineScope
 import okio.Buffer
 
 suspend fun ApplicationCall.respondGraphQL(
@@ -86,6 +96,42 @@ fun Application.apolloModule(
   }
 }
 
+fun Application.apolloSubscriptionModule(
+  executableSchema: ExecutableSchema,
+  path: String = "/subscription",
+  executionContext: (ApplicationRequest) -> ExecutionContext = { ExecutionContext.Empty }
+) {
+  install(WebSockets)
+
+  routing {
+    webSocket(path) {
+      coroutineScope {
+        val handler = SubscriptionWebSocketHandler(
+          executableSchema = executableSchema,
+          scope = this,
+          executionContext = executionContext(call.request),
+          sendMessage = {
+            when (it) {
+              is WebSocketBinaryMessage -> send(Frame.Binary(true, it.data))
+              is WebSocketTextMessage -> send(Frame.Text(it.data))
+            }
+          },
+          connectionInitHandler = {
+            ConnectionInitAck
+          }
+        )
+
+        for (frame in incoming) {
+          if (frame !is Frame.Text) {
+            continue
+          }
+          handler.handleMessage(WebSocketTextMessage(frame.readText()))
+        }
+      }
+    }
+  }
+}
+
 fun Application.apolloSandboxModule(
   title: String = "API sandbox",
   sandboxPath: String = "/sandbox",
diff --git a/apollo-execution-runtime/api/apollo-execution-runtime.api b/apollo-execution-runtime/api/apollo-execution-runtime.api
@@ -279,8 +279,8 @@ public abstract interface class com/apollographql/execution/websocket/Connection
 }
 
 public final class com/apollographql/execution/websocket/SubscriptionWebSocketHandler : com/apollographql/execution/websocket/WebSocketHandler {
-	public fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)V
-	public synthetic fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;)V
+	public synthetic fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public final fun close ()V
 	public fun handleMessage (Lcom/apollographql/execution/websocket/WebSocketMessage;)V
 }
diff --git a/apollo-execution-runtime/api/apollo-execution-runtime.klib.api b/apollo-execution-runtime/api/apollo-execution-runtime.klib.api
@@ -44,7 +44,7 @@ final class com.apollographql.execution.websocket/ConnectionInitError : com.apol
         final fun <get-payload>(): com.apollographql.apollo.api/Optional<kotlin/Any?> // com.apollographql.execution.websocket/ConnectionInitError.payload.<get-payload>|<get-payload>(){}[0]
 }
 final class com.apollographql.execution.websocket/SubscriptionWebSocketHandler : com.apollographql.execution.websocket/WebSocketHandler { // com.apollographql.execution.websocket/SubscriptionWebSocketHandler|null[0]
-    constructor <init>(com.apollographql.execution/ExecutableSchema, kotlinx.coroutines/CoroutineScope, com.apollographql.apollo.api/ExecutionContext, kotlin/Function1<com.apollographql.execution.websocket/WebSocketMessage, kotlin/Unit>, kotlin.coroutines/SuspendFunction1<kotlin/Any?, com.apollographql.execution.websocket/ConnectionInitResult> = ...) // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.<init>|<init>(com.apollographql.execution.ExecutableSchema;kotlinx.coroutines.CoroutineScope;com.apollographql.apollo.api.ExecutionContext;kotlin.Function1<com.apollographql.execution.websocket.WebSocketMessage,kotlin.Unit>;kotlin.coroutines.SuspendFunction1<kotlin.Any?,com.apollographql.execution.websocket.ConnectionInitResult>){}[0]
+    constructor <init>(com.apollographql.execution/ExecutableSchema, kotlinx.coroutines/CoroutineScope, com.apollographql.apollo.api/ExecutionContext, kotlin.coroutines/SuspendFunction1<com.apollographql.execution.websocket/WebSocketMessage, kotlin/Unit>, kotlin.coroutines/SuspendFunction1<kotlin/Any?, com.apollographql.execution.websocket/ConnectionInitResult> = ...) // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.<init>|<init>(com.apollographql.execution.ExecutableSchema;kotlinx.coroutines.CoroutineScope;com.apollographql.apollo.api.ExecutionContext;kotlin.coroutines.SuspendFunction1<com.apollographql.execution.websocket.WebSocketMessage,kotlin.Unit>;kotlin.coroutines.SuspendFunction1<kotlin.Any?,com.apollographql.execution.websocket.ConnectionInitResult>){}[0]
     final fun close() // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.close|close(){}[0]
     final fun handleMessage(com.apollographql.execution.websocket/WebSocketMessage) // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.handleMessage|handleMessage(com.apollographql.execution.websocket.WebSocketMessage){}[0]
 }
diff --git a/apollo-execution-runtime/src/commonMain/kotlin/com/apollographql/execution/internal/OperationContext.kt b/apollo-execution-runtime/src/commonMain/kotlin/com/apollographql/execution/internal/OperationContext.kt
@@ -76,7 +76,7 @@ internal class OperationContext(
       "query" -> queryRoot?.resolveRoot()
       "mutation" -> mutationRoot?.resolveRoot()
       "subscription" -> {
-        return graphqlErrorResponse("Use executeSubscription() to execute subscriptions")
+        return graphqlErrorResponse("Use subscribe() to execute subscriptions")
       }
 
       else -> {
diff --git a/apollo-execution-runtime/src/commonMain/kotlin/com/apollographql/execution/websocket/SubscriptionWebSocketHandler.kt b/apollo-execution-runtime/src/commonMain/kotlin/com/apollographql/execution/websocket/SubscriptionWebSocketHandler.kt
@@ -32,7 +32,7 @@ class SubscriptionWebSocketHandler(
     private val executableSchema: ExecutableSchema,
     private val scope: CoroutineScope,
     private val executionContext: ExecutionContext,
-    private val sendMessage: (WebSocketMessage) -> Unit,
+    private val sendMessage: suspend (WebSocketMessage) -> Unit,
     private val connectionInitHandler: ConnectionInitHandler = { ConnectionInitAck },
 ) : WebSocketHandler {
   private val lock = reentrantLock()
@@ -67,7 +67,9 @@ class SubscriptionWebSocketHandler(
           activeSubscriptions.containsKey(clientMessage.id)
         }
         if (isActive) {
-          sendMessage(SubscriptionWebsocketError(id = clientMessage.id, error = Error.Builder("Subscription ${clientMessage.id} is already active").build()).toWsMessage())
+          scope.launch {
+            sendMessage(SubscriptionWebsocketError(id = clientMessage.id, error = Error.Builder("Subscription ${clientMessage.id} is already active").build()).toWsMessage())
+          }
           return
         }
 
@@ -107,7 +109,9 @@ class SubscriptionWebSocketHandler(
       }
 
       is SubscriptionWebsocketClientMessageParseError -> {
-        sendMessage(SubscriptionWebsocketError(null, Error.Builder("Cannot handle message (${clientMessage.message})").build()).toWsMessage())
+        scope.launch {
+          sendMessage(SubscriptionWebsocketError(null, Error.Builder("Cannot handle message (${clientMessage.message})").build()).toWsMessage())
+        }
       }
     }
   }
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -23,6 +23,7 @@ ktor-server-netty = { group = "io.ktor", name = "ktor-server-netty", version.ref
 ktor-server-core = { group = "io.ktor", name = "ktor-server-core", version.ref ="ktor" }
 ktor-server-cors = { group = "io.ktor", name = "ktor-server-cors", version.ref ="ktor" }
 ktor-client-cio = { module = "io.ktor:ktor-client-cio", version.ref = "ktor" }
+ktor-server-websockets = { group = "io.ktor", name = "ktor-server-websockets", version.ref ="ktor" }
 slf4j-simple = "org.slf4j:slf4j-simple:2.0.13"
 slf4j-nop = "org.slf4j:slf4j-nop:2.0.13"
 ksp-api = { module = "com.google.devtools.ksp:symbol-processing-api", version.ref = "ksp" }

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ kotlin {`
`16`	`16`	`implementation(libs.atomicfu)`
`17`	`17`	`api(libs.coroutines)`
`18`	`18`	`api(libs.ktor.server.core)`
	`19`	`+ implementation(libs.ktor.server.websockets)`
`19`	`20`	`}`
`20`	`21`	`}`
`21`	`22`	`getByName("jvmTest") {`
Original file line number	Diff line number	Diff line change
`@@ -279,8 +279,8 @@ public abstract interface class com/apollographql/execution/websocket/Connection`
`279`	`279`	`}`
`280`	`280`
`281`	`281`	`public final class com/apollographql/execution/websocket/SubscriptionWebSocketHandler : com/apollographql/execution/websocket/WebSocketHandler {`
`282`		`- public fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;)V`
`283`		`- public synthetic fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function1;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V`
	`282`	`+ public fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;)V`
	`283`	`+ public synthetic fun <init> (Lcom/apollographql/execution/ExecutableSchema;Lkotlinx/coroutines/CoroutineScope;Lcom/apollographql/apollo/api/ExecutionContext;Lkotlin/jvm/functions/Function2;Lkotlin/jvm/functions/Function2;ILkotlin/jvm/internal/DefaultConstructorMarker;)V`
`284`	`284`	`public final fun close ()V`
`285`	`285`	`public fun handleMessage (Lcom/apollographql/execution/websocket/WebSocketMessage;)V`
`286`	`286`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ final class com.apollographql.execution.websocket/ConnectionInitError : com.apol`
`44`	`44`	`final fun <get-payload>(): com.apollographql.apollo.api/Optional<kotlin/Any?> // com.apollographql.execution.websocket/ConnectionInitError.payload.<get-payload>\|<get-payload>(){}[0]`
`45`	`45`	`}`
`46`	`46`	`final class com.apollographql.execution.websocket/SubscriptionWebSocketHandler : com.apollographql.execution.websocket/WebSocketHandler { // com.apollographql.execution.websocket/SubscriptionWebSocketHandler\|null[0]`
`47`		- constructor <init>(com.apollographql.execution/ExecutableSchema, kotlinx.coroutines/CoroutineScope, com.apollographql.apollo.api/ExecutionContext, kotlin/Function1<com.apollographql.execution.websocket/WebSocketMessage, kotlin/Unit>, kotlin.coroutines/SuspendFunction1<kotlin/Any?, com.apollographql.execution.websocket/ConnectionInitResult> = ...) // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.<init>\|<init>(com.apollographql.execution.ExecutableSchema;kotlinx.coroutines.CoroutineScope;com.apollographql.apollo.api.ExecutionContext;kotlin.Function1<com.apollographql.execution.websocket.WebSocketMessage,kotlin.Unit>;kotlin.coroutines.SuspendFunction1<kotlin.Any?,com.apollographql.execution.websocket.ConnectionInitResult>){}[0]
	`47`	+ constructor <init>(com.apollographql.execution/ExecutableSchema, kotlinx.coroutines/CoroutineScope, com.apollographql.apollo.api/ExecutionContext, kotlin.coroutines/SuspendFunction1<com.apollographql.execution.websocket/WebSocketMessage, kotlin/Unit>, kotlin.coroutines/SuspendFunction1<kotlin/Any?, com.apollographql.execution.websocket/ConnectionInitResult> = ...) // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.<init>\|<init>(com.apollographql.execution.ExecutableSchema;kotlinx.coroutines.CoroutineScope;com.apollographql.apollo.api.ExecutionContext;kotlin.coroutines.SuspendFunction1<com.apollographql.execution.websocket.WebSocketMessage,kotlin.Unit>;kotlin.coroutines.SuspendFunction1<kotlin.Any?,com.apollographql.execution.websocket.ConnectionInitResult>){}[0]
`48`	`48`	`final fun close() // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.close\|close(){}[0]`
`49`	`49`	`final fun handleMessage(com.apollographql.execution.websocket/WebSocketMessage) // com.apollographql.execution.websocket/SubscriptionWebSocketHandler.handleMessage\|handleMessage(com.apollographql.execution.websocket.WebSocketMessage){}[0]`
`50`	`50`	`}`
Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,7 @@ internal class OperationContext(`
`76`	`76`	`"query" -> queryRoot?.resolveRoot()`
`77`	`77`	`"mutation" -> mutationRoot?.resolveRoot()`
`78`	`78`	`"subscription" -> {`
`79`		`- return graphqlErrorResponse("Use executeSubscription() to execute subscriptions")`
	`79`	`+ return graphqlErrorResponse("Use subscribe() to execute subscriptions")`
`80`	`80`	`}`
`81`	`81`
`82`	`82`	`else -> {`
Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ class SubscriptionWebSocketHandler(`
`32`	`32`	`private val executableSchema: ExecutableSchema,`
`33`	`33`	`private val scope: CoroutineScope,`
`34`	`34`	`private val executionContext: ExecutionContext,`
`35`		`- private val sendMessage: (WebSocketMessage) -> Unit,`
	`35`	`+ private val sendMessage: suspend (WebSocketMessage) -> Unit,`
`36`	`36`	`private val connectionInitHandler: ConnectionInitHandler = { ConnectionInitAck },`
`37`	`37`	`) : WebSocketHandler {`
`38`	`38`	`private val lock = reentrantLock()`
`@@ -67,7 +67,9 @@ class SubscriptionWebSocketHandler(`
`67`	`67`	`activeSubscriptions.containsKey(clientMessage.id)`
`68`	`68`	`}`
`69`	`69`	`if (isActive) {`
`70`		`- sendMessage(SubscriptionWebsocketError(id = clientMessage.id, error = Error.Builder("Subscription ${clientMessage.id} is already active").build()).toWsMessage())`
	`70`	`+ scope.launch {`
	`71`	`+ sendMessage(SubscriptionWebsocketError(id = clientMessage.id, error = Error.Builder("Subscription ${clientMessage.id} is already active").build()).toWsMessage())`
	`72`	`+ }`
`71`	`73`	`return`
`72`	`74`	`}`
`73`	`75`
`@@ -107,7 +109,9 @@ class SubscriptionWebSocketHandler(`
`107`	`109`	`}`
`108`	`110`
`109`	`111`	`is SubscriptionWebsocketClientMessageParseError -> {`
`110`		`- sendMessage(SubscriptionWebsocketError(null, Error.Builder("Cannot handle message (${clientMessage.message})").build()).toWsMessage())`
	`112`	`+ scope.launch {`
	`113`	`+ sendMessage(SubscriptionWebsocketError(null, Error.Builder("Cannot handle message (${clientMessage.message})").build()).toWsMessage())`
	`114`	`+ }`
`111`	`115`	`}`
`112`	`116`	`}`
`113`	`117`	`}`