added subscribe support

Author: branden-blackline

definition.go

@@ -534,6 +534,7 @@ func defineFieldMap(ttype Named, fieldMap Fields) (FieldDefinitionMap, error) {
 			Description:       field.Description,
 			Type:              field.Type,
 			Resolve:           field.Resolve,
+			Subscribe:         field.Subscribe,
 			DeprecationReason: field.DeprecationReason,
 		}
 
@@ -606,6 +607,7 @@ type Field struct {
 	Type              Output              `json:"type"`
 	Args              FieldConfigArgument `json:"args"`
 	Resolve           FieldResolveFn      `json:"-"`
+	Subscribe         FieldResolveFn      `json:"-"`
 	DeprecationReason string              `json:"deprecationReason"`
 	Description       string              `json:"description"`
 }
@@ -625,6 +627,7 @@ type FieldDefinition struct {
 	Type              Output         `json:"type"`
 	Args              []*Argument    `json:"args"`
 	Resolve           FieldResolveFn `json:"-"`
+	Subscribe         FieldResolveFn `json:"-"`
 	DeprecationReason string         `json:"deprecationReason"`
 }
 

subscription.go

@@ -0,0 +1,217 @@
+package graphql
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/graphql-go/graphql/gqlerrors"
+	"github.com/graphql-go/graphql/language/ast"
+)
+
+type ResultIteratorFn func(count int64, result *Result, doneFunc func())
+
+type ResultIterator struct {
+	count      int64
+	ctx        context.Context
+	ch         chan *Result
+	cancelFunc context.CancelFunc
+	cancelled  bool
+	handlers   []ResultIteratorFn
+}
+
+func NewResultIterator(ctx context.Context, ch chan *Result) *ResultIterator {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	cctx, cancelFunc := context.WithCancel(ctx)
+	iterator := &ResultIterator{
+		count:      0,
+		ctx:        cctx,
+		ch:         ch,
+		cancelFunc: cancelFunc,
+		cancelled:  false,
+		handlers:   []ResultIteratorFn{},
+	}
+
+	go func() {
+		for {
+			select {
+			case <-iterator.ctx.Done():
+				return
+			case res := <-iterator.ch:
+				if iterator.cancelled {
+					return
+				}
+				iterator.count += 1
+				for _, handler := range iterator.handlers {
+					handler(iterator.count, res, iterator.Done)
+				}
+			}
+		}
+	}()
+
+	return iterator
+}
+
+func (c *ResultIterator) ForEach(handler ResultIteratorFn) {
+	c.handlers = append(c.handlers, handler)
+}
+
+func (c *ResultIterator) Done() {
+	c.cancelled = true
+	c.cancelFunc()
+}
+
+type SubscribeParams struct {
+	Schema          Schema
+	Document        *ast.Document
+	RootValue       interface{}
+	ContextValue    context.Context
+	VariableValues  map[string]interface{}
+	OperationName   string
+	FieldResolver   FieldResolveFn
+	FieldSubscriber FieldResolveFn
+}
+
+// Subscribe performs a subscribe operation
+func Subscribe(p SubscribeParams) *ResultIterator {
+	resultChannel := make(chan *Result)
+	// Use background context if no context was provided
+	ctx := p.ContextValue
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	var mapSourceToResponse = func(payload interface{}) *Result {
+		return Execute(ExecuteParams{
+			Schema:        p.Schema,
+			Root:          payload,
+			AST:           p.Document,
+			OperationName: p.OperationName,
+			Args:          p.VariableValues,
+			Context:       p.ContextValue,
+		})
+	}
+
+	go func() {
+
+		result := &Result{}
+		defer func() {
+			if err := recover(); err != nil {
+				result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
+			}
+			resultChannel <- result
+		}()
+
+		exeContext, err := buildExecutionContext(buildExecutionCtxParams{
+			Schema:        p.Schema,
+			Root:          p.RootValue,
+			AST:           p.Document,
+			OperationName: p.OperationName,
+			Args:          p.VariableValues,
+			Result:        result,
+			Context:       p.ContextValue,
+		})
+
+		if err != nil {
+			result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
+			resultChannel <- result
+			return
+		}
+
+		operationType, err := getOperationRootType(p.Schema, exeContext.Operation)
+		if err != nil {
+			result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
+			resultChannel <- result
+			return
+		}
+
+		fields := collectFields(collectFieldsParams{
+			ExeContext:   exeContext,
+			RuntimeType:  operationType,
+			SelectionSet: exeContext.Operation.GetSelectionSet(),
+		})
+
+		responseNames := []string{}
+		for name := range fields {
+			responseNames = append(responseNames, name)
+		}
+		responseName := responseNames[0]
+		fieldNodes := fields[responseName]
+		fieldNode := fieldNodes[0]
+		fieldName := fieldNode.Name.Value
+		fieldDef := getFieldDef(p.Schema, operationType, fieldName)
+
+		if fieldDef == nil {
+			err := fmt.Errorf("the subscription field %q is not defined", fieldName)
+			result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
+			resultChannel <- result
+			return
+		}
+
+		resolveFn := p.FieldSubscriber
+		if resolveFn == nil {
+			resolveFn = DefaultResolveFn
+		}
+		if fieldDef.Subscribe != nil {
+			resolveFn = fieldDef.Subscribe
+		}
+		fieldPath := &ResponsePath{
+			Key: responseName,
+		}
+
+		args := getArgumentValues(fieldDef.Args, fieldNode.Arguments, exeContext.VariableValues)
+		info := ResolveInfo{
+			FieldName:      fieldName,
+			FieldASTs:      fieldNodes,
+			Path:           fieldPath,
+			ReturnType:     fieldDef.Type,
+			ParentType:     operationType,
+			Schema:         p.Schema,
+			Fragments:      exeContext.Fragments,
+			RootValue:      exeContext.Root,
+			Operation:      exeContext.Operation,
+			VariableValues: exeContext.VariableValues,
+		}
+
+		fieldResult, err := resolveFn(ResolveParams{
+			Source:  p.RootValue,
+			Args:    args,
+			Info:    info,
+			Context: exeContext.Context,
+		})
+		if err != nil {
+			result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
+			resultChannel <- result
+			return
+		}
+
+		if fieldResult == nil {
+			err := fmt.Errorf("no field result")
+			result.Errors = append(result.Errors, gqlerrors.FormatError(err.(error)))
+			resultChannel <- result
+			return
+		}
+
+		switch fieldResult.(type) {
+		case chan interface{}:
+			for {
+				select {
+				case <-ctx.Done():
+					fmt.Printf("done context called")
+					return
+				case res := <-fieldResult.(chan interface{}):
+
+					resultChannel <- mapSourceToResponse(res)
+				}
+			}
+		default:
+			resultChannel <- mapSourceToResponse(fieldResult)
+			return
+		}
+	}()
+
+	// return a result iterator
+	return NewResultIterator(p.ContextValue, resultChannel)
+}

