feat!: Passes init payload result

This automatically propagates the init payload result from the callback into the `onExecute` and `onSubscribe` closures. Since the init callback is usually used to determine authentication and authorization, this should be usable from our execution and subscription calls, and lifecycle is most easily managed within this package

Author: NeedleInAJayStack

Sources/GraphQLTransportWS/Server.swift

@@ -6,23 +6,25 @@ import GraphQL
 /// By default, there are no authorization checks
 public class Server<
     InitPayload: Equatable & Codable & Sendable,
+    InitPayloadResult: Sendable,
     SubscriptionSequenceType: AsyncSequence & Sendable
 >: @unchecked Sendable where
     SubscriptionSequenceType.Element == GraphQLResult
 {
     // We keep this weak because we strongly inject this object into the messenger callback
     weak var messenger: Messenger?
-
-    let onExecute: (GraphQLRequest) async throws -> GraphQLResult
-    let onSubscribe: (GraphQLRequest) async throws -> SubscriptionSequenceType
-    var auth: (InitPayload) async throws -> Void
+    
+    let onInit: (InitPayload) async throws -> InitPayloadResult
+    let onExecute: (GraphQLRequest, InitPayloadResult) async throws -> GraphQLResult
+    let onSubscribe: (GraphQLRequest, InitPayloadResult) async throws -> SubscriptionSequenceType
 
     var onExit: () async throws -> Void = {}
     var onMessage: (String) async throws -> Void = { _ in }
     var onOperationComplete: (String) async throws -> Void = { _ in }
     var onOperationError: (String, [Error]) async throws -> Void = { _, _ in }
 
     var initialized = false
+    var initResult: InitPayloadResult?
 
     let decoder = JSONDecoder()
     let encoder = GraphQLJSONEncoder()
@@ -37,13 +39,14 @@ public class Server<
     ///   - onSubscribe: Callback run during `start` resolution for streaming queries. Typically this is `API.subscribe`.
     public init(
         messenger: Messenger,
-        onExecute: @escaping (GraphQLRequest) async throws -> GraphQLResult,
-        onSubscribe: @escaping (GraphQLRequest) async throws -> SubscriptionSequenceType
+        onInit: @escaping (InitPayload) async throws -> InitPayloadResult,
+        onExecute: @escaping (GraphQLRequest, InitPayloadResult) async throws -> GraphQLResult,
+        onSubscribe: @escaping (GraphQLRequest, InitPayloadResult) async throws -> SubscriptionSequenceType
     ) {
         self.messenger = messenger
+        self.onInit = onInit
         self.onExecute = onExecute
         self.onSubscribe = onSubscribe
-        auth = { _ in }
 
         messenger.onReceive { message in
             guard let messenger = self.messenger else { return }
@@ -99,13 +102,6 @@ public class Server<
         subscriptionTasks.values.forEach { $0.cancel() }
     }
 
-    /// Define a custom callback run during `connection_init` resolution that allows authorization using the `payload`.
-    /// Throw from this closure to indicate that authorization has failed.
-    /// - Parameter callback: The callback to assign
-    public func auth(_ callback: @escaping (InitPayload) async throws -> Void) {
-        auth = callback
-    }
-
     /// Define the callback run when the communication is shut down, either by the client or server
     /// - Parameter callback: The callback to assign
     public func onExit(_ callback: @escaping () -> Void) {
@@ -137,7 +133,7 @@ public class Server<
         }
 
         do {
-            try await auth(connectionInitRequest.payload)
+            initResult = try await onInit(connectionInitRequest.payload)
         } catch {
             try await self.error(.unauthorized())
             return
@@ -148,7 +144,7 @@ public class Server<
     }
 
     private func onSubscribe(_ subscribeRequest: SubscribeRequest) async throws {
-        guard initialized else {
+        guard initialized, let initResult else {
             try await error(.notInitialized())
             return
         }
@@ -171,7 +167,7 @@ public class Server<
         if isStreaming {
             subscriptionTasks[id] = Task {
                 do {
-                    let stream = try await onSubscribe(graphQLRequest)
+                    let stream = try await onSubscribe(graphQLRequest, initResult)
                     for try await event in stream {
                         try Task.checkCancellation()
                         try await self.sendNext(event, id: id)
@@ -186,7 +182,7 @@ public class Server<
             }
         } else {
             do {
-                let result = try await onExecute(graphQLRequest)
+                let result = try await onExecute(graphQLRequest, initResult)
                 try await sendNext(result, id: id)
                 try await sendComplete(id: id)
             } catch {

Tests/GraphQLTransportWSTests/GraphQLTransportWSTests.swift

@@ -8,42 +8,40 @@ import GraphQLTransportWS
 class GraphqlTransportWSTests: XCTestCase {
     var clientMessenger: TestMessenger!
     var serverMessenger: TestMessenger!
-    var server: Server<TokenInitPayload, AsyncThrowingStream<GraphQLResult, Error>>!
-    var context: TestContext!
     var subscribeReady: Bool! = false
+    
+    let context = TestContext()
+    let api = TestAPI()
 
     override func setUp() {
         // Point the client and server at each other
         clientMessenger = TestMessenger()
         serverMessenger = TestMessenger()
         clientMessenger.other = serverMessenger
         serverMessenger.other = clientMessenger
+    }
 
-        let api = TestAPI()
-        let context = TestContext()
-
-        server = .init(
+    /// Tests that trying to run methods before `connection_init` is not allowed
+    func testInitialize() async throws {
+        _ = Server<TokenInitPayload, Void, AsyncThrowingStream<GraphQLResult, Error>>(
             messenger: serverMessenger,
-            onExecute: { graphQLRequest in
-                try await api.execute(
+            onInit: { _ in },
+            onExecute: { graphQLRequest, _ in
+                try await self.api.execute(
                     request: graphQLRequest.query,
-                    context: context
+                    context: self.context
                 )
             },
-            onSubscribe: { graphQLRequest in
-                let subscription = try await api.subscribe(
+            onSubscribe: { graphQLRequest, _ in
+                let subscription = try await self.api.subscribe(
                     request: graphQLRequest.query,
-                    context: context
+                    context: self.context
                 ).get()
                 self.subscribeReady = true
                 return subscription
             }
         )
-        self.context = context
-    }
-
-    /// Tests that trying to run methods before `connection_init` is not allowed
-    func testInitialize() async throws {
+        
         let client = Client<TokenInitPayload>(messenger: clientMessenger)
         let messageStream = AsyncThrowingStream<String, any Error> { continuation in
             client.onMessage { message, _ in
@@ -78,9 +76,26 @@ class GraphqlTransportWSTests: XCTestCase {
 
     /// Tests that throwing in the authorization callback forces an unauthorized error
     func testAuthWithThrow() async throws {
-        server.auth { _ in
-            throw TestError.couldBeAnything
-        }
+        _ = Server<TokenInitPayload, Void, AsyncThrowingStream<GraphQLResult, Error>>(
+            messenger: serverMessenger,
+            onInit: { _ in
+                throw TestError.couldBeAnything
+            },
+            onExecute: { graphQLRequest, _ in
+                try await self.api.execute(
+                    request: graphQLRequest.query,
+                    context: self.context
+                )
+            },
+            onSubscribe: { graphQLRequest, _ in
+                let subscription = try await self.api.subscribe(
+                    request: graphQLRequest.query,
+                    context: self.context
+                ).get()
+                self.subscribeReady = true
+                return subscription
+            }
+        )
 
         let client = Client<TokenInitPayload>(messenger: clientMessenger)
         let messageStream = AsyncThrowingStream<String, any Error> { continuation in
@@ -111,6 +126,25 @@ class GraphqlTransportWSTests: XCTestCase {
 
     /// Tests a single-op conversation
     func testSingleOp() async throws {
+        _ = Server<TokenInitPayload, Void, AsyncThrowingStream<GraphQLResult, Error>>(
+            messenger: serverMessenger,
+            onInit: { _ in },
+            onExecute: { graphQLRequest, _ in
+                try await self.api.execute(
+                    request: graphQLRequest.query,
+                    context: self.context
+                )
+            },
+            onSubscribe: { graphQLRequest, _ in
+                let subscription = try await self.api.subscribe(
+                    request: graphQLRequest.query,
+                    context: self.context
+                ).get()
+                self.subscribeReady = true
+                return subscription
+            }
+        )
+        
         let id = UUID().description
 
         let client = Client<TokenInitPayload>(messenger: clientMessenger)
@@ -156,6 +190,25 @@ class GraphqlTransportWSTests: XCTestCase {
 
     /// Tests a streaming conversation
     func testStreaming() async throws {
+        _ = Server<TokenInitPayload, Void, AsyncThrowingStream<GraphQLResult, Error>>(
+            messenger: serverMessenger,
+            onInit: { _ in },
+            onExecute: { graphQLRequest, _ in
+                try await self.api.execute(
+                    request: graphQLRequest.query,
+                    context: self.context
+                )
+            },
+            onSubscribe: { graphQLRequest, _ in
+                let subscription = try await self.api.subscribe(
+                    request: graphQLRequest.query,
+                    context: self.context
+                ).get()
+                self.subscribeReady = true
+                return subscription
+            }
+        )
+        
         let id = UUID().description
 
         var dataIndex = 1