feat!: onWebSocketInit gets access to Request and Context

This will give access to services (like a database or HTTPClient) in order to validate the init message and do things like auth checks.

Author: NeedleInAJayStack

Sources/GraphQLHummingbird/GraphQLConfig.swift

@@ -4,6 +4,7 @@ import Hummingbird
 
 /// Configuration options for GraphQLHummingbird
 public struct GraphQLConfig<
+    Context: RequestContext,
     WebSocketInit: Equatable & Codable & Sendable,
     WebSocketInitResult: Sendable
 >: Sendable {
@@ -31,7 +32,7 @@ public struct GraphQLConfig<
         subscriptionProtocols: Set<SubscriptionProtocol> = [.websocket],
         websocket: WebSocket = .init(
             // Including this strongly-typed argument is required to avoid compiler failures on Swift 6.2.3.
-            onWebSocketInit: { (_: EmptyWebSocketInit) in }
+            onWebSocketInit: { (_: EmptyWebSocketInit, _: Request, _: Context) in }
         ),
         additionalValidationRules: [@Sendable (ValidationContext) -> Visitor] = []
     ) {
@@ -97,14 +98,14 @@ public struct GraphQLConfig<
     }
 
     public struct WebSocket: Sendable {
-        let onWebSocketInit: @Sendable (WebSocketInit) async throws -> WebSocketInitResult
+        let onWebSocketInit: @Sendable (WebSocketInit, Request, Context) async throws -> WebSocketInitResult
 
         /// GraphQL over WebSocket configuration
         /// - Parameter onWebSocketInit: A custom callback run during `connection_init` resolution that allows
         /// authorization using the `payload` field of the `connection_init` message.
         /// Throw from this closure to indicate that authorization has failed.
         public init(
-            onWebSocketInit: @Sendable @escaping (WebSocketInit) async throws -> WebSocketInitResult = { (_: EmptyWebSocketInit) in }
+            onWebSocketInit: @Sendable @escaping (WebSocketInit, Request, Context) async throws -> WebSocketInitResult = { (_: EmptyWebSocketInit, _: Request, _: Context) in }
         ) {
             self.onWebSocketInit = onWebSocketInit
         }

Sources/GraphQLHummingbird/GraphQLHandler.swift

@@ -9,7 +9,7 @@ struct GraphQLHandler<
 >: Sendable {
     let schema: GraphQLSchema
     let rootValue: any Sendable
-    let config: GraphQLConfig<WebSocketInit, WebSocketInitResult>
+    let config: GraphQLConfig<Context, WebSocketInit, WebSocketInitResult>
     let computeContext: @Sendable (GraphQLContextComputationInputs<Context, WebSocketInitResult>) async throws -> GraphQLContext
 }
 

Sources/GraphQLHummingbird/Router+graphql.swift

@@ -19,7 +19,7 @@ public extension RouterMethods {
         _ path: RouterPath = "graphql",
         schema: GraphQLSchema,
         rootValue: any Sendable = (),
-        config: GraphQLConfig<EmptyWebSocketInit, Void> = .init(),
+        config: GraphQLConfig<Context, EmptyWebSocketInit, Void> = .init(),
         computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context, Void>) async throws -> GraphQLContext
     ) -> Self {
         // https://github.com/graphql/graphql-over-http/blob/main/spec/GraphQLOverHTTP.md#request
@@ -84,7 +84,7 @@ public extension RouterMethods where Context: WebSocketRequestContext {
         _ path: RouterPath = "graphql",
         schema: GraphQLSchema,
         rootValue: any Sendable = (),
-        config: GraphQLConfig<WebSocketInit, WebSocketInitResult> = GraphQLConfig<EmptyWebSocketInit, Void>(),
+        config: GraphQLConfig<Context, WebSocketInit, WebSocketInitResult>,
         computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context, WebSocketInitResult>) async throws -> GraphQLContext
     ) -> Self {
         let handler = GraphQLHandler<Context, GraphQLContext, WebSocketInit, WebSocketInitResult>(
@@ -108,4 +108,34 @@ public extension RouterMethods where Context: WebSocketRequestContext {
         }
         return self
     }
+
+    /// Registers a graphql websocket route that responds using the provided schema.
+    ///
+    /// WebSocket requests support the
+    /// [`graphql-transport-ws`](https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md)
+    /// and [`graphql-ws`](https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md)
+    /// subprotocols.
+    ///
+    /// - Parameters:
+    ///   - path: The route that should respond to GraphQL requests. Both `GET` and `POST` routes are registered.
+    ///   - schema: The GraphQL schema that should be used to respond to requests.
+    ///   - rootValue: The `rootValue` GraphQL execution arg. This is the object passed to the root resolvers.
+    ///   - computeContext: A closure used to compute the GraphQL context from incoming requests. This must be provided.
+    @discardableResult
+    func graphqlWebSocket<GraphQLContext: Sendable>(
+        _ path: RouterPath = "graphql",
+        schema: GraphQLSchema,
+        rootValue: any Sendable = (),
+        computeContext: @Sendable @escaping (GraphQLContextComputationInputs<Context, Void>) async throws -> GraphQLContext
+    ) -> Self {
+        // This is just an overload that allows not passing `config`, since we cannot use a default argument that
+        // uses `Context` without getting `error: generic parameter 'Self' could not be inferred` compilation errors
+        graphqlWebSocket(
+            path,
+            schema: schema,
+            rootValue: rootValue,
+            config: GraphQLConfig<Context, EmptyWebSocketInit, Void>(),
+            computeContext: computeContext
+        )
+    }
 }

Sources/GraphQLHummingbird/WebSocket/GraphQLHandler+handleWebSocket.swift

@@ -32,7 +32,7 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
             let server = GraphQLTransportWS.Server<WebSocketInit, WebSocketInitResult, AsyncThrowingStream<GraphQLResult, Error>>(
                 messenger: messenger,
                 onInit: { initPayload in
-                    try await config.websocket.onWebSocketInit(initPayload)
+                    try await config.websocket.onWebSocketInit(initPayload, context.request, context.requestContext)
                 },
                 onExecute: { graphQLRequest, initResult in
                     let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(
@@ -75,7 +75,7 @@ extension GraphQLHandler where Context: WebSocketRequestContext {
             let server = GraphQLWS.Server<WebSocketInit, WebSocketInitResult, AsyncThrowingStream<GraphQLResult, Error>>(
                 messenger: messenger,
                 onInit: { initPayload in
-                    try await config.websocket.onWebSocketInit(initPayload)
+                    try await config.websocket.onWebSocketInit(initPayload, context.request, context.requestContext)
                 },
                 onExecute: { graphQLRequest, initResult in
                     let graphQLContextComputationInputs = GraphQLContextComputationInputs<Context, WebSocketInitResult>(