feat: implement apollographql/subscriptions-transport-ws protocol

Author: pankgeorg

src/GraphQLClient.jl

@@ -11,6 +11,7 @@ export query, mutate, open_subscription, Client, GQLEnum, Alias,
     create_introspected_struct, list_all_introspected_objects, global_graphql_client,
     @gql_str
 
+include("constants.jl")
 # Types
 include("client.jl")
 include("types.jl")

src/constants.jl

@@ -0,0 +1,19 @@
+# GQL over WS Protocol constants
+# https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md
+
+const GQL_CLIENT_CONNECTION_INIT = "connection_init"
+const GQL_SERVER_CONNECTION_ACK = "connection_ack"
+const GQL_SERVER_CONNECTION_ERROR = "connection_error"
+const GQL_SERVER_CONNECTION_KEEP_ALIVE = "ka"
+const GQL_CLIENT_START = "start"
+const GQL_CLIENT_STOP = "stop"
+const GQL_CLIENT_CONNECTION_TERMINATE = "connection_terminate"
+const GQL_SERVER_DATA = "data"
+const GQL_SERVER_ERROR = "error"
+const GQL_SERVER_COMPLETE = "complete"
+
+# Subscription tracker
+const SUBSCRIPTION_STATUS_OPEN = "open"
+const SUBSCRIPTION_STATUS_ERROR = "errored"
+const SUBSCRIPTION_STATUS_CLOSED = "closed"
+

src/gqlresponse.jl

