draft: resolving defer response

Author: devsergiy

execution/engine/execution_engine.go

@@ -155,6 +155,7 @@ func (e *ExecutionEngine) Execute(ctx context.Context, operation *graphql.Reques
 			astnormalization.WithRemoveFragmentDefinitions(),
 			astnormalization.WithRemoveUnusedVariables(),
 			astnormalization.WithInlineFragmentSpreads(),
+			astnormalization.WithInlineDefer(),
 		)
 		if err != nil {
 			return err
@@ -229,6 +230,9 @@ func (e *ExecutionEngine) Execute(ctx context.Context, operation *graphql.Reques
 	case *plan.SynchronousResponsePlan:
 		_, err := e.resolver.ResolveGraphQLResponse(execContext.resolveContext, p.Response, nil, writer)
 		return err
+	case *plan.DeferResponsePlan:
+		_, err := e.resolver.ResolveGraphQLDeferResponse(execContext.resolveContext, p.Response, writer)
+		return err
 	case *plan.SubscriptionResponsePlan:
 		return e.resolver.ResolveGraphQLSubscription(execContext.resolveContext, p.Response, writer)
 	default:

execution/engine/execution_engine_test.go

@@ -220,6 +220,7 @@ type _executionTestOptions struct {
 	apolloRouterCompatibilitySubrequestHTTPError bool
 	propagateFetchReasons                        bool
 	validateRequiredExternalFields               bool
+	streamingResponse                            bool
 }
 
 type executionTestOptions func(*_executionTestOptions)
@@ -244,6 +245,12 @@ func validateRequiredExternalFields() executionTestOptions {
 	}
 }
 
+func withStreamingResponse() executionTestOptions {
+	return func(options *_executionTestOptions) {
+		options.streamingResponse = true
+	}
+}
+
 func TestExecutionEngine_Execute(t *testing.T) {
 	run := func(testCase ExecutionEngineTestCase, withError bool, expectedErrorMessage string, options ...executionTestOptions) func(t *testing.T) {
 		t.Helper()
@@ -289,6 +296,14 @@ func TestExecutionEngine_Execute(t *testing.T) {
 
 			operation := testCase.operation(t)
 			resultWriter := graphql.NewEngineResultWriter()
+
+			streamingBuf := bytes.NewBuffer(nil)
+			if opts.streamingResponse {
+				resultWriter.SetFlushCallback(func(data []byte) {
+					streamingBuf.Write(data)
+				})
+			}
+
 			execCtx, execCtxCancel := context.WithCancel(context.Background())
 			defer execCtxCancel()
 			err = engine.Execute(execCtx, &operation, &resultWriter, testCase.engineOptions...)
@@ -311,7 +326,12 @@ func TestExecutionEngine_Execute(t *testing.T) {
 			}
 
 			if testCase.expectedResponse != "" {
-				assert.Equal(t, testCase.expectedResponse, actualResponse)
+				if opts.streamingResponse {
+					streamingResponse := streamingBuf.String()
+					assert.Equal(t, testCase.expectedResponse, streamingResponse)
+				} else {
+					assert.Equal(t, testCase.expectedResponse, actualResponse)
+				}
 			}
 
 			if withError {
@@ -5603,6 +5623,100 @@ func TestExecutionEngine_Execute(t *testing.T) {
 			}, withFetchReasons(), validateRequiredExternalFields()))
 		})
 	})
