feat: allows hook in the subscriptions

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
 }
 
+// SubscriptionOnStartFn 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 an error.
+type SubscriptionOnStartFn func(ctx resolve.StartupHookContext, input []byte) (err error)
+
 type SubscriptionConfiguration struct {
 	URL           string
 	Header        http.Header
@@ -119,6 +125,8 @@ type SubscriptionConfiguration struct {
 	// these headers by itself.
 	ForwardedClientHeaderRegularExpressions []RegularExpression
 	WsSubProtocol                           string
+	// StartupHooks contains the method called when a subscription is started
+	StartupHooks []SubscriptionOnStartFn
 }
 
 type FetchConfiguration struct {

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

@@ -449,7 +449,8 @@ func (p *Planner[T]) ConfigureSubscription() plan.SubscriptionConfiguration {
 	return plan.SubscriptionConfiguration{
 		Input: string(input),
 		DataSource: &SubscriptionSource{
-			client: p.subscriptionClient,
+			client:                 p.subscriptionClient,
+			subscriptionOnStartFns: p.config.subscription.StartupHooks,
 		},
 		Variables:      p.variables,
 		PostProcessing: DefaultPostProcessingConfiguration,
@@ -1953,7 +1954,8 @@ type RegularExpression struct {
 }
 
 type SubscriptionSource struct {
-	client GraphQLSubscriptionClient
+	client                 GraphQLSubscriptionClient
+	subscriptionOnStartFns []SubscriptionOnStartFn
 }
 
 func (s *SubscriptionSource) AsyncStart(ctx *resolve.Context, id uint64, input []byte, updater resolve.SubscriptionUpdater) error {
@@ -2003,3 +2005,16 @@ func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte,
 	}
 	return s.client.UniqueRequestID(ctx, options, xxh)
 }
+
+// SubscriptionOnStart is called when a subscription is started.
+// Each hook is called in a separate goroutine.
+func (s *SubscriptionSource) SubscriptionOnStart(ctx resolve.StartupHookContext, input []byte) error {
+	for _, fn := range s.subscriptionOnStartFns {
+		err := fn(ctx, input)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

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

@@ -4006,7 +4006,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,
 				},
@@ -8239,7 +8239,6 @@ func (f *FailingSubscriptionClient) SubscribeAsync(ctx *resolve.Context, id uint
 }
 
 func (f *FailingSubscriptionClient) Unsubscribe(id uint64) {
-
 }
 
 func (f *FailingSubscriptionClient) Subscribe(ctx *resolve.Context, options GraphQLSubscriptionOptions, updater resolve.SubscriptionUpdater) error {
@@ -8272,6 +8271,19 @@ func (t *testSubscriptionUpdaterChan) Update(data []byte) {
 	t.updates <- string(data)
 }
 
+// empty method to satisfy the interface, not used in this tests
+func (t *testSubscriptionUpdaterChan) UpdateSubscription(id resolve.SubscriptionIdentifier, data []byte) {
+}
+
+// empty method to satisfy the interface, not used in this tests
+func (t *testSubscriptionUpdaterChan) CloseSubscription(kind resolve.SubscriptionCloseKind, id resolve.SubscriptionIdentifier) {
+}
+
+// empty method to satisfy the interface, not used in this tests
+func (t *testSubscriptionUpdaterChan) Subscriptions() map[context.Context]resolve.SubscriptionIdentifier {
+	return make(map[context.Context]resolve.SubscriptionIdentifier)
+}
+
 func (t *testSubscriptionUpdaterChan) Complete() {
 	close(t.complete)
 }
@@ -8397,6 +8409,19 @@ func (t *testSubscriptionUpdater) Close(kind resolve.SubscriptionCloseKind) {
 	t.closed = true
 }
 
+// empty method to satisfy the interface, not used in this tests
+func (t *testSubscriptionUpdater) CloseSubscription(kind resolve.SubscriptionCloseKind, id resolve.SubscriptionIdentifier) {
+}
+
+// empty method to satisfy the interface, not used in this tests
+func (t *testSubscriptionUpdater) Subscriptions() map[context.Context]resolve.SubscriptionIdentifier {
+	return make(map[context.Context]resolve.SubscriptionIdentifier)
+}
+
+// empty method to satisfy the interface, not used in this tests
+func (t *testSubscriptionUpdater) UpdateSubscription(id resolve.SubscriptionIdentifier, data []byte) {
+}
+
 func TestSubscriptionSource_Start(t *testing.T) {
 	chatServer := httptest.NewServer(subscriptiontesting.ChatGraphQLEndpointHandler())
 	defer chatServer.Close()
@@ -8970,6 +8995,60 @@ func TestSanitizeKey(t *testing.T) {
 	}
 }
 
+func TestSubscriptionSource_SubscriptionOnStart(t *testing.T) {
+
+	t.Run("SubscriptionOnStart calls subscriptionOnStartFns", func(t *testing.T) {
+		ctx := resolve.StartupHookContext{
+			Context: context.Background(),
+			Updater: func(data []byte) {},
+		}
+
+		type fnData struct {
+			ctx   resolve.StartupHookContext
+			input []byte
+		}
+
+		startFnCalled := make(chan fnData, 1)
+		subscriptionSource := SubscriptionSource{
+			subscriptionOnStartFns: []SubscriptionOnStartFn{
+				func(ctx resolve.StartupHookContext, input []byte) error {
+					startFnCalled <- fnData{ctx, input}
+					return nil
+				},
+			},
+		}
+
+		err := subscriptionSource.SubscriptionOnStart(ctx, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`))
+		require.NoError(t, err)
+		var called fnData
+		select {
+		case called = <-startFnCalled:
+		case <-time.After(1 * time.Second):
+			t.Fatal("SubscriptionOnStartFn was not called")
+		}
+		assert.Equal(t, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`), called.input)
+	})
+
+	t.Run("SubscriptionOnStart calls subscriptionOnStartFns and returns error if one of the functions returns an error", func(t *testing.T) {
+		ctx := resolve.StartupHookContext{
+			Context: context.Background(),
+			Updater: func(data []byte) {},
+		}
+
+		subscriptionSource := SubscriptionSource{
+			subscriptionOnStartFns: []SubscriptionOnStartFn{
+				func(ctx resolve.StartupHookContext, input []byte) error {
+					return errors.New("test error")
+				},
+			},
+		}
+
+		err := subscriptionSource.SubscriptionOnStart(ctx, []byte(`{"variables": {}, "extensions": {}, "operationName": "LiveMessages", "query": "subscription LiveMessages { messageAdded(roomName: \"#test\") { text createdBy } }"}`))
+		require.Error(t, err)
+		assert.ErrorContains(t, err, "test error")
+	})
+}
+
 const interfaceSelectionSchema = `
 
 scalar String