feat: add subscriptions hook

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

@@ -9,6 +9,7 @@ import (
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/astparser"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/asttransform"
 	grpcdatasource "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/grpc_datasource"
+	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/federation"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/operationreport"
 )
@@ -103,6 +104,13 @@ type SingleTypeField struct {
 	FieldName string
 }
 
+// OnSubscriptionStartFn defines a hook function that is called when a subscription starts.
+// It receives the resolve context and the input of the subscription.
+// The function can return a boolean indicating if the subscription should be closed, and an error.
+// The error is propagated to the client.
+// If close is true, the subscription is closed.
+type OnSubscriptionStartFn func(ctx *resolve.Context, input []byte) (close bool, err error)
+
 type SubscriptionConfiguration struct {
 	URL           string
 	Header        http.Header
@@ -119,6 +127,7 @@ type SubscriptionConfiguration struct {
 	// these headers by itself.
 	ForwardedClientHeaderRegularExpressions []RegularExpression
 	WsSubProtocol                           string
+	OnSubscriptionStartFns                  []OnSubscriptionStartFn
 }
 
 type FetchConfiguration struct {

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

@@ -450,7 +450,8 @@ func (p *Planner[T]) ConfigureSubscription() plan.SubscriptionConfiguration {
 	return plan.SubscriptionConfiguration{
 		Input: string(input),
 		DataSource: &SubscriptionSource{
-			client: p.subscriptionClient,
+			client:                 p.subscriptionClient,
+			onSubscriptionStartFns: p.config.subscription.OnSubscriptionStartFns,
 		},
 		Variables:      p.variables,
 		PostProcessing: DefaultPostProcessingConfiguration,
@@ -1939,7 +1940,8 @@ type RegularExpression struct {
 }
 
 type SubscriptionSource struct {
-	client GraphQLSubscriptionClient
+	client                 GraphQLSubscriptionClient
+	onSubscriptionStartFns []OnSubscriptionStartFn
 }
 
 func (s *SubscriptionSource) AsyncStart(ctx *resolve.Context, id uint64, input []byte, updater resolve.SubscriptionUpdater) error {
@@ -1989,3 +1991,13 @@ func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte,
 	}
 	return s.client.UniqueRequestID(ctx, options, xxh)
 }
+
+func (s *SubscriptionSource) OnSubscriptionStart(ctx *resolve.Context, input []byte) (close bool, err error) {
+	for _, fn := range s.onSubscriptionStartFns {
+		close, err = fn(ctx, input)
+		if err != nil || close {
+			return close, err
+		}
+	}
+	return
+}

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

@@ -4019,7 +4019,7 @@ func TestGraphQLDataSource(t *testing.T) {
 				Trigger: resolve.GraphQLSubscriptionTrigger{
 					Input: []byte(`{"url":"wss://swapi.com/graphql","body":{"query":"subscription{remainingJedis}"}}`),
 					Source: &SubscriptionSource{
-						NewGraphQLSubscriptionClient(http.DefaultClient, http.DefaultClient, ctx),
+						client: NewGraphQLSubscriptionClient(http.DefaultClient, http.DefaultClient, ctx),
 					},
 					PostProcessing: DefaultPostProcessingConfiguration,
 				},
@@ -8902,6 +8902,56 @@ func TestSanitizeKey(t *testing.T) {
 	}
 }
 
+func TestSubscriptionSource_OnSubscriptionStart(t *testing.T) {
+
+	t.Run("OnSubscriptionStart calls onSubscriptionStartFns", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		startFnCalled := make(chan struct {
+			ctx   *resolve.Context
+			input []byte
+		}, 1)
+		subscriptionSource := SubscriptionSource{
+			onSubscriptionStartFns: []OnSubscriptionStartFn{
+				func(ctx *resolve.Context, input []byte) (bool, error) {
+					startFnCalled <- struct {
+						ctx   *resolve.Context
+						input []byte
+					}{ctx, input}
+					return false, nil
+				},
+			},
+		}
+
+		close, err := subscriptionSource.OnSubscriptionStart(ctx, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`))
+		require.NoError(t, err)
+		assert.False(t, close)
+		require.Len(t, startFnCalled, 1)
+		called := <-startFnCalled
+		assert.Equal(t, ctx, called.ctx)
+		assert.Equal(t, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`), called.input)
+	})
+
+	t.Run("OnSubscriptionStart calls onSubscriptionStartFns and returns error if one of the functions returns an error", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		subscriptionSource := SubscriptionSource{
+			onSubscriptionStartFns: []OnSubscriptionStartFn{
+				func(ctx *resolve.Context, input []byte) (bool, error) {
+					return false, errors.New("test error")
+				},
+			},
+		}
+
+		close, err := subscriptionSource.OnSubscriptionStart(ctx, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`))
+		assert.False(t, close)
+		require.Error(t, err)
+		assert.ErrorContains(t, err, "test error")
+	})
+}
+
 const interfaceSelectionSchema = `
 
 scalar String

v2/pkg/engine/resolve/context.go

@@ -32,6 +32,8 @@ type Context struct {
 	fieldRenderer FieldValueRenderer
 
 	subgraphErrors error
+
+	subscriptionUpdater func(data []byte)
 }
 
 type ExecutionOptions struct {
@@ -104,6 +106,22 @@ func (c *Context) SetEngineLoaderHooks(hooks LoaderHooks) {
 	c.LoaderHooks = hooks
 }
 
+// SetSubscriptionUpdater add a function that will be called when EmitSubscriptionUpdate is called
+// usually it is set by the resolver when a subscription is started, but it can be set to nil to disable the feature
+// or set to a different function to extend the behaviour or make it easy to test the subscription update logic
+func (c *Context) SetSubscriptionUpdater(fn func(data []byte)) {
+	c.subscriptionUpdater = fn
+}
+
+// EmitSubscriptionUpdate emits a subscription update to the client
+// if the emitSubscriptionUpdate function is not set, the update is not sent to the client
+// this is used to allow external code to emit updates on this subscription
+func (c *Context) EmitSubscriptionUpdate(data []byte) {
+	if c.subscriptionUpdater != nil {
+		c.subscriptionUpdater(data)
+	}
+}
+
 type RateLimitOptions struct {
 	// Enable switches rate limiting on or off
 	Enable bool

v2/pkg/engine/resolve/datasource.go

@@ -26,3 +26,16 @@ type AsyncSubscriptionDataSource interface {
 	AsyncStop(id uint64)
 	UniqueRequestID(ctx *Context, input []byte, xxh *xxhash.Digest) (err error)
 }
+
+// SubscriptionDataSourceHookable is a hookable interface for subscription data sources.
+// It is used to call a function when a subscription is started.
+// This is useful for data sources that need to do some work when a subscription is started,
+// e.g. to establish a connection to the data source or to emit updates to the client.
+// The function is called with the context and the input of the subscription.
+// The function is called before the subscription is started and can be used to emit updates to the client.
+type SubscriptionDataSourceHookable interface {
+	// OnSubscriptionStart is called when a new subscription is created
+	// If close is true, the subscription is closed.
+	// If an error is returned, the error is propagated to the client.
+	OnSubscriptionStart(ctx *Context, input []byte) (close bool, err error)
+}

v2/pkg/engine/resolve/resolve.go

@@ -569,6 +569,16 @@ func (r *Resolver) handleTriggerComplete(triggerID uint64) {
 	r.completeTrigger(triggerID)
 }
 
+// callSubscriptionDataSourceStartHook is used to call the OnSubscriptionStart method of the subscription data source
+// if the subscription data source implements the SubscriptionDataSourceHook interface.
+// This is used to allow external code to emit updates on this subscription.
+func callSubscriptionDataSourceStartHook(ctx *Context, source SubscriptionDataSource, input []byte) (close bool, err error) {
+	if hook, ok := source.(SubscriptionDataSourceHookable); ok {
+		return hook.OnSubscriptionStart(ctx, input)
+	}
+	return false, nil
+}
+
 func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription) {
 	var (
 		err error
@@ -590,6 +600,12 @@ func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription)
 		s.heartbeat = true
 	}
 
+	// Set the subscription updater to the resolver to allow external code to emit updates
+	// on this subscription
+	add.ctx.SetSubscriptionUpdater(func(data []byte) {
+		r.handleTriggerUpdateSubscription(add.ctx, s, data)
+	})
+
 	// Start the dedicated worker goroutine where the subscription updates are processed
 	// and writes are written to the client in a single threaded manner
 	go s.startWorker()
@@ -603,6 +619,13 @@ func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription)
 		if r.options.Debug {
 			fmt.Printf("resolver:trigger:subscription:added:%d:%d\n", triggerID, add.id.SubscriptionID)
 		}
+		close, errStartHook := callSubscriptionDataSourceStartHook(add.ctx, add.resolve.Trigger.Source, add.input)
+		if errStartHook != nil {
+			r.asyncErrorWriter.WriteError(add.ctx, errStartHook, add.resolve.Response, add.writer)
+		}
+		if close {
+			r.closeTrigger(triggerID, SubscriptionCloseKindNormal)
+		}
 		return
 	}
 
@@ -639,6 +662,15 @@ func (r *Resolver) handleAddSubscription(triggerID uint64, add *addSubscription)
 		asyncDataSource = async
 	}
 
+	close, errStartHook := callSubscriptionDataSourceStartHook(add.ctx, add.resolve.Trigger.Source, add.input)
+	if errStartHook != nil {
+		r.asyncErrorWriter.WriteError(add.ctx, errStartHook, add.resolve.Response, add.writer)
+	}
+	if close {
+		r.closeTrigger(triggerID, SubscriptionCloseKindNormal)
+		return
+	}
+
 	go func() {
 		if r.options.Debug {
 			fmt.Printf("resolver:trigger:start:%d\n", triggerID)
@@ -757,6 +789,34 @@ func (r *Resolver) handleRemoveClient(id int64) {
 	}
 }
 
+func (r *Resolver) handleTriggerUpdateSubscription(ctx *Context, sub *sub, data []byte) {
+	if err := ctx.ctx.Err(); err != nil {
+		return // no need to schedule an event update when the client already disconnected
+	}
+	skip, err := sub.resolve.Filter.SkipEvent(ctx, data, r.triggerUpdateBuf)
+	if err != nil {
+		r.asyncErrorWriter.WriteError(ctx, err, sub.resolve.Response, sub.writer)
+		return
+	}
+	if skip {
+		return
+	}
+
+	fn := func() {
+		r.executeSubscriptionUpdate(ctx, sub, data)
+	}
+
+	select {
+	case <-r.ctx.Done():
+		// Skip sending all events if the resolver is shutting down
+		return
+	case <-ctx.ctx.Done():
+		// Skip sending the event if the client disconnected
+	case sub.workChan <- workItem{fn, false}:
+		// Send the event to the subscription worker
+	}
+}
+
 func (r *Resolver) handleTriggerUpdate(id uint64, data []byte) {
 	trig, ok := r.triggers[id]
 	if !ok {
@@ -767,32 +827,7 @@ func (r *Resolver) handleTriggerUpdate(id uint64, data []byte) {
 	}
 
 	for c, s := range trig.subscriptions {
-		c, s := c, s
-		if err := c.ctx.Err(); err != nil {
-			continue // no need to schedule an event update when the client already disconnected
-		}
-		skip, err := s.resolve.Filter.SkipEvent(c, data, r.triggerUpdateBuf)
-		if err != nil {
-			r.asyncErrorWriter.WriteError(c, err, s.resolve.Response, s.writer)
-			continue
-		}
-		if skip {
-			continue
-		}
-
-		fn := func() {
-			r.executeSubscriptionUpdate(c, s, data)
-		}
-
-		select {
-		case <-r.ctx.Done():
-			// Skip sending all events if the resolver is shutting down
-			return
-		case <-c.ctx.Done():
-			// Skip sending the event if the client disconnected
-		case s.workChan <- workItem{fn, false}:
-			// Send the event to the subscription worker
-		}
+		r.handleTriggerUpdateSubscription(c, s, data)
 	}
 }