add SUBSCRIPTION Directive

codegen/generated!.gotpl

@@ -166,7 +166,7 @@ func (e *executableSchema) Subscription(ctx context.Context, op *ast.OperationDe
 	{{- if .SubscriptionRoot }}
 		ec := executionContext{graphql.GetRequestContext(ctx), e}
 
-		{{ if .MutationDirectives -}}
+		{{ if .SubscriptionDirectives -}}
 			next := ec._{{.SubscriptionRoot.Name}}Middleware(ctx, op)
 		{{- else -}}
 			next := ec._{{.SubscriptionRoot.Name}}(ctx, op.SelectionSet)
@@ -247,7 +247,7 @@ func (ec *executionContext) _{{.QueryRoot.Name}}Middleware(ctx context.Context,
 {{end}}
 
 {{ if and .SubscriptionDirectives .SubscriptionRoot }}
-func (ec *executionContext) _{{.SubscriptionRoot.Name}}Middleware(ctx context.Context, obj *ast.OperationDefinition) graphql.Marshaler {
+func (ec *executionContext) _{{.SubscriptionRoot.Name}}Middleware(ctx context.Context, obj *ast.OperationDefinition) func() graphql.Marshaler {
 
 	next := func(ctx context.Context) (interface{}, error){
 		return ec._{{.SubscriptionRoot.Name}}(ctx, obj.SelectionSet),nil
@@ -261,7 +261,9 @@ func (ec *executionContext) _{{.SubscriptionRoot.Name}}Middleware(ctx context.Co
 				args, err := ec.{{ $directive.ArgsFunc }}(ctx,rawArgs)
 				if err != nil {
 					ec.Error(ctx, err)
-					return graphql.Null
+					return func() graphql.Marshaler {
+						return graphql.Null
+					}
 				}
 			{{- end }}
 			n := next
@@ -274,13 +276,17 @@ func (ec *executionContext) _{{.SubscriptionRoot.Name}}Middleware(ctx context.Co
 	tmp, err := next(ctx)
 	if err != nil {
 		ec.Error(ctx, err)
-		return graphql.Null
+		return func() graphql.Marshaler {
+			return graphql.Null
+		}
 	}
-	if data, ok := tmp.(graphql.Marshaler); ok {
+	if data, ok := tmp.(func() graphql.Marshaler); ok {
 		return data
 	}
 	ec.Errorf(ctx, `unexpected type %T from directive, should be graphql.Marshaler`, tmp)
-	return graphql.Null
+	return func() graphql.Marshaler {
+		return graphql.Null
+	}
 }
 {{end}}
 

example/chat/chat_test.go

@@ -15,15 +15,16 @@ func TestChatSubscriptions(t *testing.T) {
 	srv := httptest.NewServer(handler.GraphQL(NewExecutableSchema(New())))
 	c := client.New(srv.URL)
 
-	sub := c.Websocket(`subscription { messageAdded(roomName:"#gophers") { text createdBy } }`)
+	sub := c.Websocket(`subscription @user(username:"vektah") { messageAdded(roomName:"#gophers") { text createdBy } }`)
 	defer sub.Close()
 
 	go func() {
 		var resp interface{}
 		time.Sleep(10 * time.Millisecond)
 		err := c.Post(`mutation { 
 				a:post(text:"Hello!", roomName:"#gophers", username:"vektah") { id } 
-				b:post(text:"Whats up?", roomName:"#gophers", username:"vektah") { id } 
+				b:post(text:"Hello Vektah!", roomName:"#gophers", username:"andrey") { id } 
+				c:post(text:"Whats up?", roomName:"#gophers", username:"vektah") { id } 
 			}`, &resp)
 		assert.NoError(t, err)
 	}()

example/chat/resolvers.go

@@ -3,10 +3,12 @@
 package chat
 
 import (
-	context "context"
+	"context"
 	"math/rand"
 	"sync"
 	"time"
+
+	"github.com/99designs/gqlgen/graphql"
 )
 
 type resolver struct {
@@ -31,13 +33,28 @@ func New() Config {
 		Resolvers: &resolver{
 			Rooms: map[string]*Chatroom{},
 		},
+		Directives: DirectiveRoot{
+			User: func(ctx context.Context, obj interface{}, next graphql.Resolver, username string) (res interface{}, err error) {
+				return next(context.WithValue(ctx, "username", username))
+			},
+		},
+	}
+}
+
+func getUsername(ctx context.Context) string {
+	if username, ok := ctx.Value("username").(string); ok {
+		return username
 	}
+	return ""
 }
 
 type Chatroom struct {
 	Name      string
 	Messages  []Message
-	Observers map[string]chan *Message
+	Observers map[string]struct {
+		Username string
+		Message  chan *Message
+	}
 }
 
 type mutationResolver struct{ *resolver }
@@ -46,7 +63,13 @@ func (r *mutationResolver) Post(ctx context.Context, text string, username strin
 	r.mu.Lock()
 	room := r.Rooms[roomName]
 	if room == nil {
-		room = &Chatroom{Name: roomName, Observers: map[string]chan *Message{}}
+		room = &Chatroom{
+			Name: roomName,
+			Observers: map[string]struct {
+				Username string
+				Message  chan *Message
+			}{},
+		}
 		r.Rooms[roomName] = room
 	}
 	r.mu.Unlock()
@@ -61,7 +84,9 @@ func (r *mutationResolver) Post(ctx context.Context, text string, username strin
 	room.Messages = append(room.Messages, message)
 	r.mu.Lock()
 	for _, observer := range room.Observers {
-		observer <- &message
+		if observer.Username == "" || observer.Username == message.CreatedBy {
+			observer.Message <- &message
+		}
 	}
 	r.mu.Unlock()
 	return &message, nil
@@ -73,7 +98,13 @@ func (r *queryResolver) Room(ctx context.Context, name string) (*Chatroom, error
 	r.mu.Lock()
 	room := r.Rooms[name]
 	if room == nil {
-		room = &Chatroom{Name: name, Observers: map[string]chan *Message{}}
+		room = &Chatroom{
+			Name: name,
+			Observers: map[string]struct {
+				Username string
+				Message  chan *Message
+			}{},
+		}
 		r.Rooms[name] = room
 	}
 	r.mu.Unlock()
@@ -87,7 +118,13 @@ func (r *subscriptionResolver) MessageAdded(ctx context.Context, roomName string
 	r.mu.Lock()
 	room := r.Rooms[roomName]
 	if room == nil {
-		room = &Chatroom{Name: roomName, Observers: map[string]chan *Message{}}
+		room = &Chatroom{
+			Name: roomName,
+			Observers: map[string]struct {
+				Username string
+				Message  chan *Message
+			}{},
+		}
 		r.Rooms[roomName] = room
 	}
 	r.mu.Unlock()
@@ -103,7 +140,10 @@ func (r *subscriptionResolver) MessageAdded(ctx context.Context, roomName string
 	}()
 
 	r.mu.Lock()
-	room.Observers[id] = events
+	room.Observers[id] = struct {
+		Username string
+		Message  chan *Message
+	}{Username: getUsername(ctx), Message: events}
 	r.mu.Unlock()
 
 	return events, nil

example/chat/schema.graphql

@@ -23,3 +23,5 @@ type Subscription {
 }
 
 scalar Time
+
+directive @user(username: String!) on SUBSCRIPTION