diff --git a/Sources/GraphQLWS/Server.swift b/Sources/GraphQLWS/Server.swift index b043ed3..fb9ea17 100644 --- a/Sources/GraphQLWS/Server.swift +++ b/Sources/GraphQLWS/Server.swift @@ -19,6 +19,8 @@ public class Server { var auth: (InitPayload) throws -> Void = { _ in } var onExit: () -> Void = { } var onMessage: (String) -> Void = { _ in } + var onOperationComplete: (String) -> Void = { _ in } + var onOperationError: (String) -> Void = { _ in } var initialized = false @@ -66,6 +68,7 @@ public class Server { return } + // handle incoming message switch request.type { case .GQL_CONNECTION_INIT: guard let connectionInitRequest = try? self.decoder.decode(ConnectionInitRequest.self, from: json) else { @@ -84,7 +87,7 @@ public class Server { self.error(.invalidRequestFormat(messageType: .GQL_STOP)) return } - self.onStop(stopRequest, messenger) + self.onOperationComplete(stopRequest.id) case .GQL_CONNECTION_TERMINATE: guard let connectionTerminateRequest = try? self.decoder.decode(ConnectionTerminateRequest.self, from: json) else { self.error(.invalidRequestFormat(messageType: .GQL_CONNECTION_TERMINATE)) @@ -116,6 +119,18 @@ public class Server { self.onMessage = callback } + /// Define the callback run on the completion a full operation (query/mutation, end of subscription) + /// - Parameter callback: The callback to assign + public func onOperationComplete(_ callback: @escaping (String) -> Void) { + self.onOperationComplete = callback + } + + /// Define the callback to run on error of any full operation (failed query, interrupted subscription) + /// - Parameter callback: The callback to assign + public func onOperationError(_ callback: @escaping (String) -> Void) { + self.onOperationError = callback + } + private func onConnectionInit(_ connectionInitRequest: ConnectionInitRequest, _ messenger: Messenger) { guard !initialized else { self.error(.tooManyInitializations()) @@ -201,13 +216,6 @@ public class Server { } } - private func onStop(_: StopRequest, _ messenger: Messenger) { - guard initialized else { - self.error(.notInitialized()) - return - } - } - private func onConnectionTerminate(_: ConnectionTerminateRequest, _ messenger: Messenger) { onExit() _ = messenger.close() @@ -256,6 +264,7 @@ public class Server { id: id ).toJSON(encoder) ) + onOperationComplete(id) } /// Send an `error` response through the messenger @@ -267,6 +276,7 @@ public class Server { id: id ).toJSON(encoder) ) + onOperationError(id) } /// Send an `error` response through the messenger