+
+	t.Run("defer", func(t *testing.T) {
+		t.Run("simple", func(t *testing.T) {
+
+			definition := `
+				type User {
+					id: ID!
+					name: String!
+					title: String!
+				}
+
+				type Query {
+					user: User!
+				}
+			`
+
+			makeDataSource := func(t *testing.T, expectFetchReasons bool) []plan.DataSource {
+				return []plan.DataSource{
+					mustGraphqlDataSourceConfiguration(t,
+						"id-1",
+						mustFactory(t,
+							testConditionalNetHttpClient(t, conditionalTestCase{
+								expectedHost: "first",
+								expectedPath: "/",
+								responses: map[string]sendResponse{
+									`{"query":"{user {name}}"}`: {
+										statusCode: 200,
+										body:       `{"data":{"user":{"name":"Black"}}}`,
+									},
+									`{"query":"{user {title}}"}`: {
+										statusCode: 200,
+										body:       `{"data":{"user":{"title":"Sabbat"}}}`,
+									},
+								},
+							}),
+						),
+						&plan.DataSourceMetadata{
+							RootNodes: []plan.TypeField{
+								{
+									TypeName:   "Query",
+									FieldNames: []string{"user"},
+								},
+							},
+							ChildNodes: []plan.TypeField{
+								{
+									TypeName:   "User",
+									FieldNames: []string{"id", "title", "name"},
+								},
+							},
+						},
+						mustConfiguration(t, graphql_datasource.ConfigurationInput{
+							Fetch: &graphql_datasource.FetchConfiguration{
+								URL:    "https://first/",
+								Method: "POST",
+							},
+							SchemaConfiguration: mustSchemaConfig(
+								t,
+								&graphql_datasource.FederationConfiguration{
+									Enabled:    true,
+									ServiceSDL: definition,
+								},
+								definition,
+							),
+						}),
+					),
+				}
+			}
+
+			t.Run("run", runWithoutError(ExecutionEngineTestCase{
+				schema: func(t *testing.T) *graphql.Schema {
+					t.Helper()
+					parseSchema, err := graphql.NewSchemaFromString(definition)
+					require.NoError(t, err)
+					return parseSchema
+				}(t),
+				operation: func(t *testing.T) graphql.Request {
+					return graphql.Request{
+						OperationName: "DeferUserTitle",
+						Query: `
+						query DeferUserTitle {
+							user {
+								name
+								... @defer {
+									title
+								}
+							}
+						}`,
+					}
+				},
+				dataSources:      makeDataSource(t, false),
+				expectedResponse: `{"data":{"user":{"name":"Black"}}}{"name":"Black"}{"data":{{"name":"Black"}}}`,
+			}, withStreamingResponse()))
+		})
+	})
 }
 
 func testNetHttpClient(t *testing.T, testCase roundTripperTestCase) *http.Client {

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go

@@ -150,6 +150,11 @@ func (p *Planner[T]) EnterDirective(ref int) {
 }
 
 func (p *Planner[T]) addDirectiveToNode(directiveRef int, node ast.Node) {
+	// do not propagate internal directives to upstream query document
+	if bytes.Equal(p.visitor.Operation.DirectiveNameBytes(directiveRef), literal.DEFER_INTERNAL) {
+		return
+	}
+
 	directiveName := p.visitor.Operation.DirectiveNameString(directiveRef)
 	operationType := ast.OperationTypeQuery
 	if !p.dataSourcePlannerConfig.IsNested {

v2/pkg/engine/resolve/loader.go

@@ -190,14 +190,19 @@ func (l *Loader) Free() {
 }
 
 func (l *Loader) LoadGraphQLResponseData(ctx *Context, response *GraphQLResponse, resolvable *Resolvable) (err error) {
+	l.Init(ctx, response.Info, resolvable)
+
+	return l.ResolveFetchNode(response.Fetches)
+}
+
+func (l *Loader) Init(ctx *Context, responseInfo *GraphQLResponseInfo, resolvable *Resolvable) {
 	l.resolvable = resolvable
 	l.ctx = ctx
-	l.info = response.Info
+	l.info = responseInfo
 	l.taintedObjs = make(taintedObjects)
-	return l.resolveFetchNode(response.Fetches)
 }
 
-func (l *Loader) resolveFetchNode(node *FetchTreeNode) error {
+func (l *Loader) ResolveFetchNode(node *FetchTreeNode) error {
 	if node == nil {
 		return nil
 	}
@@ -264,7 +269,7 @@ func (l *Loader) resolveParallel(nodes []*FetchTreeNode) error {
 
 func (l *Loader) resolveSerial(nodes []*FetchTreeNode) error {
 	for i := range nodes {
-		err := l.resolveFetchNode(nodes[i])
+		err := l.ResolveFetchNode(nodes[i])
 		if err != nil {
 			return errors.WithStack(err)
 		}
@@ -629,6 +634,7 @@ func (l *Loader) mergeResult(fetchItem *FetchItem, res *result, items []*astjson
 		if responseData.Type() != astjson.TypeObject {
 			return l.renderErrorsFailedToFetch(fetchItem, res, invalidGraphQLResponseShape)
 		}
+		// TODO: unclear why we doing this
 		l.resolvable.data = responseData
 		return nil
 	}

v2/pkg/engine/resolve/resolvable.go

@@ -50,6 +50,8 @@ type Resolvable struct {
 	wroteErrors         bool
 	wroteData           bool
 	skipValueCompletion bool
+	deferMode           bool
+	deferID             string
 
 	typeNames [][]byte
 
@@ -262,6 +264,10 @@ func (r *Resolvable) printData(root *Object) {
 	r.printBytes(colon)
 	r.printBytes(lBrace)
 	r.print = true
+
+	if r.deferMode && r.deferID != "" {
+		r.print = false
+	}
 	_ = r.walkObject(root, r.data)
 	r.print = false
 	r.printBytes(rBrace)
@@ -615,6 +621,21 @@ func (r *Resolvable) walkObject(obj *Object, parent *astjson.Value) bool {
 		}
 	}
 
+	hasDeferredFields := false
+	for i := range obj.Fields {
+		if obj.Fields[i].Defer != nil && obj.Fields[i].Defer.DeferID == r.deferID {
+			hasDeferredFields = true
+		}
+	}
+
+	if hasDeferredFields {
+		// enable printing
+		if !r.print {
+			r.print = true
+			defer func() { r.print = false }()
+		}
+	}
+
 	if r.print && !isRoot {
 		r.printBytes(lBrace)
 	}
@@ -625,6 +646,22 @@ func (r *Resolvable) walkObject(obj *Object, parent *astjson.Value) bool {
 		r.typeNames = r.typeNames[:len(r.typeNames)-1]
 	}()
 	for i := range obj.Fields {
+		if r.deferMode {
+
+			// for initial response render only fields without a defer id
+			if r.deferID == "" && obj.Fields[i].Defer != nil {
+				continue
+			}
+
+			// skip fields with different defer id in deferred response
+			if r.deferID != "" && obj.Fields[i].Defer != nil && obj.Fields[i].Defer.DeferID == r.deferID {
+				continue
+			}
+
+			// when defer id matches render field - render it
+			// walk objects to find other defers
+		}
+
 		if obj.Fields[i].ParentOnTypeNames != nil {
 			if r.skipFieldOnParentTypeNames(obj.Fields[i]) {
 				continue

v2/pkg/engine/resolve/resolve.go

@@ -291,6 +291,71 @@ func (r *Resolver) ResolveGraphQLResponse(ctx *Context, response *GraphQLRespons
 	return resp, err
 }
 
+func (r *Resolver) ResolveGraphQLDeferResponse(ctx *Context, response *GraphQLDeferResponse, writer DeferResponseWriter) (*GraphQLResolveInfo, error) {
+	resp := &GraphQLResolveInfo{}
+
+	start := time.Now()
+	<-r.maxConcurrency
+	resp.ResolveAcquireWaitTime = time.Since(start)
+	defer func() {
+		r.maxConcurrency <- struct{}{}
+	}()
+
+	t := newTools(r.options, r.allowedErrorExtensionFields, r.allowedErrorFields)
+
+	err := t.resolvable.Init(ctx, nil, response.Response.Info.OperationType)
+	if err != nil {
+		return nil, err
+	}
+
+	if !ctx.ExecutionOptions.SkipLoader {
+		t.loader.Init(ctx, response.Response.Info, t.resolvable)
+
+		// fetch initial response
+		if err := t.loader.ResolveFetchNode(response.Response.Fetches); err != nil {
+			return nil, err
+		}
+
+		t.resolvable.deferMode = true
+		t.resolvable.deferID = ""
+
+		// render initial response
+		err = t.resolvable.Resolve(ctx.ctx, response.Response.Data, response.Response.Fetches, writer)
+		if err != nil {
+			return nil, err
+		}
+
+		err = writer.Flush()
+		if err != nil {
+			return nil, err
+		}
+
+		// fetch deferred responses
+
+		for _, deferGroup := range response.Defers {
+			if err := t.loader.ResolveFetchNode(deferGroup.Fetches); err != nil {
+				return nil, err
+			}
+
+			// render deferred response
+			t.resolvable.deferID = deferGroup.DeferID
+			err = t.resolvable.Resolve(ctx.ctx, response.Response.Data, deferGroup.Fetches, writer)
+			if err != nil {
+				return nil, err
+			}
+
+			// flush after each deferred response
+
+			err = writer.Flush()
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	return resp, err
+}
+
 type trigger struct {
 	id            uint64
 	cancel        context.CancelFunc

v2/pkg/engine/resolve/response.go

@@ -63,6 +63,11 @@ type ResponseWriter interface {
 	io.Writer
 }
 
+type DeferResponseWriter interface {
+	ResponseWriter
+	Flush() error
+}
+
 type SubscriptionCloseKind struct {
 	WSCode ws.StatusCode
 	Reason string