Preserve websocket session in FallbackSubscriptionProtocolHandler fix #137

Author: oliemansm

src/main/java/graphql/servlet/internal/FallbackSubscriptionProtocolHandler.java

@@ -1,5 +1,7 @@
 package graphql.servlet.internal;
 
+import graphql.servlet.GraphQLSingleInvocationInput;
+
 import javax.websocket.Session;
 import javax.websocket.server.HandshakeRequest;
 import java.io.IOException;
@@ -20,18 +22,23 @@ public FallbackSubscriptionProtocolHandler(SubscriptionHandlerInput subscription
 
     @Override
     public void onMessage(HandshakeRequest request, Session session, WsSessionSubscriptions subscriptions, String text) throws Exception {
+        GraphQLSingleInvocationInput graphQLSingleInvocationInput = createInvocationInput(session, text);
         subscribe(
             session,
-            input.getQueryInvoker().query(
-                input.getInvocationInputFactory().create(
-                    input.getGraphQLObjectMapper().readGraphQLRequest(text)
-                )
-            ),
+            input.getQueryInvoker().query(graphQLSingleInvocationInput),
             subscriptions,
             UUID.randomUUID().toString()
         );
     }
 
+    private GraphQLSingleInvocationInput createInvocationInput(Session session, String text) throws IOException {
+        GraphQLRequest graphQLRequest = input.getGraphQLObjectMapper().readGraphQLRequest(text);
+        HandshakeRequest handshakeRequest = (HandshakeRequest) session.getUserProperties()
+                .get(HandshakeRequest.class.getName());
+
+        return input.getInvocationInputFactory().create(graphQLRequest, session, handshakeRequest);
+    }
+
     @Override
     protected void sendDataMessage(Session session, String id, Object payload) {
         sender.send(session, payload);