@@ -99,9 +99,9 @@ end
 Struct for subsriptions that wraps a `GQLReponse{T}` alongside various metadata.
 """
 struct GQLSubscriptionResponse{T}
-    id::String
+    id::Union{String, Nothing}
     type::String
-    payload::GQLResponse{T}
+    payload::Union{GQLResponse{T}, Nothing}
 end
 StructTypes.StructType(::Type{<:GQLSubscriptionResponse}) = StructTypes.Struct()
 

src/subscriptions.jl

@@ -26,7 +26,7 @@ The subscription uses the `ws_endpoint` field of the `client.`
 This function is designed to be used with the `do` keyword.
 
 # Arguments
-- `fn::Function`: function to be run on each result, recieves the response from the
+- `fn::Function`: function to be run on each result, receives the response from the
     subscription`. Must return a boolean to indicate whether or not to close the subscription,
     with `true` closing the subscription.
 - `client::Client`: GraphQL client (optional). If not supplied, [`global_graphql_client`](@ref) is used.
@@ -87,61 +87,105 @@ function open_subscription(fn::Function,
 
     sub_id = string(length(keys(subscription_tracker[])) + 1)
     sub_id *= "-" * string(Threads.threadid())
-    message = Dict(
-        "id" => string(sub_id), 
-        "type" => "start", 
-        "payload" => payload
-    )
-    message_str = JSON3.write(message)
+
     throw_if_assigned = Ref{GraphQLError}()
-    HTTP.WebSockets.open(client.ws_endpoint; retry=retry, headers=client.headers, suppress_close_error=false) do ws
+    headers = Dict(client.headers)
+    # We currently implement the `apollographql/subscriptions-transport-ws` which is default in hasura and others
+    # Defined here https://github.com/apollographql/subscriptions-transport-ws/blob/master/PROTOCOL.md
+    # TODO: Add support for the newer `graphql-transport-ws` from the graphql-ws library.
+    # ()
+    headers["Sec-WebSocket-Protocol"] = "apollographql/subscriptions-transport-ws"
+
+    HTTP.WebSockets.open(client.ws_endpoint; retry=retry, headers=headers) do ws
         # Start sub
         output_info(verbose) && println("Starting $(get_name(subscription_name)) subscription with ID $sub_id")
-        HTTP.send(ws, message_str)
-        subscription_tracker[][sub_id] = "open"
-
+        HTTP.send(ws, JSON3.write(Dict("id" => sub_id, "type" => GQL_CLIENT_CONNECTION_INIT)))
         # Init function
         if !isnothing(initfn)
             output_debug(verbose) && println("Running subscription initialisation function")
             initfn()
         end
 
+        data = readfromwebsocket(ws, stopfn, subtimeout)
+        try checkreturn(data, verbose, sub_id)
+        catch e
+            e isa Interrupt && return
+        end
+        response = JSON3.read(data, GQLSubscriptionResponse{output_type})
+        while response.type == GQL_SERVER_CONNECTION_KEEP_ALIVE
+            data = readfromwebsocket(ws, stopfn, subtimeout)
+            try checkreturn(data, verbose, sub_id)
+            catch e
+                e isa Interrupt && return
+            end
+            response = JSON3.read(data, GQLSubscriptionResponse{output_type})
+        end
+        if response.type == GQL_SERVER_CONNECTION_ERROR && throw_on_execution_error
+            subscription_tracker[][sub_id] = SUBSCRIPTION_STATUS_ERROR
+            throw(GraphQLError("Error while establishing connection.", payload))
+        end
+
+        start_message = Dict(
+            "id" => string(sub_id), 
+            "type" => GQL_CLIENT_START, 
+            "payload" => payload
+        )
+        message_str = JSON3.write(start_message)
+        HTTP.send(ws, message_str)
+        subscription_tracker[][sub_id] = SUBSCRIPTION_STATUS_OPEN
+
         # Get listening
         output_debug(verbose) && println("Listening to $(get_name(subscription_name)) with ID $sub_id...")
 
         # Run function
-        finish = false
-        while !finish
+        while true
             data = readfromwebsocket(ws, stopfn, subtimeout)
-            if data === :timeout
-                output_info(verbose) && println("Subscription $sub_id timed out")
-                break
-            elseif data === :stopfn
-                output_info(verbose) && println("Subscription $sub_id stopped by the stop function supplied")
-                break
+            try checkreturn(data, verbose, sub_id)
+            catch e
+                e isa Interrupt && break
             end
+            # data = String(data)
+            # println(data)
             response = JSON3.read(data, GQLSubscriptionResponse{output_type})
+
+            response.type == GQL_SERVER_CONNECTION_KEEP_ALIVE && continue
+            response.type == GQL_SERVER_COMPLETE              && break
+            response.type == GQL_SERVER_CONNECTION_ERROR      && begin
+                throw_if_assigned[] = GraphQLError("Error during subscription. Server reporeted connection error")
+            end
+            response.type == GQL_SERVER_ERROR                 && begin
+                throw_if_assigned[] = GraphQLError("Error during subscription - GQL_SERVER_ERROR.", response.payload)
+            end
+            # response.type == GQL_SERVER_DATA
             payload = response.payload
             if !isnothing(payload.errors) && !isempty(payload.errors) && throw_on_execution_error
-                subscription_tracker[][sub_id] = "errored"
+                subscription_tracker[][sub_id] = SUBSCRIPTION_STATUS_ERROR
                 throw_if_assigned[] = GraphQLError("Error during subscription.", payload)
                 break
             end
             # Handle multiple subs, do we need this?
             if response.id == string(sub_id)
-                output_debug(verbose) && println("Result recieved on subscription with ID $sub_id")
+                output_debug(verbose) && println("Result received on subscription with ID $sub_id")
                 finish = fn(payload)
                 if !isa(finish, Bool)
-                    subscription_tracker[][sub_id] = "errored"
-                    error("Subscription function must return a boolean")
+                    subscription_tracker[][sub_id] = SUBSCRIPTION_STATUS_ERROR
+                    throw_if_assigned[] = ErrorException("Subscription function must return a boolean")
+                end
+                if finish
+                    # Protocol says we need to let the server know we're unsubscribing
+                    output_debug(verbose) && println("Finished. Closing subscription")
+                    HTTP.send(ws, JSON3.write(Dict("id" => sub_id, "type" => GQL_CLIENT_STOP)))
+                    HTTP.send(ws, JSON3.write(Dict("id" => sub_id, "type" => GQL_CLIENT_CONNECTION_TERMINATE)))
+                    close(ws)
+                    break
                 end
             end
         end
     end
-    # We can't throw errors from the ws handle function in HTTP.jl 1.0, as they get digested.
+    # We can't throw errors from the ws handle function in HTTP.jl 1.#, as they get digested.
     isassigned(throw_if_assigned) && throw(throw_if_assigned[])
     output_debug(verbose) && println("Finished. Closing subscription")
-    subscription_tracker[][sub_id] = "closed"
+    subscription_tracker[][sub_id] = SUBSCRIPTION_STATUS_CLOSED
     return
 end
 
@@ -158,7 +202,7 @@ function clear_subscriptions()
     end
 end
 
-function async_reader_with_timeout(ws::HTTP.WebSocket, subtimeout)::Channel
+function async_reader_with_timeout(ws::HTTP.WebSockets.WebSocket, subtimeout)::Channel
     ch = Channel(1)
     task = @async begin
         reader_task = current_task()
@@ -225,3 +269,15 @@ function readfromwebsocket(ws::HTTP.WebSockets.WebSocket, stopfn, subtimeout)
     end
     return data
 end
+struct Interrupt <: Exception
+end
+
+function checkreturn(data, verbose, sub_id)
+    if data === :timeout
+        output_info(verbose) && println("Subscription $sub_id timed out")
+        throw(Interrupt())
+    elseif data === :stopfn
+        output_info(verbose) && println("Subscription $sub_id stopped by the stop function supplied")
+        throw(Interrupt())
+    end
+end