global subscription context config option (#4237)

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

# Conflicts:
#	codegen/field.go
#	codegen/internal!.gotpl
#	codegen/object.go
#	docs/content/reference/subscription-context.md
#	graphql/event.go
#	graphql/event_test.go

Author: StevenACoffman

.golangci.yml

@@ -19,6 +19,7 @@ formatters:
         - codegen/testserver/followschema/resolver.go
         - codegen/testserver/singlefile/resolver.go
         - codegen/testserver/usefunctionsyntaxforexecutioncontext/resolver.go
+        - codegen/testserver/subscriptioncontextfield/resolver.go
         - generated
   enable:
     - golines

_examples/chat/generated.go

@@ -49,6 +49,7 @@ type Config struct {
 	OmitPanicHandler                     bool                       `yaml:"omit_panic_handler,omitempty"`
 	OmitEnumJSONMarshalers               bool                       `yaml:"omit_enum_json_marshalers,omitempty"`
 	UseFunctionSyntaxForExecutionContext bool                       `yaml:"use_function_syntax_for_execution_context,omitempty"`
+	SubscriptionContextField             bool                       `yaml:"subscription_context_field,omitempty"`
 	// If this is set to true, argument directives that
 	// decorate a field with a null value will still be called.
 	//

codegen/config/config.go

@@ -78,7 +78,17 @@ func {{.FuncReceiver}}_mutationMiddleware(ctx context.Context, {{.ECFuncParam}}o
 {{ end }}
 
 {{ if .AllDirectives.LocationDirectives "SUBSCRIPTION" }}
-func {{$.FuncReceiver}}_subscriptionMiddleware(ctx context.Context, {{$.ECFuncParam}}obj *ast.OperationDefinition, next func(ctx context.Context) (any, error)) func(ctx context.Context) graphql.Marshaler {
+{{- /* The subscription middleware returns a per-event-context tuple handler when the
+   subscription_context_field option is on, and a plain marshaler handler otherwise. The shape
+   is decided once here; the runtime supplies the matching null/error handler so its closure is
+   not duplicated into the generated executor. */ -}}
+{{- $retType := "graphql.Marshaler" -}}
+{{- $nullHandler := "graphql.NullStream()" -}}
+{{- if $.Config.SubscriptionContextField -}}
+	{{- $retType = "(context.Context, graphql.Marshaler)" -}}
+	{{- $nullHandler = "graphql.NullEventStream()" -}}
+{{- end -}}
+func {{$.FuncReceiver}}_subscriptionMiddleware(ctx context.Context, {{$.ECFuncParam}}obj *ast.OperationDefinition, next func(ctx context.Context) (any, error)) func(ctx context.Context) {{$retType}} {
 	for _, d := range obj.Directives {
 		switch d.Name {
 		{{- range $directive := .AllDirectives.LocationDirectives "SUBSCRIPTION" }}
@@ -88,9 +98,7 @@ func {{$.FuncReceiver}}_subscriptionMiddleware(ctx context.Context, {{$.ECFuncPa
 				args, err := {{$.ECDot}}{{ $directive.ArgsFunc }}(ctx,{{$.ECArg}}rawArgs)
 				if err != nil {
 					ec.Error(ctx, err)
-					return func(ctx context.Context) graphql.Marshaler {
-						return graphql.Null
-					}
+					return {{$nullHandler}}
 				}
 			{{- end }}
 			n := next
@@ -103,17 +111,13 @@ func {{$.FuncReceiver}}_subscriptionMiddleware(ctx context.Context, {{$.ECFuncPa
 	tmp, err := next(ctx)
 	if err != nil {
 		ec.Error(ctx, err)
-		return func(ctx context.Context) graphql.Marshaler {
-			return graphql.Null
-		}
+		return {{$nullHandler}}
 	}
-	if data, ok := tmp.(func(ctx context.Context) graphql.Marshaler); ok {
+	if data, ok := tmp.(func(ctx context.Context) {{$retType}}); ok {
 		return data
 	}
 	graphql.AddErrorf(ctx, `unexpected type %T from directive, should be graphql.Marshaler`, tmp)
-	return func(ctx context.Context) graphql.Marshaler {
-		return graphql.Null
-	}
+	return {{$nullHandler}}
 }
 {{ end }}
 

codegen/directives_.gotpl

@@ -41,6 +41,10 @@ type Field struct {
 	HasHaser         bool   // Whether a haser method is available (e.g., HasName())
 	HaserMethodName  string // Name of the haser method
 	Batch            bool   // Enable batch resolver for this field
+	// SubscriptionContextField mirrors the global subscription_context_field config
+	// option, resolved once at build time so UsesSubscriptionContext and the methods
+	// that depend on it stay nullary instead of threading the flag through the call chain.
+	SubscriptionContextField bool
 }
 
 func (b *builder) buildField(obj *Object, field *ast.FieldDefinition) (*Field, error) {
@@ -50,12 +54,13 @@ func (b *builder) buildField(obj *Object, field *ast.FieldDefinition) (*Field, e
 	}
 
 	f := Field{
-		FieldDefinition: field,
-		Object:          obj,
-		Directives:      dirs,
-		GoFieldName:     templates.ToGo(field.Name),
-		GoFieldType:     GoFieldVariable,
-		GoReceiverName:  "obj",
+		FieldDefinition:          field,
+		Object:                   obj,
+		Directives:               dirs,
+		GoFieldName:              templates.ToGo(field.Name),
+		GoFieldType:              GoFieldVariable,
+		GoReceiverName:           "obj",
+		SubscriptionContextField: b.Config.SubscriptionContextField,
 	}
 
 	if field.DefaultValue != nil {
@@ -563,14 +568,18 @@ func (f *Field) HasDirectives() bool {
 }
 
 // UsesSubscriptionContext reports whether this field is on the Subscription
-// root type and is annotated with @subscriptionContext. Codegen uses this to
+// root type and is annotated with @subscriptionContext, or the global
+// subscription_context_field option is enabled. Codegen uses this to
 // emit a resolver returning <-chan graphql.Event[T] instead of <-chan T, and
 // to thread per-event context into the AroundResponses interceptor chain.
 // Returns false for non-subscription fields even if they carry the directive.
 func (f *Field) UsesSubscriptionContext() bool {
 	if !f.Object.Stream {
 		return false
 	}
+	if f.SubscriptionContextField {
+		return true
+	}
 	for _, d := range f.FieldDefinition.Directives {
 		if d.Name == config.DirSubscriptionContext {
 			return true

codegen/field.go

@@ -9,7 +9,7 @@ func {{$.FuncReceiver}}_{{$object.Name}}_{{$field.Name}}(ctx context.Context, {{
 	{{- if $field.TypeReference.IsRoot }}
 		fc, err := {{$.ECDot}}{{ $field.FieldContextFunc }}(ctx, {{$.ECArg}}field)
 		if err != nil {
-			return {{ $null }}
+			return {{- if $field.UsesSubscriptionContext }} ctx,{{- end}} {{ $null }}
 		}
 		ctx = graphql.WithFieldContext(ctx, fc)
 		{{- if not $.Config.OmitPanicHandler }}

codegen/field.gotpl

@@ -156,7 +156,7 @@ func (e *executableSchema) Exec(ctx context.Context) graphql.ResponseHandler {
 	}
 }
 
-{{- if and .SubscriptionRoot .SubscriptionRoot.HasSubscriptionContextField }}
+{{- if and .SubscriptionRoot (.SubscriptionRoot.HasSubscriptionContextField) }}
 
 // ExecWithEventContext is the per-event-context-aware response handler. The
 // graphql executor type-asserts the executable schema for this method when

codegen/internal!.gotpl

@@ -140,7 +140,8 @@ func (o *Object) HasDirectives() bool {
 }
 
 // HasSubscriptionContextField reports whether this object is a streaming
-// (subscription) root with at least one field annotated @subscriptionContext.
+// (subscription) root with at least one field annotated @subscriptionContext,
+// or the global subscription_context_field option is enabled.
 // Codegen uses this to decide whether to emit the event-context-aware
 // dispatcher and the optional ExecWithEventContext method on the
 // generated executableSchema. Returns false for non-streaming objects.