subscription_test.go

@@ -0,0 +1,134 @@
+package graphql
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/graphql-go/graphql/language/parser"
+	"github.com/graphql-go/graphql/language/source"
+)
+
+func TestSubscription(t *testing.T) {
+	var maxPublish = 5
+	m := make(chan interface{})
+
+	source1 := source.NewSource(&source.Source{
+		Body: []byte(`subscription {
+			watch_count
+		}`),
+		Name: "GraphQL request",
+	})
+
+	source2 := source.NewSource(&source.Source{
+		Body: []byte(`subscription {
+			watch_should_fail
+		}`),
+		Name: "GraphQL request",
+	})
+
+	document1, _ := parser.Parse(parser.ParseParams{Source: source1})
+	document2, _ := parser.Parse(parser.ParseParams{Source: source2})
+
+	schema, err := NewSchema(SchemaConfig{
+		Query: NewObject(ObjectConfig{
+			Name: "Query",
+			Fields: Fields{
+				"hello": &Field{
+					Type: String,
+					Resolve: func(p ResolveParams) (interface{}, error) {
+						return "world", nil
+					},
+				},
+			},
+		}),
+		Subscription: NewObject(ObjectConfig{
+			Name: "Subscription",
+			Fields: Fields{
+				"watch_count": &Field{
+					Type: String,
+					Resolve: func(p ResolveParams) (interface{}, error) {
+						return fmt.Sprintf("count=%v", p.Source), nil
+					},
+					Subscribe: func(p ResolveParams) (interface{}, error) {
+						return m, nil
+					},
+				},
+				"watch_should_fail": &Field{
+					Type: String,
+					Resolve: func(p ResolveParams) (interface{}, error) {
+						return fmt.Sprintf("count=%v", p.Source), nil
+					},
+					Subscribe: func(p ResolveParams) (interface{}, error) {
+						return nil, nil
+					},
+				},
+			},
+		}),
+	})
+
+	if err != nil {
+		t.Errorf("failed to create schema: %v", err)
+		return
+	}
+
+	failIterator := Subscribe(SubscribeParams{
+		Schema:   schema,
+		Document: document2,
+	})
+
+	// test a subscribe that should fail due to no return value
+	failIterator.ForEach(func(count int64, res *Result, doneFunc func()) {
+		if !res.HasErrors() {
+			t.Errorf("subscribe failed to catch nil result from subscribe")
+			doneFunc()
+			return
+		}
+		doneFunc()
+		return
+	})
+
+	resultIterator := Subscribe(SubscribeParams{
+		Schema:       schema,
+		Document:     document1,
+		ContextValue: context.Background(),
+	})
+
+	resultIterator.ForEach(func(count int64, res *Result, doneFunc func()) {
+		if res.HasErrors() {
+			t.Errorf("subscribe error(s): %v", res.Errors)
+			doneFunc()
+			return
+		}
+
+		if res.Data != nil {
+			data := res.Data.(map[string]interface{})["watch_count"]
+			expected := fmt.Sprintf("count=%d", count)
+			actual := fmt.Sprintf("%v", data)
+			if actual != expected {
+				t.Errorf("subscription result error: expected %q, actual %q", expected, actual)
+				doneFunc()
+				return
+			}
+
+			// test the done func by quitting after 3 iterations
+			// the publisher will publish up to 5
+			if count >= int64(maxPublish-2) {
+				doneFunc()
+				return
+			}
+		}
+	})
+
+	// start publishing
+	go func() {
+		for i := 1; i <= maxPublish; i++ {
+			time.Sleep(200 * time.Millisecond)
+			m <- i
+		}
+	}()
+
+	// give time for the test to complete
+	time.Sleep(1 * time.Second)
+}