Enable Subscription Resolver to return websocket error message (#2506)

* Enanble Subscription Resolver to return websocket error message

* add PR link

* lint

* fmt and regenerate

Signed-off-by: Steve Coffman <steve@khanacademy.org>

Signed-off-by: Steve Coffman <steve@khanacademy.org>
Co-authored-by: Zhixin Wen <zwen@nuro.ai>
Co-authored-by: Steve Coffman <steve@khanacademy.org>

_examples/chat/chat_test.go

@@ -2,13 +2,14 @@ package chat
 
 import (
 	"fmt"
+	"runtime"
+	"sync"
+	"testing"
+
 	"github.com/99designs/gqlgen/client"
 	"github.com/99designs/gqlgen/graphql/handler"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"runtime"
-	"sync"
-	"testing"
 )
 
 func TestChatSubscriptions(t *testing.T) {

graphql/handler/transport/websocket.go

@@ -350,6 +350,7 @@ func (c *wsConnection) subscribe(start time.Time, msg *message) {
 	c.mu.Unlock()
 
 	go func() {
+		ctx = withSubscriptionErrorContext(ctx)
 		defer func() {
 			if r := recover(); r != nil {
 				err := rc.Recover(ctx, r)
@@ -362,7 +363,11 @@ func (c *wsConnection) subscribe(start time.Time, msg *message) {
 				}
 				c.sendError(msg.id, gqlerr)
 			}
-			c.complete(msg.id)
+			if errs := getSubscriptionError(ctx); len(errs) != 0 {
+				c.sendError(msg.id, errs...)
+			} else {
+				c.complete(msg.id)
+			}
 			c.mu.Lock()
 			delete(c.active, msg.id)
 			c.mu.Unlock()

graphql/handler/transport/websocket_resolver_error.go

@@ -0,0 +1,69 @@
+package transport
+
+import (
+	"context"
+
+	"github.com/vektah/gqlparser/v2/gqlerror"
+)
+
+// A private key for context that only this package can access. This is important
+// to prevent collisions between different context uses
+var wsSubscriptionErrorCtxKey = &wsSubscriptionErrorContextKey{"subscription-error"}
+
+type wsSubscriptionErrorContextKey struct {
+	name string
+}
+
+type subscriptionError struct {
+	errs []*gqlerror.Error
+}
+
+// AddSubscriptionError is used to let websocket return an error message after subscription resolver returns a channel.
+// for example:
+//
+//	func (r *subscriptionResolver) Method(ctx context.Context) (<-chan *model.Message, error) {
+//		ch := make(chan *model.Message)
+//		go func() {
+//	     defer func() {
+//				close(ch)
+//	     }
+//			// some kind of block processing (e.g.: gRPC client streaming)
+//			stream, err := gRPCClientStreamRequest(ctx)
+//			if err != nil {
+//				   transport.AddSubscriptionError(ctx, err)
+//	            return // must return and close channel so websocket can send error back
+//	     }
+//			for {
+//				m, err := stream.Recv()
+//				if err == io.EOF {
+//					return
+//				}
+//				if err != nil {
+//				   transport.AddSubscriptionError(ctx, err)
+//	            return // must return and close channel so websocket can send error back
+//				}
+//				ch <- m
+//			}
+//		}()
+//
+//		return ch, nil
+//	}
+//
+// see https://github.com/99designs/gqlgen/pull/2506 for more details
+func AddSubscriptionError(ctx context.Context, err *gqlerror.Error) {
+	subscriptionErrStruct := getSubscriptionErrorStruct(ctx)
+	subscriptionErrStruct.errs = append(subscriptionErrStruct.errs, err)
+}
+
+func withSubscriptionErrorContext(ctx context.Context) context.Context {
+	return context.WithValue(ctx, wsSubscriptionErrorCtxKey, &subscriptionError{})
+}
+
+func getSubscriptionErrorStruct(ctx context.Context) *subscriptionError {
+	v, _ := ctx.Value(wsSubscriptionErrorCtxKey).(*subscriptionError)
+	return v
+}
+
+func getSubscriptionError(ctx context.Context) []*gqlerror.Error {
+	return getSubscriptionErrorStruct(ctx).errs
+}

plugin/federation/testdata/entityresolver/generated/errors.go

@@ -3,7 +3,6 @@ package generated
 import "errors"
 
 // Errors defined for retained code that we want to stick around between generations.
-//
 var (
 	ErrResolvingHelloWithErrorsByName         = errors.New("error resolving HelloWithErrorsByName")
 	ErrEmptyKeyResolvingHelloWithErrorsByName = errors.New("error (empty key) resolving HelloWithErrorsByName")

plugin/resolvergen/testdata/return_values/return_values_test.go

@@ -1,9 +1,10 @@
 package return_values
 
 import (
-	"github.com/stretchr/testify/require"
 	"reflect"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 //go:generate rm -f resolvers.go