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,11 @@ type SingleTypeField struct {
 	FieldName string
 }
 
+// OnSubscriptionStartFn defines a hook function that is called when a subscription starts.
+// It receives the resolve context and can return initial events or an error.
+// If an error is returned, the subscription will not start.
+type OnSubscriptionStartFn func(ctx *resolve.Context) ([][]byte, error)
+
 type SubscriptionConfiguration struct {
 	URL           string
 	Header        http.Header
@@ -119,6 +125,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,
@@ -1932,7 +1933,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 {
@@ -1944,6 +1946,15 @@ func (s *SubscriptionSource) AsyncStart(ctx *resolve.Context, id uint64, input [
 	if options.Body.Query == "" {
 		return resolve.ErrUnableToResolve
 	}
+	for _, fn := range s.onSubscriptionStartFns {
+		events, err := fn(ctx)
+		if err != nil {
+			return err
+		}
+		for _, event := range events {
+			updater.Update(event)
+		}
+	}
 	return s.client.SubscribeAsync(ctx, id, options, updater)
 }
 
@@ -1963,6 +1974,15 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r
 	if options.Body.Query == "" {
 		return resolve.ErrUnableToResolve
 	}
+	for _, fn := range s.onSubscriptionStartFns {
+		events, err := fn(ctx)
+		if err != nil {
+			return err
+		}
+		for _, event := range events {
+			updater.Update(event)
+		}
+	}
 	return s.client.Subscribe(ctx, options, updater)
 }
 

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,174 @@ func TestSanitizeKey(t *testing.T) {
 	}
 }
 
+func TestSubscriptionSource_OnSubscriptionStartFns(t *testing.T) {
+	chatServer := httptest.NewServer(subscriptiontesting.ChatGraphQLEndpointHandler())
+	defer chatServer.Close()
+
+	sendChatMessage := func(t *testing.T, username, message string) {
+		time.Sleep(200 * time.Millisecond)
+		httpClient := http.Client{}
+		req, err := http.NewRequest(
+			http.MethodPost,
+			chatServer.URL,
+			bytes.NewBufferString(fmt.Sprintf(`{"variables": {}, "operationName": "SendMessage", "query": "mutation SendMessage { post(roomName: \"#test\", username: \"%s\", text: \"%s\") { id } }"}`, username, message)),
+		)
+		require.NoError(t, err)
+
+		req.Header.Set("Content-Type", "application/json")
+		resp, err := httpClient.Do(req)
+		require.NoError(t, err)
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+	}
+
+	chatServerSubscriptionOptions := func(t *testing.T, body string) []byte {
+		var gqlBody GraphQLBody
+		_ = json.Unmarshal([]byte(body), &gqlBody)
+		options := GraphQLSubscriptionOptions{
+			URL:    chatServer.URL,
+			Body:   gqlBody,
+			Header: nil,
+		}
+
+		optionsBytes, err := json.Marshal(options)
+		require.NoError(t, err)
+
+		return optionsBytes
+	}
+
+	newSubscriptionSource := func(ctx context.Context) SubscriptionSource {
+		httpClient := http.Client{}
+		subscriptionSource := SubscriptionSource{client: NewGraphQLSubscriptionClient(&httpClient, http.DefaultClient, ctx)}
+		return subscriptionSource
+	}
+
+	t.Run("should call the onSubscriptionStartFns at subscription source Start", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		updater := &testSubscriptionUpdater{}
+
+		startFnCalled := make(chan *resolve.Context, 1)
+		source := newSubscriptionSource(ctx.Context())
+		source.onSubscriptionStartFns = []OnSubscriptionStartFn{
+			func(ctx *resolve.Context) ([][]byte, error) {
+				startFnCalled <- ctx
+				return nil, nil
+			},
+		}
+		chatSubscriptionOptions := chatServerSubscriptionOptions(t, `{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`)
+
+		err := source.Start(ctx, chatSubscriptionOptions, updater)
+		require.NoError(t, err)
+
+		username := "myuser"
+		message := "hello world!"
+		go sendChatMessage(t, username, message)
+		updater.AwaitUpdates(t, time.Second, 1)
+		assert.Len(t, updater.updates, 1)
+		assert.Equal(t, `{"data":{"messageAdded":{"text":"hello world!","createdBy":"myuser"}}}`, updater.updates[0])
+
+		select {
+		case <-startFnCalled:
+		case <-time.After(time.Second):
+			t.Fatalf("start fn not called")
+		}
+	})
+
+	t.Run("should call the onSubscriptionStartFns at subscription source AsyncStart", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		updater := &testSubscriptionUpdater{}
+
+		startFnCalled := make(chan *resolve.Context, 1)
+		source := newSubscriptionSource(ctx.Context())
+		source.onSubscriptionStartFns = []OnSubscriptionStartFn{
+			func(ctx *resolve.Context) ([][]byte, error) {
+				startFnCalled <- ctx
+				return nil, nil
+			},
+		}
+		chatSubscriptionOptions := chatServerSubscriptionOptions(t, `{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`)
+
+		err := source.AsyncStart(ctx, 0, chatSubscriptionOptions, updater)
+		require.NoError(t, err)
+
+		username := "myuser"
+		message := "hello world!"
+		go sendChatMessage(t, username, message)
+		updater.AwaitUpdates(t, time.Second, 1)
+		assert.Len(t, updater.updates, 1)
+		assert.Equal(t, `{"data":{"messageAdded":{"text":"hello world!","createdBy":"myuser"}}}`, updater.updates[0])
+
+		select {
+		case <-startFnCalled:
+		case <-time.After(time.Second):
+			t.Fatalf("start fn not called")
+		}
+	})
+
+	t.Run("if the onSubscriptionStartFns returns an error, the subscription should not async start and return the error", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		updater := &testSubscriptionUpdater{}
+
+		source := newSubscriptionSource(ctx.Context())
+		source.onSubscriptionStartFns = []OnSubscriptionStartFn{
+			func(ctx *resolve.Context) ([][]byte, error) {
+				return nil, errors.New("test error")
+			},
+		}
+		chatSubscriptionOptions := chatServerSubscriptionOptions(t, `{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`)
+
+		err := source.AsyncStart(ctx, 0, chatSubscriptionOptions, updater)
+		require.Error(t, err)
+		require.ErrorContains(t, err, "test error")
+		assert.Len(t, updater.updates, 0)
+	})
+
+	t.Run("if the onSubscriptionStartFns returns an error, the subscription should not start and return the error", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		updater := &testSubscriptionUpdater{}
+
+		source := newSubscriptionSource(ctx.Context())
+		source.onSubscriptionStartFns = []OnSubscriptionStartFn{
+			func(ctx *resolve.Context) ([][]byte, error) {
+				return nil, errors.New("test error")
+			},
+		}
+		chatSubscriptionOptions := chatServerSubscriptionOptions(t, `{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`)
+
+		err := source.Start(ctx, chatSubscriptionOptions, updater)
+		require.Error(t, err)
+		require.ErrorContains(t, err, "test error")
+		assert.Len(t, updater.updates, 0)
+	})
+
+	t.Run("if the onSubscriptionStartFns returns a message, the message should be sent to the client", func(t *testing.T) {
+		ctx := resolve.NewContext(context.Background())
+		defer ctx.Context().Done()
+
+		updater := &testSubscriptionUpdater{}
+
+		source := newSubscriptionSource(ctx.Context())
+		source.onSubscriptionStartFns = []OnSubscriptionStartFn{
+			func(ctx *resolve.Context) ([][]byte, error) {
+				return [][]byte{[]byte(`{"data":{"messageAdded":{"text":"hello world!","createdBy":"injection"}}}`)}, nil
+			},
+		}
+		chatSubscriptionOptions := chatServerSubscriptionOptions(t, `{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`)
+
+		err := source.AsyncStart(ctx, 0, chatSubscriptionOptions, updater)
+		require.NoError(t, err)
+		assert.Len(t, updater.updates, 1)
+		assert.Equal(t, `{"data":{"messageAdded":{"text":"hello world!","createdBy":"injection"}}}`, updater.updates[0])
+	})
+}
+
 const interfaceSelectionSchema = `
 
 scalar String