fix: pass test, but Post and websocket not optimized

99designs · Nov 12, 2024 · 5091920 · 5091920
1 parent cd88c7b
commit 5091920
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 8 deletions.
diff --git a/graphql/executor/executor.go b/graphql/executor/executor.go
@@ -63,12 +63,23 @@ func (e *Executor) CreateOperationContext(
 	opCtx := e.pool.Get().(*graphql.OperationContext)
 	defer e.pool.Put(opCtx)
 
-	opCtx.DisableIntrospection = true
-	opCtx.RecoverFunc = e.recoverFunc
-	opCtx.ResolverMiddleware = e.ext.fieldMiddleware
-	opCtx.RootResolverMiddleware = e.ext.rootFieldMiddleware
-	opCtx.Stats.Read = params.ReadTime
-	opCtx.Stats.OperationStart = graphql.GetStartTime(ctx)
+	if params.Headers.Get("Connection") == "Upgrade" && params.Headers.Get("Upgrade") == "websocket" {
+		copyCtx := *opCtx
+		copyCtx.DisableIntrospection = true
+		copyCtx.RecoverFunc = e.recoverFunc
+		copyCtx.ResolverMiddleware = e.ext.fieldMiddleware
+		copyCtx.RootResolverMiddleware = e.ext.rootFieldMiddleware
+		copyCtx.Stats.Read = params.ReadTime
+		copyCtx.Stats.OperationStart = graphql.GetStartTime(ctx)
+		opCtx = &copyCtx
+	} else {
+		opCtx.DisableIntrospection = true
+		opCtx.RecoverFunc = e.recoverFunc
+		opCtx.ResolverMiddleware = e.ext.fieldMiddleware
+		opCtx.RootResolverMiddleware = e.ext.rootFieldMiddleware
+		opCtx.Stats.Read = params.ReadTime
+		opCtx.Stats.OperationStart = graphql.GetStartTime(ctx)
+	}
 
 	ctx = graphql.WithOperationContext(ctx, opCtx)
 

diff --git a/graphql/handler/transport/http_post.go b/graphql/handler/transport/http_post.go
@@ -56,6 +56,8 @@ func (h POST) Do(w http.ResponseWriter, r *http.Request, exec graphql.GraphExecu
 		Start: start,
 		End:   graphql.Now(),
 	}
+	params.Headers.Set("Connection", "Upgrade")
+	params.Headers.Set("Upgrade", "websocket")
 
 	bodyString, err := getRequestBody(r)
 	if err != nil {

diff --git a/graphql/handler/transport/websocket.go b/graphql/handler/transport/websocket.go
@@ -375,7 +375,13 @@ func (c *wsConnection) closeOnCancel(ctx context.Context) {
 
 func (c *wsConnection) subscribe(start time.Time, msg *message) {
 	ctx := graphql.StartOperationTrace(c.ctx)
-	var params *graphql.RawParams
+	params := &graphql.RawParams{
+		Headers: http.Header{
+			"Upgrade": []string{"websocket"},
+			"Connection": []string{"Upgrade"},
+		},
+	}
+
 	if err := jsonDecode(bytes.NewReader(msg.payload), &params); err != nil {
 		c.sendError(msg.id, &gqlerror.Error{Message: "invalid json"})
 		c.complete(msg.id)
@@ -386,7 +392,7 @@ func (c *wsConnection) subscribe(start time.Time, msg *message) {
 		Start: start,
 		End:   graphql.Now(),
 	}
-
+	
 	rc, err := c.exec.CreateOperationContext(ctx, params)
 	if err != nil {
 		resp := c.exec.DispatchError(graphql.WithOperationContext(ctx, rc), err)