From d1ae694c6f0c5e7ff9fa6b3b898012e83cec173d Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 3 Mar 2025 13:01:11 +0100 Subject: [PATCH 01/49] chore: initial and incomplete refactor --- router/core/executor.go | 7 +- router/core/factoryresolver.go | 79 +----- router/core/graph_server.go | 206 ++++++-------- router/core/plan_generator.go | 24 -- router/pkg/pubsub/datasource/datasource.go | 71 +++++ router/pkg/pubsub/datasource/old.go | 290 ++++++++++++++++++++ router/pkg/pubsub/datasource/providers.go | 1 + router/pkg/pubsub/kafka/kafka_datasource.go | 122 ++++++++ router/pkg/pubsub/nats/nats_datasource.go | 124 +++++++++ 9 files changed, 697 insertions(+), 227 deletions(-) create mode 100644 router/pkg/pubsub/datasource/datasource.go create mode 100644 router/pkg/pubsub/datasource/old.go create mode 100644 router/pkg/pubsub/datasource/providers.go create mode 100644 router/pkg/pubsub/kafka/kafka_datasource.go create mode 100644 router/pkg/pubsub/nats/nats_datasource.go diff --git a/router/core/executor.go b/router/core/executor.go index 3344caa83a..d3e35b7fea 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -51,7 +51,6 @@ type ExecutorBuildOptions struct { EngineConfig *nodev1.EngineConfiguration Subgraphs []*nodev1.Subgraph RouterEngineConfig *RouterEngineConfiguration - PubSubProviders *EnginePubSubProviders Reporter resolve.Reporter ApolloCompatibilityFlags config.ApolloCompatibilityFlags ApolloRouterCompatibilityFlags config.ApolloRouterCompatibilityFlags @@ -59,7 +58,7 @@ type ExecutorBuildOptions struct { } func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *ExecutorBuildOptions) (*Executor, error) { - planConfig, err := b.buildPlannerConfiguration(ctx, opts.EngineConfig, opts.Subgraphs, opts.RouterEngineConfig, opts.PubSubProviders) + planConfig, err := b.buildPlannerConfiguration(ctx, opts.EngineConfig, opts.Subgraphs, opts.RouterEngineConfig) if err != nil { return nil, fmt.Errorf("failed to build planner configuration: %w", err) } @@ -268,7 +267,7 @@ func buildKafkaOptions(eventSource config.KafkaEventSource) ([]kgo.Opt, error) { return opts, nil } -func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineCfg *RouterEngineConfiguration, pubSubProviders *EnginePubSubProviders) (*plan.Configuration, error) { +func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineCfg *RouterEngineConfiguration) (*plan.Configuration, error) { // this loader is used to take the engine config and create a plan config // the plan config is what the engine uses to turn a GraphQL Request into an execution plan // the plan config is stateful as it carries connection pools and other things @@ -280,8 +279,6 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con b.logger, routerEngineCfg.Execution.EnableSingleFlight, routerEngineCfg.Execution.EnableNetPoll, - pubSubProviders.nats, - pubSubProviders.kafka, )) // this generates the plan config using the data source factories from the config package diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 3b0dd558bc..1e44cda7f9 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -8,9 +8,9 @@ import ( "net/url" "slices" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" - "github.com/buger/jsonparser" + pubsubDatasource "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" "github.com/wundergraph/cosmo/router/pkg/config" @@ -18,7 +18,6 @@ import ( "go.uber.org/zap" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/graphql_datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/staticdatasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" @@ -35,7 +34,6 @@ type Loader struct { type FactoryResolver interface { ResolveGraphqlFactory(subgraphName string) (plan.PlannerFactory[graphql_datasource.Configuration], error) ResolveStaticFactory() (plan.PlannerFactory[staticdatasource.Configuration], error) - ResolvePubsubFactory() (plan.PlannerFactory[pubsub_datasource.Configuration], error) } type ApiTransportFactory interface { @@ -49,7 +47,6 @@ type DefaultFactoryResolver struct { transportOptions *TransportOptions static *staticdatasource.Factory[staticdatasource.Configuration] - pubsub *pubsub_datasource.Factory[pubsub_datasource.Configuration] log *zap.Logger engineCtx context.Context @@ -68,8 +65,6 @@ func NewDefaultFactoryResolver( log *zap.Logger, enableSingleFlight bool, enableNetPoll bool, - natsPubSubBySourceID map[string]pubsub_datasource.NatsPubSub, - kafkaPubSubBySourceID map[string]pubsub_datasource.KafkaPubSub, ) *DefaultFactoryResolver { transportFactory := NewTransport(transportOptions) @@ -105,7 +100,6 @@ func NewDefaultFactoryResolver( transportFactory: transportFactory, transportOptions: transportOptions, static: &staticdatasource.Factory[staticdatasource.Configuration]{}, - pubsub: pubsub_datasource.NewFactory(ctx, natsPubSubBySourceID, kafkaPubSubBySourceID), log: log, factoryLogger: factoryLogger, engineCtx: ctx, @@ -139,10 +133,6 @@ func (d *DefaultFactoryResolver) ResolveStaticFactory() (factory plan.PlannerFac return d.static, nil } -func (d *DefaultFactoryResolver) ResolvePubsubFactory() (factory plan.PlannerFactory[pubsub_datasource.Configuration], err error) { - return d.pubsub, nil -} - func NewLoader(includeInfo bool, resolver FactoryResolver) *Loader { return &Loader{ resolver: resolver, @@ -394,72 +384,11 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod } case nodev1.DataSourceKind_PUBSUB: - var eventConfigurations []pubsub_datasource.EventConfiguration - - for _, eventConfiguration := range in.GetCustomEvents().GetNats() { - eventType, err := pubsub_datasource.EventTypeFromString(eventConfiguration.EngineEventConfiguration.Type.String()) - if err != nil { - return nil, fmt.Errorf("invalid event type %q for data source %q: %w", eventConfiguration.EngineEventConfiguration.Type.String(), in.Id, err) - } - - var streamConfiguration *pubsub_datasource.NatsStreamConfiguration - if eventConfiguration.StreamConfiguration != nil { - streamConfiguration = &pubsub_datasource.NatsStreamConfiguration{ - Consumer: eventConfiguration.StreamConfiguration.GetConsumerName(), - StreamName: eventConfiguration.StreamConfiguration.GetStreamName(), - ConsumerInactiveThreshold: eventConfiguration.StreamConfiguration.GetConsumerInactiveThreshold(), - } - } - - eventConfigurations = append(eventConfigurations, pubsub_datasource.EventConfiguration{ - Metadata: &pubsub_datasource.EventMetadata{ - ProviderID: eventConfiguration.EngineEventConfiguration.GetProviderId(), - Type: eventType, - TypeName: eventConfiguration.EngineEventConfiguration.GetTypeName(), - FieldName: eventConfiguration.EngineEventConfiguration.GetFieldName(), - }, - Configuration: &pubsub_datasource.NatsEventConfiguration{ - StreamConfiguration: streamConfiguration, - Subjects: eventConfiguration.GetSubjects(), - }, - }) - } - - for _, eventConfiguration := range in.GetCustomEvents().GetKafka() { - eventType, err := pubsub_datasource.EventTypeFromString(eventConfiguration.EngineEventConfiguration.Type.String()) - if err != nil { - return nil, fmt.Errorf("invalid event type %q for data source %q: %w", eventConfiguration.EngineEventConfiguration.Type.String(), in.Id, err) - } - - eventConfigurations = append(eventConfigurations, pubsub_datasource.EventConfiguration{ - Metadata: &pubsub_datasource.EventMetadata{ - ProviderID: eventConfiguration.EngineEventConfiguration.GetProviderId(), - Type: eventType, - TypeName: eventConfiguration.EngineEventConfiguration.GetTypeName(), - FieldName: eventConfiguration.EngineEventConfiguration.GetFieldName(), - }, - Configuration: &pubsub_datasource.KafkaEventConfiguration{ - Topics: eventConfiguration.GetTopics(), - }, - }) - } - - factory, err := l.resolver.ResolvePubsubFactory() + var err error + out, err = pubsubDatasource.GetDataSourcesFromConfig(context.Background(), in, l.dataSourceMetaData(in), routerEngineConfig.Events) if err != nil { return nil, err } - - out, err = plan.NewDataSourceConfiguration[pubsub_datasource.Configuration]( - in.Id, - factory, - l.dataSourceMetaData(in), - pubsub_datasource.Configuration{ - Events: eventConfigurations, - }, - ) - if err != nil { - return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) - } default: return nil, fmt.Errorf("unknown data source type %q", in.Kind) } diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 0f76ffd784..2f985bfbf9 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -18,8 +18,6 @@ import ( "github.com/golang-jwt/jwt/v5" "github.com/klauspost/compress/gzhttp" "github.com/klauspost/compress/gzip" - "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" "go.opentelemetry.io/otel/attribute" otelmetric "go.opentelemetry.io/otel/metric" @@ -30,8 +28,6 @@ import ( "go.uber.org/zap/zapcore" "golang.org/x/exp/maps" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" - "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/common" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" rjwt "github.com/wundergraph/cosmo/router/internal/jwt" @@ -46,9 +42,6 @@ import ( "github.com/wundergraph/cosmo/router/pkg/logging" rmetric "github.com/wundergraph/cosmo/router/pkg/metric" "github.com/wundergraph/cosmo/router/pkg/otel" - "github.com/wundergraph/cosmo/router/pkg/pubsub" - "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" - pubsubNats "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/router/pkg/statistics" rtrace "github.com/wundergraph/cosmo/router/pkg/trace" ) @@ -65,11 +58,6 @@ type ( HealthChecks() health.Checker } - EnginePubSubProviders struct { - nats map[string]pubsub_datasource.NatsPubSub - kafka map[string]pubsub_datasource.KafkaPubSub - } - // graphServer is the swappable implementation of a Graph instance which is an HTTP mux with middlewares. // Everytime a schema is updated, the old graph server is shutdown and a new graph server is created. // For feature flags, a graphql server has multiple mux and is dynamically switched based on the feature flag header or cookie. @@ -78,7 +66,6 @@ type ( *Config context context.Context cancelFunc context.CancelFunc - pubSubProviders *EnginePubSubProviders engineStats statistics.EngineStatistics playgroundHandler func(http.Handler) http.Handler publicKey *ecdsa.PublicKey @@ -126,10 +113,6 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC graphMuxList: make([]*graphMux, 0, 1), routerListenAddr: r.listenAddr, hostName: r.hostName, - pubSubProviders: &EnginePubSubProviders{ - nats: map[string]pubsub_datasource.NatsPubSub{}, - kafka: map[string]pubsub_datasource.KafkaPubSub{}, - }, } baseOtelAttributes := []attribute.KeyValue{ @@ -880,10 +863,10 @@ func (s *graphServer) buildGraphMux(ctx context.Context, SubgraphErrorPropagation: s.subgraphErrorPropagation, } - err = s.buildPubSubConfiguration(ctx, engineConfig, routerEngineConfig) - if err != nil { - return nil, fmt.Errorf("failed to build pubsub configuration: %w", err) - } + //err = s.buildPubSubConfiguration(ctx, engineConfig, routerEngineConfig) + //if err != nil { + // return nil, fmt.Errorf("failed to build pubsub configuration: %w", err) + //} ecb := &ExecutorConfigurationBuilder{ introspection: s.introspection, @@ -919,7 +902,6 @@ func (s *graphServer) buildGraphMux(ctx context.Context, EngineConfig: engineConfig, Subgraphs: configSubgraphs, RouterEngineConfig: routerEngineConfig, - PubSubProviders: s.pubSubProviders, Reporter: s.engineStats, ApolloCompatibilityFlags: s.apolloCompatibilityFlags, ApolloRouterCompatibilityFlags: s.apolloRouterCompatibilityFlags, @@ -1170,85 +1152,85 @@ func (s *graphServer) buildGraphMux(ctx context.Context, return gm, nil } -func (s *graphServer) buildPubSubConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, routerEngineCfg *RouterEngineConfiguration) error { - datasourceConfigurations := engineConfig.GetDatasourceConfigurations() - for _, datasourceConfiguration := range datasourceConfigurations { - if datasourceConfiguration.CustomEvents == nil { - continue - } - - for _, eventConfiguration := range datasourceConfiguration.GetCustomEvents().GetNats() { - - providerID := eventConfiguration.EngineEventConfiguration.GetProviderId() - // if this source name's provider has already been initiated, do not try to initiate again - _, ok := s.pubSubProviders.nats[providerID] - if ok { - continue - } - - for _, eventSource := range routerEngineCfg.Events.Providers.Nats { - if eventSource.ID == eventConfiguration.EngineEventConfiguration.GetProviderId() { - options, err := buildNatsOptions(eventSource, s.logger) - if err != nil { - return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", providerID, err) - } - natsConnection, err := nats.Connect(eventSource.URL, options...) - if err != nil { - return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", providerID, err) - } - js, err := jetstream.New(natsConnection) - if err != nil { - return err - } - - s.pubSubProviders.nats[providerID] = pubsubNats.NewConnector(s.logger, natsConnection, js, s.hostName, s.routerListenAddr).New(ctx) - - break - } - } - - _, ok = s.pubSubProviders.nats[providerID] - if !ok { - return fmt.Errorf("failed to find Nats provider with ID \"%s\". Ensure the provider definition is part of the config", providerID) - } - } - - for _, eventConfiguration := range datasourceConfiguration.GetCustomEvents().GetKafka() { - - providerID := eventConfiguration.EngineEventConfiguration.GetProviderId() - // if this source name's provider has already been initiated, do not try to initiate again - _, ok := s.pubSubProviders.kafka[providerID] - if ok { - continue - } - - for _, eventSource := range routerEngineCfg.Events.Providers.Kafka { - if eventSource.ID == providerID { - options, err := buildKafkaOptions(eventSource) - if err != nil { - return fmt.Errorf("failed to build options for Kafka provider with ID \"%s\": %w", providerID, err) - } - ps, err := kafka.NewConnector(s.logger, options) - if err != nil { - return fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", providerID, err) - } - - s.pubSubProviders.kafka[providerID] = ps.New(ctx) - - break - } - } - - _, ok = s.pubSubProviders.kafka[providerID] - if !ok { - return fmt.Errorf("failed to find Kafka provider with ID \"%s\". Ensure the provider definition is part of the config", providerID) - } - } - - } - - return nil -} +//func (s *graphServer) buildPubSubConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, routerEngineCfg *RouterEngineConfiguration) error { +// datasourceConfigurations := engineConfig.GetDatasourceConfigurations() +// for _, datasourceConfiguration := range datasourceConfigurations { +// if datasourceConfiguration.CustomEvents == nil { +// continue +// } +// +// for _, eventConfiguration := range datasourceConfiguration.GetCustomEvents().GetNats() { +// +// providerID := eventConfiguration.EngineEventConfiguration.GetProviderId() +// // if this source name's provider has already been initiated, do not try to initiate again +// _, ok := s.pubSubProviders.nats[providerID] +// if ok { +// continue +// } +// +// for _, eventSource := range routerEngineCfg.Events.Providers.Nats { +// if eventSource.ID == eventConfiguration.EngineEventConfiguration.GetProviderId() { +// options, err := buildNatsOptions(eventSource, s.logger) +// if err != nil { +// return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", providerID, err) +// } +// natsConnection, err := nats.Connect(eventSource.URL, options...) +// if err != nil { +// return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", providerID, err) +// } +// js, err := jetstream.New(natsConnection) +// if err != nil { +// return err +// } +// +// s.pubSubProviders.nats[providerID] = pubsubNats.NewConnector(s.logger, natsConnection, js, s.hostName, s.routerListenAddr).New(ctx) +// +// break +// } +// } +// +// _, ok = s.pubSubProviders.nats[providerID] +// if !ok { +// return fmt.Errorf("failed to find Nats provider with ID \"%s\". Ensure the provider definition is part of the config", providerID) +// } +// } +// +// for _, eventConfiguration := range datasourceConfiguration.GetCustomEvents().GetKafka() { +// +// providerID := eventConfiguration.EngineEventConfiguration.GetProviderId() +// // if this source name's provider has already been initiated, do not try to initiate again +// _, ok := s.pubSubProviders.kafka[providerID] +// if ok { +// continue +// } +// +// for _, eventSource := range routerEngineCfg.Events.Providers.Kafka { +// if eventSource.ID == providerID { +// options, err := buildKafkaOptions(eventSource) +// if err != nil { +// return fmt.Errorf("failed to build options for Kafka provider with ID \"%s\": %w", providerID, err) +// } +// ps, err := kafka.NewConnector(s.logger, options) +// if err != nil { +// return fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", providerID, err) +// } +// +// s.pubSubProviders.kafka[providerID] = ps.New(ctx) +// +// break +// } +// } +// +// _, ok = s.pubSubProviders.kafka[providerID] +// if !ok { +// return fmt.Errorf("failed to find Kafka provider with ID \"%s\". Ensure the provider definition is part of the config", providerID) +// } +// } +// +// } +// +// return nil +//} // wait waits for all in-flight requests to finish. Similar to http.Server.Shutdown we wait in intervals + jitter // to make the shutdown process more efficient. @@ -1327,28 +1309,6 @@ func (s *graphServer) Shutdown(ctx context.Context) error { } } - if s.pubSubProviders != nil { - - s.logger.Debug("Shutting down pubsub providers") - - for _, pubSub := range s.pubSubProviders.nats { - if p, ok := pubSub.(pubsub.Lifecycle); ok { - if err := p.Shutdown(ctx); err != nil { - s.logger.Error("Failed to shutdown Nats pubsub provider", zap.Error(err)) - finalErr = errors.Join(finalErr, err) - } - } - } - for _, pubSub := range s.pubSubProviders.kafka { - if p, ok := pubSub.(pubsub.Lifecycle); ok { - if err := p.Shutdown(ctx); err != nil { - s.logger.Error("Failed to shutdown Kafka pubsub provider", zap.Error(err)) - finalErr = errors.Join(finalErr, err) - } - } - } - } - // Shutdown all graphs muxes to release resources // e.g. planner cache s.graphMuxListLock.Lock() diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index ecaa3dd500..33a52b504e 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -8,14 +8,12 @@ import ( "os" log "github.com/jensneuse/abstractlogger" - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/astnormalization" "github.com/wundergraph/graphql-go-tools/v2/pkg/astparser" "github.com/wundergraph/graphql-go-tools/v2/pkg/asttransform" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/graphql_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/introspection_datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/postprocess" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" @@ -114,27 +112,6 @@ func (pg *PlanGenerator) loadConfiguration(configFilePath string) error { return err } - natSources := map[string]pubsub_datasource.NatsPubSub{} - kafkaSources := map[string]pubsub_datasource.KafkaPubSub{} - for _, ds := range routerConfig.GetEngineConfig().GetDatasourceConfigurations() { - if ds.GetKind() != nodev1.DataSourceKind_PUBSUB || ds.GetCustomEvents() == nil { - continue - } - for _, natConfig := range ds.GetCustomEvents().GetNats() { - providerId := natConfig.GetEngineEventConfiguration().GetProviderId() - if _, ok := natSources[providerId]; !ok { - natSources[providerId] = nil - } - } - for _, kafkaConfig := range ds.GetCustomEvents().GetKafka() { - providerId := kafkaConfig.GetEngineEventConfiguration().GetProviderId() - if _, ok := kafkaSources[providerId]; !ok { - kafkaSources[providerId] = nil - } - } - } - pubSubFactory := pubsub_datasource.NewFactory(context.Background(), natSources, kafkaSources) - var netPollConfig graphql_datasource.NetPollConfiguration netPollConfig.ApplyDefaults() @@ -152,7 +129,6 @@ func (pg *PlanGenerator) loadConfiguration(configFilePath string) error { streamingClient: http.DefaultClient, subscriptionClient: subscriptionClient, transportOptions: &TransportOptions{SubgraphTransportOptions: NewSubgraphTransportOptions(config.TrafficShapingRules{})}, - pubsub: pubSubFactory, }) // this generates the plan configuration using the data source factories from the config package diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go new file mode 100644 index 0000000000..db7d8c9f0c --- /dev/null +++ b/router/pkg/pubsub/datasource/datasource.go @@ -0,0 +1,71 @@ +package datasource + +import ( + "context" + "fmt" + + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" +) + +type PubSubImplementer[F any] interface { + VerifyConfig(in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error + GetFactory(executionContext context.Context, config config.EventsConfiguration) F +} + +func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) (plan.DataSource, error) { + var pubSubs []any + if natsData := in.GetCustomEvents().GetNats(); natsData != nil { + n := nats.NewPubSub() + err := n.VerifyConfig(in, dsMeta, config) + if err != nil { + return nil, err + } + factory := nats.NewFactory(ctx, config) + ds, err := plan.NewDataSourceConfiguration[nats.Configuration]( + in.Id, + factory, + dsMeta, + nats.Configuration{}, + ) + if err != nil { + return nil, err + } + pubSubs = append(pubSubs, ds) + } + if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { + k := kafka.NewPubSub() + err := k.VerifyConfig(in, dsMeta, config) + if err != nil { + return nil, err + } + factory := k.GetFactory(ctx, config) + ds, err := plan.NewDataSourceConfiguration[kafka.Configuration]( + in.Id, + factory, + dsMeta, + kafka.Configuration{}, + ) + + if err != nil { + return nil, err + } + + pubSubs = append(pubSubs, ds) + } + + if len(pubSubs) == 0 { + return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) + } + + factory := NewFactory(ctx, pubSubs) + return plan.NewDataSourceConfiguration[Configuration]( + in.Id, + factory, + dsMeta, + Configuration{}, + ) +} diff --git a/router/pkg/pubsub/datasource/old.go b/router/pkg/pubsub/datasource/old.go new file mode 100644 index 0000000000..d1f53bb21c --- /dev/null +++ b/router/pkg/pubsub/datasource/old.go @@ -0,0 +1,290 @@ +package datasource + +import ( + "context" + "fmt" + "regexp" + "strings" + + "github.com/jensneuse/abstractlogger" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type EventType string + +const ( + EventTypePublish EventType = "publish" + EventTypeRequest EventType = "request" + EventTypeSubscribe EventType = "subscribe" +) + +var eventSubjectRegex = regexp.MustCompile(`{{ args.([a-zA-Z0-9_]+) }}`) + +func EventTypeFromString(s string) (EventType, error) { + et := EventType(strings.ToLower(s)) + switch et { + case EventTypePublish, EventTypeRequest, EventTypeSubscribe: + return et, nil + default: + return "", fmt.Errorf("invalid event type: %q", s) + } +} + +type Configuration struct { + pubSubs []Factory +} + +type Planner struct { + id int + pubSubs []any + eventManager any + rootFieldRef int + variables resolve.Variables + visitor *plan.Visitor + config Configuration +} + +func (p *Planner) SetID(id int) { + p.id = id +} + +func (p *Planner) ID() (id int) { + return p.id +} + +func (p *Planner) EnterField(ref int) { + if p.rootFieldRef != -1 { + // This is a nested field; nothing needs to be done + return + } + p.rootFieldRef = ref + + fieldName := p.visitor.Operation.FieldNameString(ref) + typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("nope fieldName %s and typeName %s", fieldName, typeName)) + + //switch v := eventConfig.Configuration.(type) { + //case *NatsEventConfiguration: + // em := &NatsEventManager{ + // visitor: p.visitor, + // variables: &p.variables, + // eventMetadata: *eventConfig.Metadata, + // eventConfiguration: v, + // } + // p.eventManager = em + // + // switch eventConfig.Metadata.Type { + // case EventTypePublish, EventTypeRequest: + // em.handlePublishAndRequestEvent(ref) + // case EventTypeSubscribe: + // em.handleSubscriptionEvent(ref) + // default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Nats", eventConfig.Metadata.Type)) + // } + //case *KafkaEventConfiguration: + // em := &KafkaEventManager{ + // visitor: p.visitor, + // variables: &p.variables, + // eventMetadata: *eventConfig.Metadata, + // eventConfiguration: v, + // } + // p.eventManager = em + // + // switch eventConfig.Metadata.Type { + // case EventTypePublish: + // em.handlePublishEvent(ref) + // case EventTypeSubscribe: + // em.handleSubscriptionEvent(ref) + // default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Kafka", eventConfig.Metadata.Type)) + // } + //default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid event configuration type: %T", v)) + //} +} + +func (p *Planner) EnterDocument(_, _ *ast.Document) { + p.rootFieldRef = -1 + p.eventManager = nil +} + +func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], dataSourcePlannerConfiguration plan.DataSourcePlannerConfiguration) error { + p.visitor = visitor + visitor.Walker.RegisterEnterFieldVisitor(p) + visitor.Walker.RegisterEnterDocumentVisitor(p) + p.config = configuration.CustomConfiguration() + return nil +} + +func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { + if p.eventManager == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event manager is nil")) + return resolve.FetchConfiguration{} + } + + //var dataSource resolve.DataSource + // + //switch v := p.eventManager.(type) { + //case *NatsEventManager: + // pubsub, ok := p.natsPubSubByProviderID[v.eventMetadata.ProviderID] + // if !ok { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) + // return resolve.FetchConfiguration{} + // } + // + // switch v.eventMetadata.Type { + // case EventTypePublish: + // dataSource = &NatsPublishDataSource{ + // pubSub: pubsub, + // } + // case EventTypeRequest: + // dataSource = &NatsRequestDataSource{ + // pubSub: pubsub, + // } + // default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Nats", v.eventMetadata.Type)) + // return resolve.FetchConfiguration{} + // } + // + // return resolve.FetchConfiguration{ + // Input: v.publishAndRequestEventConfiguration.MarshalJSONTemplate(), + // Variables: p.variables, + // DataSource: dataSource, + // PostProcessing: resolve.PostProcessingConfiguration{ + // MergePath: []string{v.eventMetadata.FieldName}, + // }, + // } + // + //case *KafkaEventManager: + // pubsub, ok := p.kafkaPubSubByProviderID[v.eventMetadata.ProviderID] + // if !ok { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) + // return resolve.FetchConfiguration{} + // } + // + // switch v.eventMetadata.Type { + // case EventTypePublish: + // dataSource = &KafkaPublishDataSource{ + // pubSub: pubsub, + // } + // case EventTypeRequest: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("event type \"%s\" is not supported for Kafka", v.eventMetadata.Type)) + // return resolve.FetchConfiguration{} + // default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", v.eventMetadata.Type)) + // return resolve.FetchConfiguration{} + // } + // + // return resolve.FetchConfiguration{ + // Input: v.publishEventConfiguration.MarshalJSONTemplate(), + // Variables: p.variables, + // DataSource: dataSource, + // PostProcessing: resolve.PostProcessingConfiguration{ + // MergePath: []string{v.eventMetadata.FieldName}, + // }, + // } + // + //default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event manager type: %T", p.eventManager)) + //} + + return resolve.FetchConfiguration{} +} + +func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { + if p.eventManager == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) + return plan.SubscriptionConfiguration{} + } + + //switch v := p.eventManager.(type) { + //case *NatsEventManager: + // pubsub, ok := p.natsPubSubByProviderID[v.eventMetadata.ProviderID] + // if !ok { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) + // return plan.SubscriptionConfiguration{} + // } + // object, err := json.Marshal(v.subscriptionEventConfiguration) + // if err != nil { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) + // return plan.SubscriptionConfiguration{} + // } + // return plan.SubscriptionConfiguration{ + // Input: string(object), + // Variables: p.variables, + // DataSource: &NatsSubscriptionSource{ + // pubSub: pubsub, + // }, + // PostProcessing: resolve.PostProcessingConfiguration{ + // MergePath: []string{v.eventMetadata.FieldName}, + // }, + // } + //case *KafkaEventManager: + // pubsub, ok := p.kafkaPubSubByProviderID[v.eventMetadata.ProviderID] + // if !ok { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) + // return plan.SubscriptionConfiguration{} + // } + // object, err := json.Marshal(v.subscriptionEventConfiguration) + // if err != nil { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) + // return plan.SubscriptionConfiguration{} + // } + // return plan.SubscriptionConfiguration{ + // Input: string(object), + // Variables: p.variables, + // DataSource: &KafkaSubscriptionSource{ + // pubSub: pubsub, + // }, + // PostProcessing: resolve.PostProcessingConfiguration{ + // MergePath: []string{v.eventMetadata.FieldName}, + // }, + // } + //default: + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: invalid event manager type: %T", p.eventManager)) + //} + + return plan.SubscriptionConfiguration{} +} + +func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { + return plan.DataSourcePlanningBehavior{ + MergeAliasedRootNodes: false, + OverrideFieldPathFromAlias: false, + IncludeTypeNameFields: true, + } +} + +func (p *Planner) DownstreamResponseFieldAlias(_ int) (alias string, exists bool) { + return "", false +} + +func NewFactory(executionContext context.Context, pubSubs []any) *Factory { + return &Factory{ + executionContext: executionContext, + pubSubs: pubSubs, + } +} + +type Factory struct { + executionContext context.Context + pubSubs []any +} + +func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { + return &Planner{ + pubSubs: f.pubSubs, + } +} + +func (f *Factory) Context() context.Context { + return f.executionContext +} + +func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { + return nil, false +} diff --git a/router/pkg/pubsub/datasource/providers.go b/router/pkg/pubsub/datasource/providers.go new file mode 100644 index 0000000000..0d7172bf8d --- /dev/null +++ b/router/pkg/pubsub/datasource/providers.go @@ -0,0 +1 @@ +package datasource diff --git a/router/pkg/pubsub/kafka/kafka_datasource.go b/router/pkg/pubsub/kafka/kafka_datasource.go new file mode 100644 index 0000000000..ae59f6bb97 --- /dev/null +++ b/router/pkg/pubsub/kafka/kafka_datasource.go @@ -0,0 +1,122 @@ +package kafka + +import ( + "bytes" + "context" + "fmt" + + "github.com/jensneuse/abstractlogger" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type Kafka struct{} + +func (k *Kafka) VerifyConfig(in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { + providers := map[string]bool{} + for _, provider := range config.Providers.Kafka { + providers[provider.ID] = true + } + for _, event := range in.CustomEvents.GetKafka() { + if !providers[event.EngineEventConfiguration.ProviderId] { + return fmt.Errorf("failed to find Kafka provider with ID %s", event.EngineEventConfiguration.ProviderId) + } + } + return nil +} + +func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { + return NewFactory(executionContext, config) +} + +func NewPubSub() Kafka { + return Kafka{} +} + +type Configuration struct { + Data string `json:"data"` +} + +type Planner struct { + id int + config Configuration + eventsConfig config.EventsConfiguration +} + +func (p *Planner) SetID(id int) { + p.id = id +} + +func (p *Planner) ID() (id int) { + return p.id +} + +func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { + // skip, not required + return +} + +func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { + return plan.DataSourcePlanningBehavior{ + MergeAliasedRootNodes: false, + OverrideFieldPathFromAlias: false, + } +} + +func (p *Planner) Register(_ *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { + p.config = configuration.CustomConfiguration() + return nil +} + +func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { + return resolve.FetchConfiguration{ + Input: p.config.Data, + DataSource: Source{}, + } +} + +func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { + return plan.SubscriptionConfiguration{ + Input: p.config.Data, + } +} + +type Source struct{} + +func (Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { + _, err = out.Write(input) + return +} + +func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { + panic("not implemented") +} + +func NewFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { + return &Factory{ + executionContext: executionContext, + config: config, + } +} + +type Factory struct { + executionContext context.Context + config config.EventsConfiguration +} + +func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { + return &Planner{} +} + +func (f *Factory) Context() context.Context { + return f.executionContext +} + +func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { + return nil, false +} diff --git a/router/pkg/pubsub/nats/nats_datasource.go b/router/pkg/pubsub/nats/nats_datasource.go new file mode 100644 index 0000000000..8c893a0567 --- /dev/null +++ b/router/pkg/pubsub/nats/nats_datasource.go @@ -0,0 +1,124 @@ +package nats + +import ( + "bytes" + "context" + "fmt" + + "github.com/jensneuse/abstractlogger" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type Nats struct{} + +func (n *Nats) VerifyConfig(in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { + providers := map[string]bool{} + for _, provider := range config.Providers.Nats { + providers[provider.ID] = true + } + for _, event := range in.CustomEvents.GetNats() { + if !providers[event.EngineEventConfiguration.ProviderId] { + return fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) + } + } + return nil +} + +func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { + return NewFactory(executionContext, config) +} + +func NewPubSub() Nats { + return Nats{} +} + +type Configuration struct { + Data string `json:"data"` +} + +type Planner struct { + id int + config Configuration +} + +func (p *Planner) SetID(id int) { + p.id = id +} + +func (p *Planner) ID() (id int) { + return p.id +} + +func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { + // skip, not required + return +} + +func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { + return plan.DataSourcePlanningBehavior{ + MergeAliasedRootNodes: false, + OverrideFieldPathFromAlias: false, + } +} + +func (p *Planner) Register(_ *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { + p.config = configuration.CustomConfiguration() + return nil +} + +func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { + return resolve.FetchConfiguration{ + Input: p.config.Data, + DataSource: Source{}, + } +} + +func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { + return plan.SubscriptionConfiguration{ + Input: p.config.Data, + } +} + +type Source struct{} + +func (Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { + _, err = out.Write(input) + return +} + +func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { + panic("not implemented") +} + +func NewFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { + return &Factory{ + executionContext: executionContext, + eventsConfiguration: config, + } +} + +type Factory struct { + config Configuration + eventsConfiguration config.EventsConfiguration + executionContext context.Context +} + +func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { + return &Planner{ + config: f.config, + } +} + +func (f *Factory) Context() context.Context { + return f.executionContext +} + +func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { + return nil, false +} From bca3e2bf6b998e35df240803927b09647a52e990 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 18 Mar 2025 16:41:14 +0100 Subject: [PATCH 02/49] feat(pubsub): refactor pubsub handling and add Kafka/NATS support --- router-tests/events/kafka_events_test.go | 64 +++++ router/core/errors.go | 4 +- router/core/executor.go | 79 +----- router/core/factoryresolver.go | 23 +- router/pkg/pubsub/datasource.go | 35 +++ router/pkg/pubsub/datasource/datasource.go | 68 +---- router/pkg/pubsub/{ => datasource}/error.go | 2 +- router/pkg/pubsub/datasource/old.go | 20 +- router/pkg/pubsub/datasource/providers.go | 1 - router/pkg/pubsub/kafka/kafka.go | 20 +- router/pkg/pubsub/kafka/kafka_datasource.go | 268 ++++++++++++++++-- .../pubsub/kafka/subscription_datasource.go | 87 ++++++ router/pkg/pubsub/nats/nats.go | 27 +- router/pkg/pubsub/nats/nats_datasource.go | 141 ++++++++- router/pkg/pubsub/utils/utils.go | 38 +++ 15 files changed, 673 insertions(+), 204 deletions(-) create mode 100644 router/pkg/pubsub/datasource.go rename router/pkg/pubsub/{ => datasource}/error.go (93%) delete mode 100644 router/pkg/pubsub/datasource/providers.go create mode 100644 router/pkg/pubsub/kafka/subscription_datasource.go create mode 100644 router/pkg/pubsub/utils/utils.go diff --git a/router-tests/events/kafka_events_test.go b/router-tests/events/kafka_events_test.go index 32e8ace611..36dd03b086 100644 --- a/router-tests/events/kafka_events_test.go +++ b/router-tests/events/kafka_events_test.go @@ -1050,6 +1050,70 @@ func TestKafkaEvents(t *testing.T) { xEnv.WaitForConnectionCount(0, KafkaWaitTimeout) }) }) + + t.Run("mutate", func(t *testing.T) { + t.Parallel() + + topics := []string{"employeeUpdated", "employeeUpdatedTwo"} + + testenv.Run(t, &testenv.Config{ + RouterConfigJSONTemplate: testenv.ConfigWithEdfsKafkaJSONTemplate, + EnableKafka: true, + }, func(t *testing.T, xEnv *testenv.Environment) { + ensureTopicExists(t, xEnv, topics...) + + var subscriptionOne struct { + employeeUpdatedMyKafka struct { + ID float64 `graphql:"id"` + Details struct { + Forename string `graphql:"forename"` + Surname string `graphql:"surname"` + } `graphql:"details"` + } `graphql:"employeeUpdatedMyKafka(employeeID: 3)"` + } + + surl := xEnv.GraphQLWebSocketSubscriptionURL() + client := graphql.NewSubscriptionClient(surl) + t.Cleanup(func() { + _ = client.Close() + }) + + var counter atomic.Uint32 + + subscriptionOneID, err := client.Subscribe(&subscriptionOne, nil, func(dataValue []byte, errValue error) error { + defer counter.Add(1) + require.NoError(t, errValue) + require.JSONEq(t, `{"updateEmployeeMyKafka":{"id":1,"details":{"forename":"Jens","surname":"Neuse"}}}`, string(dataValue)) + return nil + }) + require.NoError(t, err) + require.NotEmpty(t, subscriptionOneID) + + go func() { + clientErr := client.Run() + require.NoError(t, clientErr) + }() + + go func() { + require.Eventually(t, func() bool { + return counter.Load() == 1 + }, KafkaWaitTimeout, time.Millisecond*100) + _ = client.Close() + }() + + xEnv.WaitForSubscriptionCount(1, KafkaWaitTimeout) + + // Send a mutation to trigger the first subscription + resOne := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ + Query: `mutation { updateEmployeeMyKafka(employeeID: 3, update: {name: "name test"}) { success } }`, + }) + require.JSONEq(t, `{"data":{"updateEmployeeMyKafka":{"id":3}}}`, resOne.Body) + + xEnv.WaitForMessagesSent(1, KafkaWaitTimeout) + xEnv.WaitForSubscriptionCount(0, KafkaWaitTimeout) + xEnv.WaitForConnectionCount(0, KafkaWaitTimeout) + }) + }) } func TestFlakyKafkaEvents(t *testing.T) { diff --git a/router/core/errors.go b/router/core/errors.go index 118c7256a1..f770c52fed 100644 --- a/router/core/errors.go +++ b/router/core/errors.go @@ -12,7 +12,7 @@ import ( rErrors "github.com/wundergraph/cosmo/router/internal/errors" "github.com/wundergraph/cosmo/router/internal/persistedoperation" "github.com/wundergraph/cosmo/router/internal/unique" - "github.com/wundergraph/cosmo/router/pkg/pubsub" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" rtrace "github.com/wundergraph/cosmo/router/pkg/trace" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/graphql_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" @@ -73,7 +73,7 @@ func getErrorType(err error) errorType { return errorTypeContextTimeout } } - var edfsErr *pubsub.Error + var edfsErr *datasource.Error if errors.As(err, &edfsErr) { return errorTypeEDFS } diff --git a/router/core/executor.go b/router/core/executor.go index d3e35b7fea..d26203bba3 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -2,15 +2,10 @@ package core import ( "context" - "crypto/tls" - "errors" "fmt" "net/http" "time" - "github.com/nats-io/nats.go" - "github.com/twmb/franz-go/pkg/kgo" - "github.com/twmb/franz-go/pkg/sasl/plain" "go.uber.org/zap" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" @@ -195,78 +190,6 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor }, nil } -func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { - opts := []nats.Option{ - nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), - nats.ReconnectJitter(500*time.Millisecond, 2*time.Second), - nats.ClosedHandler(func(conn *nats.Conn) { - logger.Info("NATS connection closed", zap.String("provider_id", eventSource.ID), zap.Error(conn.LastError())) - }), - nats.ConnectHandler(func(nc *nats.Conn) { - logger.Info("NATS connection established", zap.String("provider_id", eventSource.ID), zap.String("url", nc.ConnectedUrlRedacted())) - }), - nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { - if err != nil { - logger.Error("NATS disconnected; will attempt to reconnect", zap.Error(err), zap.String("provider_id", eventSource.ID)) - } else { - logger.Info("NATS disconnected", zap.String("provider_id", eventSource.ID)) - } - }), - nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { - if errors.Is(err, nats.ErrSlowConsumer) { - logger.Warn( - "NATS slow consumer detected. Events are being dropped. Please consider increasing the buffer size or reducing the number of messages being sent.", - zap.Error(err), - zap.String("provider_id", eventSource.ID), - ) - } else { - logger.Error("NATS error", zap.Error(err)) - } - }), - nats.ReconnectHandler(func(conn *nats.Conn) { - logger.Info("NATS reconnected", zap.String("provider_id", eventSource.ID), zap.String("url", conn.ConnectedUrlRedacted())) - }), - } - - if eventSource.Authentication != nil { - if eventSource.Authentication.Token != nil { - opts = append(opts, nats.Token(*eventSource.Authentication.Token)) - } else if eventSource.Authentication.UserInfo.Username != nil && eventSource.Authentication.UserInfo.Password != nil { - opts = append(opts, nats.UserInfo(*eventSource.Authentication.UserInfo.Username, *eventSource.Authentication.UserInfo.Password)) - } - } - - return opts, nil -} - -// buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. -// Only general options like TLS, SASL, etc. are configured here. Specific options like topics, etc. are -// configured in the KafkaPubSub implementation. -func buildKafkaOptions(eventSource config.KafkaEventSource) ([]kgo.Opt, error) { - opts := []kgo.Opt{ - kgo.SeedBrokers(eventSource.Brokers...), - // Ensure proper timeouts are set - kgo.ProduceRequestTimeout(10 * time.Second), - kgo.ConnIdleTimeout(60 * time.Second), - } - - if eventSource.TLS != nil && eventSource.TLS.Enabled { - opts = append(opts, - // Configure TLS. Uses SystemCertPool for RootCAs by default. - kgo.DialTLSConfig(new(tls.Config)), - ) - } - - if eventSource.Authentication != nil && eventSource.Authentication.SASLPlain.Username != nil && eventSource.Authentication.SASLPlain.Password != nil { - opts = append(opts, kgo.SASL(plain.Auth{ - User: *eventSource.Authentication.SASLPlain.Username, - Pass: *eventSource.Authentication.SASLPlain.Password, - }.AsMechanism())) - } - - return opts, nil -} - func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineCfg *RouterEngineConfiguration) (*plan.Configuration, error) { // this loader is used to take the engine config and create a plan config // the plan config is what the engine uses to turn a GraphQL Request into an execution plan @@ -279,7 +202,7 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con b.logger, routerEngineCfg.Execution.EnableSingleFlight, routerEngineCfg.Execution.EnableNetPoll, - )) + ), b.logger) // this generates the plan config using the data source factories from the config package planConfig, err := loader.Load(engineConfig, subgraphs, routerEngineCfg) diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 1e44cda7f9..5ae818a7f3 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -9,7 +9,7 @@ import ( "slices" "github.com/buger/jsonparser" - pubsubDatasource "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" "github.com/wundergraph/cosmo/router/pkg/config" @@ -29,6 +29,7 @@ type Loader struct { resolver FactoryResolver // includeInfo controls whether additional information like type usage and field usage is included in the plan includeInfo bool + logger *zap.Logger } type FactoryResolver interface { @@ -133,10 +134,11 @@ func (d *DefaultFactoryResolver) ResolveStaticFactory() (factory plan.PlannerFac return d.static, nil } -func NewLoader(includeInfo bool, resolver FactoryResolver) *Loader { +func NewLoader(includeInfo bool, resolver FactoryResolver, logger *zap.Logger) *Loader { return &Loader{ resolver: resolver, includeInfo: includeInfo, + logger: logger, } } @@ -245,7 +247,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod } for _, in := range engineConfig.DatasourceConfigurations { - var out plan.DataSource + var outs []plan.DataSource switch in.Kind { case nodev1.DataSourceKind_STATIC: @@ -254,6 +256,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod return nil, err } + var out plan.DataSource out, err = plan.NewDataSourceConfiguration[staticdatasource.Configuration]( in.Id, factory, @@ -265,6 +268,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod if err != nil { return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) } + outs = append(outs, out) case nodev1.DataSourceKind_GRAPHQL: header := http.Header{} @@ -372,6 +376,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod return nil, err } + var out plan.DataSource out, err = plan.NewDataSourceConfigurationWithName[graphql_datasource.Configuration]( in.Id, dataSourceName, @@ -382,10 +387,18 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod if err != nil { return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) } + outs = append(outs, out) case nodev1.DataSourceKind_PUBSUB: var err error - out, err = pubsubDatasource.GetDataSourcesFromConfig(context.Background(), in, l.dataSourceMetaData(in), routerEngineConfig.Events) + + outs, err = pubsub.GetDataSourcesFromConfig( + context.Background(), + in, + l.dataSourceMetaData(in), + routerEngineConfig.Events, + l.logger, + ) if err != nil { return nil, err } @@ -393,7 +406,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod return nil, fmt.Errorf("unknown data source type %q", in.Kind) } - outConfig.DataSources = append(outConfig.DataSources, out) + outConfig.DataSources = append(outConfig.DataSources, outs...) } return &outConfig, nil } diff --git a/router/pkg/pubsub/datasource.go b/router/pkg/pubsub/datasource.go new file mode 100644 index 0000000000..c5bf21a575 --- /dev/null +++ b/router/pkg/pubsub/datasource.go @@ -0,0 +1,35 @@ +package pubsub + +import ( + "context" + "fmt" + + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "go.uber.org/zap" + + // Register all PubSub implementations + _ "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" + _ "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" +) + +func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) ([]plan.DataSource, error) { + var dataSources []plan.DataSource + for _, pubSub := range datasource.GetRegisteredPubSubs() { + ds, err := pubSub(ctx, in, dsMeta, config, logger) + if err != nil { + return nil, err + } + if ds != nil { + dataSources = append(dataSources, ds) + } + } + + if len(dataSources) == 0 { + return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) + } + + return dataSources, nil +} diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index db7d8c9f0c..20a7ee12ab 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -2,70 +2,26 @@ package datasource import ( "context" - "fmt" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" - "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "go.uber.org/zap" ) -type PubSubImplementer[F any] interface { - VerifyConfig(in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error - GetFactory(executionContext context.Context, config config.EventsConfiguration) F -} - -func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) (plan.DataSource, error) { - var pubSubs []any - if natsData := in.GetCustomEvents().GetNats(); natsData != nil { - n := nats.NewPubSub() - err := n.VerifyConfig(in, dsMeta, config) - if err != nil { - return nil, err - } - factory := nats.NewFactory(ctx, config) - ds, err := plan.NewDataSourceConfiguration[nats.Configuration]( - in.Id, - factory, - dsMeta, - nats.Configuration{}, - ) - if err != nil { - return nil, err - } - pubSubs = append(pubSubs, ds) - } - if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { - k := kafka.NewPubSub() - err := k.VerifyConfig(in, dsMeta, config) - if err != nil { - return nil, err - } - factory := k.GetFactory(ctx, config) - ds, err := plan.NewDataSourceConfiguration[kafka.Configuration]( - in.Id, - factory, - dsMeta, - kafka.Configuration{}, - ) +type Getter func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) - if err != nil { - return nil, err - } +var pubSubs []Getter - pubSubs = append(pubSubs, ds) - } +func RegisterPubSub(pubSub Getter) { + pubSubs = append(pubSubs, pubSub) +} - if len(pubSubs) == 0 { - return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) - } +func GetRegisteredPubSubs() []Getter { + return pubSubs +} - factory := NewFactory(ctx, pubSubs) - return plan.NewDataSourceConfiguration[Configuration]( - in.Id, - factory, - dsMeta, - Configuration{}, - ) +type PubSubImplementer[F any] interface { + PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error + GetFactory(executionContext context.Context, config config.EventsConfiguration) F } diff --git a/router/pkg/pubsub/error.go b/router/pkg/pubsub/datasource/error.go similarity index 93% rename from router/pkg/pubsub/error.go rename to router/pkg/pubsub/datasource/error.go index f6220fb7b1..f09b271688 100644 --- a/router/pkg/pubsub/error.go +++ b/router/pkg/pubsub/datasource/error.go @@ -1,4 +1,4 @@ -package pubsub +package datasource type Error struct { Internal error diff --git a/router/pkg/pubsub/datasource/old.go b/router/pkg/pubsub/datasource/old.go index d1f53bb21c..b4528243f5 100644 --- a/router/pkg/pubsub/datasource/old.go +++ b/router/pkg/pubsub/datasource/old.go @@ -33,8 +33,16 @@ func EventTypeFromString(s string) (EventType, error) { } } +type EventMetadata struct { + ProviderID string `json:"providerId"` + Type EventType `json:"type"` + TypeName string `json:"typeName"` + FieldName string `json:"fieldName"` +} + type Configuration struct { - pubSubs []Factory + Metadata *EventMetadata `json:"metadata"` + pubSubs []plan.Planner } type Planner struct { @@ -65,6 +73,16 @@ func (p *Planner) EnterField(ref int) { fieldName := p.visitor.Operation.FieldNameString(ref) typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + //var eventConfig PubSubber + //for _, cfg := range p.config.pubSubs { + // if pubSub, ok := cfg.(PubSubber); ok { + // if pubSub.MatchFieldNameAndType(fieldName, typeName) { + // eventConfig = &cfg + // break + // } + // } + //} + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("nope fieldName %s and typeName %s", fieldName, typeName)) //switch v := eventConfig.Configuration.(type) { diff --git a/router/pkg/pubsub/datasource/providers.go b/router/pkg/pubsub/datasource/providers.go deleted file mode 100644 index 0d7172bf8d..0000000000 --- a/router/pkg/pubsub/datasource/providers.go +++ /dev/null @@ -1 +0,0 @@ -package datasource diff --git a/router/pkg/pubsub/kafka/kafka.go b/router/pkg/pubsub/kafka/kafka.go index 8d2a645e63..6a9260342a 100644 --- a/router/pkg/pubsub/kafka/kafka.go +++ b/router/pkg/pubsub/kafka/kafka.go @@ -10,17 +10,12 @@ import ( "github.com/twmb/franz-go/pkg/kerr" "github.com/twmb/franz-go/pkg/kgo" - "github.com/wundergraph/cosmo/router/pkg/pubsub" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" ) var ( - _ pubsub_datasource.KafkaConnector = (*connector)(nil) - _ pubsub_datasource.KafkaPubSub = (*kafkaPubSub)(nil) - _ pubsub.Lifecycle = (*kafkaPubSub)(nil) - errClientClosed = errors.New("client closed") ) @@ -30,8 +25,7 @@ type connector struct { logger *zap.Logger } -func NewConnector(logger *zap.Logger, opts []kgo.Opt) (pubsub_datasource.KafkaConnector, error) { - +func NewConnector(logger *zap.Logger, opts []kgo.Opt) (*connector, error) { writeClient, err := kgo.NewClient(append(opts, // For observability, we set the client ID to "router" kgo.ClientID("cosmo.router.producer"))..., @@ -47,7 +41,7 @@ func NewConnector(logger *zap.Logger, opts []kgo.Opt) (pubsub_datasource.KafkaCo }, nil } -func (c *connector) New(ctx context.Context) pubsub_datasource.KafkaPubSub { +func (c *connector) New(ctx context.Context) *kafkaPubSub { ctx, cancel := context.WithCancel(ctx) @@ -132,7 +126,7 @@ func (p *kafkaPubSub) topicPoller(ctx context.Context, client *kgo.Client, updat // Subscribe subscribes to the given topics and updates the subscription updater. // The engine already deduplicates subscriptions with the same topics, stream configuration, extensions, headers, etc. -func (p *kafkaPubSub) Subscribe(ctx context.Context, event pubsub_datasource.KafkaSubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { +func (p *kafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), @@ -140,8 +134,6 @@ func (p *kafkaPubSub) Subscribe(ctx context.Context, event pubsub_datasource.Kaf zap.Strings("topics", event.Topics), ) - log.Debug("subscribe") - // Create a new client for the topic client, err := kgo.NewClient(append(p.opts, kgo.ConsumeTopics(event.Topics...), @@ -181,7 +173,7 @@ func (p *kafkaPubSub) Subscribe(ctx context.Context, event pubsub_datasource.Kaf // Publish publishes the given event to the Kafka topic in a non-blocking way. // Publish errors are logged and returned as a pubsub error. // The event is written with a dedicated write client. -func (p *kafkaPubSub) Publish(ctx context.Context, event pubsub_datasource.KafkaPublishEventConfiguration) error { +func (p *kafkaPubSub) Publish(ctx context.Context, event PublishEventConfiguration) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "publish"), @@ -209,7 +201,7 @@ func (p *kafkaPubSub) Publish(ctx context.Context, event pubsub_datasource.Kafka if pErr != nil { log.Error("publish error", zap.Error(pErr)) - return pubsub.NewError(fmt.Sprintf("error publishing to Kafka topic %s", event.Topic), pErr) + return datasource.NewError(fmt.Sprintf("error publishing to Kafka topic %s", event.Topic), pErr) } return nil diff --git a/router/pkg/pubsub/kafka/kafka_datasource.go b/router/pkg/pubsub/kafka/kafka_datasource.go index ae59f6bb97..780bc13c50 100644 --- a/router/pkg/pubsub/kafka/kafka_datasource.go +++ b/router/pkg/pubsub/kafka/kafka_datasource.go @@ -3,11 +3,19 @@ package kafka import ( "bytes" "context" + "crypto/tls" + "encoding/json" "fmt" + "time" "github.com/jensneuse/abstractlogger" + "github.com/twmb/franz-go/pkg/kgo" + "github.com/twmb/franz-go/pkg/sasl/plain" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" + "go.uber.org/zap" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" @@ -15,37 +23,126 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type Kafka struct{} +func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { + if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { + k := NewPubSub(logger) + err := k.PrepareProviders(ctx, in, dsMeta, config) + if err != nil { + return nil, err + } + factory := k.GetFactory(ctx, config) + ds, err := plan.NewDataSourceConfiguration[Configuration]( + in.Id, + factory, + dsMeta, + Configuration{ + EventConfiguration: kafkaData, + Logger: logger, + }, + ) + + if err != nil { + return nil, err + } + + return ds, nil + } + + return nil, nil +} + +func init() { + datasource.RegisterPubSub(GetDataSource) +} + +type Kafka struct { + logger *zap.Logger + providers map[string]*kafkaPubSub +} + +// buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. +// Only general options like TLS, SASL, etc. are configured here. Specific options like topics, etc. are +// configured in the KafkaPubSub implementation. +func buildKafkaOptions(eventSource config.KafkaEventSource) ([]kgo.Opt, error) { + opts := []kgo.Opt{ + kgo.SeedBrokers(eventSource.Brokers...), + // Ensure proper timeouts are set + kgo.ProduceRequestTimeout(10 * time.Second), + kgo.ConnIdleTimeout(60 * time.Second), + } + + if eventSource.TLS != nil && eventSource.TLS.Enabled { + opts = append(opts, + // Configure TLS. Uses SystemCertPool for RootCAs by default. + kgo.DialTLSConfig(new(tls.Config)), + ) + } + + if eventSource.Authentication != nil && eventSource.Authentication.SASLPlain.Username != nil && eventSource.Authentication.SASLPlain.Password != nil { + opts = append(opts, kgo.SASL(plain.Auth{ + User: *eventSource.Authentication.SASLPlain.Username, + Pass: *eventSource.Authentication.SASLPlain.Password, + }.AsMechanism())) + } + + return opts, nil +} -func (k *Kafka) VerifyConfig(in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { - providers := map[string]bool{} +func (k *Kafka) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { + definedProviders := make(map[string]bool) for _, provider := range config.Providers.Kafka { - providers[provider.ID] = true + definedProviders[provider.ID] = true } + usedProviders := make(map[string]bool) for _, event := range in.CustomEvents.GetKafka() { - if !providers[event.EngineEventConfiguration.ProviderId] { + if !definedProviders[event.EngineEventConfiguration.ProviderId] { return fmt.Errorf("failed to find Kafka provider with ID %s", event.EngineEventConfiguration.ProviderId) } + usedProviders[event.EngineEventConfiguration.ProviderId] = true + } + for _, provider := range config.Providers.Kafka { + if !usedProviders[provider.ID] { + continue + } + options, err := buildKafkaOptions(provider) + if err != nil { + return fmt.Errorf("failed to build options for Kafka provider with ID \"%s\": %w", provider.ID, err) + } + ps, err := NewConnector(k.logger, options) + if err != nil { + return fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", provider.ID, err) + } + k.providers[provider.ID] = ps.New(ctx) } return nil } func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { - return NewFactory(executionContext, config) + return NewFactory(executionContext, config, k.providers) } -func NewPubSub() Kafka { - return Kafka{} +func NewPubSub(logger *zap.Logger) Kafka { + return Kafka{ + providers: map[string]*kafkaPubSub{}, + logger: logger, + } } type Configuration struct { - Data string `json:"data"` + Data string `json:"data"` + EventConfiguration []*nodev1.KafkaEventConfiguration + Logger *zap.Logger } type Planner struct { id int config Configuration eventsConfig config.EventsConfiguration + eventConfig *nodev1.KafkaEventConfiguration + rootFieldRef int + variables resolve.Variables + visitor *plan.Visitor + providers map[string]*kafkaPubSub } func (p *Planner) SetID(id int) { @@ -68,49 +165,186 @@ func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { } } -func (p *Planner) Register(_ *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { +func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { + p.visitor = visitor + visitor.Walker.RegisterEnterFieldVisitor(p) + visitor.Walker.RegisterEnterDocumentVisitor(p) p.config = configuration.CustomConfiguration() return nil } func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { + if p.eventConfig == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) + return resolve.FetchConfiguration{} + } + + var dataSource resolve.DataSource + providerId := p.eventConfig.GetEngineEventConfiguration().GetProviderId() + typeName := p.eventConfig.GetEngineEventConfiguration().GetType() + topics := p.eventConfig.GetTopics() + pubsub, ok := p.providers[providerId] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) + return resolve.FetchConfiguration{} + } + + switch p.eventConfig.GetEngineEventConfiguration().GetType() { + case nodev1.EventType_PUBLISH: + dataSource = &KafkaPublishDataSource{ + pubSub: pubsub, + } + default: + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String())) + return resolve.FetchConfiguration{} + } + + if len(topics) != 1 { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("publish and request events should define one subject but received %d", len(topics))) + return resolve.FetchConfiguration{} + } + + topic := topics[0] + + event, eventErr := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, p.variables) + if eventErr != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to build event data bytes: %w", eventErr)) + return resolve.FetchConfiguration{} + } + + evtCfg := PublishEventConfiguration{ + ProviderID: providerId, + Topic: topic, + Data: event, + } + return resolve.FetchConfiguration{ - Input: p.config.Data, - DataSource: Source{}, + Input: evtCfg.MarshalJSONTemplate(), + Variables: p.variables, + DataSource: dataSource, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, } } func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { + if p.eventConfig == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) + return plan.SubscriptionConfiguration{} + } + providerId := p.eventConfig.GetEngineEventConfiguration().GetProviderId() + pubsub, ok := p.providers[providerId] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) + return plan.SubscriptionConfiguration{} + } + evtCfg := SubscriptionEventConfiguration{ + ProviderID: providerId, + Topics: p.eventConfig.GetTopics(), + } + object, err := json.Marshal(evtCfg) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) + return plan.SubscriptionConfiguration{} + } + return plan.SubscriptionConfiguration{ - Input: p.config.Data, + Input: string(object), + Variables: p.variables, + DataSource: &SubscriptionSource{ + pubSub: pubsub, + }, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner) EnterDocument(_, _ *ast.Document) { + p.rootFieldRef = -1 + p.eventConfig = nil +} + +func (p *Planner) EnterField(ref int) { + if p.rootFieldRef != -1 { + // This is a nested field; nothing needs to be done + return + } + p.rootFieldRef = ref + + fieldName := p.visitor.Operation.FieldNameString(ref) + typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + + var eventConfig *nodev1.KafkaEventConfiguration + for _, cfg := range p.config.EventConfiguration { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + eventConfig = cfg + break + } + } + + if eventConfig == nil { + return + } + + p.eventConfig = eventConfig + + providerId := eventConfig.GetEngineEventConfiguration().GetProviderId() + + switch eventConfig.GetEngineEventConfiguration().GetType() { + case nodev1.EventType_PUBLISH: + if len(p.eventConfig.GetTopics()) != 1 { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("publish events should define one subject but received %d", len(p.eventConfig.GetTopics()))) + return + } + _, found := p.providers[providerId] + if !found { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("unable to publish events of provider %d", len(p.eventConfig.GetTopics()))) + } + _ = PublishEventConfiguration{ + ProviderID: providerId, + Topic: eventConfig.GetTopics()[0], + Data: json.RawMessage("[]"), + } + p.config.Logger.Warn("Publishing!") + // provider.Publish(provider.ctx, pubCfg) + case nodev1.EventType_SUBSCRIBE: + p.config.Logger.Warn("Subscribing!") + default: + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Kafka", eventConfig.GetEngineEventConfiguration().GetType())) } } type Source struct{} -func (Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { +func (s *Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { _, err = out.Write(input) return } -func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { +func (s *Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { panic("not implemented") } -func NewFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { +func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*kafkaPubSub) *Factory { return &Factory{ + providers: providers, executionContext: executionContext, config: config, } } type Factory struct { + providers map[string]*kafkaPubSub executionContext context.Context config config.EventsConfiguration } func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { - return &Planner{} + return &Planner{ + providers: f.providers, + } } func (f *Factory) Context() context.Context { diff --git a/router/pkg/pubsub/kafka/subscription_datasource.go b/router/pkg/pubsub/kafka/subscription_datasource.go new file mode 100644 index 0000000000..3109d3cf35 --- /dev/null +++ b/router/pkg/pubsub/kafka/subscription_datasource.go @@ -0,0 +1,87 @@ +package kafka + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + + "github.com/buger/jsonparser" + "github.com/cespare/xxhash/v2" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type SubscriptionEventConfiguration struct { + ProviderID string `json:"providerId"` + Topics []string `json:"topics"` +} + +type PublishEventConfiguration struct { + ProviderID string `json:"providerId"` + Topic string `json:"topic"` + Data json.RawMessage `json:"data"` +} + +func (s *PublishEventConfiguration) MarshalJSONTemplate() string { + return fmt.Sprintf(`{"topic":"%s", "data": %s, "providerId":"%s"}`, s.Topic, s.Data, s.ProviderID) +} + +type SubscriptionSource struct { + pubSub *kafkaPubSub +} + +func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { + + val, _, _, err := jsonparser.Get(input, "topics") + if err != nil { + return err + } + + _, err = xxh.Write(val) + if err != nil { + return err + } + + val, _, _, err = jsonparser.Get(input, "providerId") + if err != nil { + return err + } + + _, err = xxh.Write(val) + return err +} + +func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater resolve.SubscriptionUpdater) error { + var subscriptionConfiguration SubscriptionEventConfiguration + err := json.Unmarshal(input, &subscriptionConfiguration) + if err != nil { + return err + } + + return s.pubSub.Subscribe(ctx.Context(), subscriptionConfiguration, updater) +} + +type KafkaPublishDataSource struct { + pubSub *kafkaPubSub +} + +func (s *KafkaPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { + var publishConfiguration PublishEventConfiguration + err := json.Unmarshal(input, &publishConfiguration) + if err != nil { + return err + } + + if err := s.pubSub.Publish(ctx, publishConfiguration); err != nil { + _, err = io.WriteString(out, `{"success": false}`) + return err + } + _, err = io.WriteString(out, `{"success": true}`) + return err +} + +func (s *KafkaPublishDataSource) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { + panic("not implemented") +} diff --git a/router/pkg/pubsub/nats/nats.go b/router/pkg/pubsub/nats/nats.go index 921b199e40..c609b55d3f 100644 --- a/router/pkg/pubsub/nats/nats.go +++ b/router/pkg/pubsub/nats/nats.go @@ -4,22 +4,17 @@ import ( "context" "errors" "fmt" + "io" + "sync" + "time" + "github.com/cespare/xxhash/v2" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" - "github.com/wundergraph/cosmo/router/pkg/pubsub" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" - "io" - "sync" - "time" -) - -var ( - _ pubsub_datasource.NatsConnector = (*connector)(nil) - _ pubsub_datasource.NatsPubSub = (*natsPubSub)(nil) - _ pubsub.Lifecycle = (*natsPubSub)(nil) ) type connector struct { @@ -30,7 +25,7 @@ type connector struct { routerListenAddr string } -func NewConnector(logger *zap.Logger, conn *nats.Conn, js jetstream.JetStream, hostName string, routerListenAddr string) pubsub_datasource.NatsConnector { +func NewConnector(logger *zap.Logger, conn *nats.Conn, js jetstream.JetStream, hostName string, routerListenAddr string) *connector { return &connector{ conn: conn, logger: logger, @@ -40,7 +35,7 @@ func NewConnector(logger *zap.Logger, conn *nats.Conn, js jetstream.JetStream, h } } -func (c *connector) New(ctx context.Context) pubsub_datasource.NatsPubSub { +func (c *connector) New(ctx context.Context) *natsPubSub { return &natsPubSub{ ctx: ctx, conn: c.conn, @@ -113,7 +108,7 @@ func (p *natsPubSub) Subscribe(ctx context.Context, event pubsub_datasource.Nats consumer, err := p.js.CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) if err != nil { log.Error("error creating or updating consumer", zap.Error(err)) - return pubsub.NewError(fmt.Sprintf(`failed to create or update consumer for stream "%s"`, event.StreamConfiguration.StreamName), err) + return datasource.NewError(fmt.Sprintf(`failed to create or update consumer for stream "%s"`, event.StreamConfiguration.StreamName), err) } p.closeWg.Add(1) @@ -164,7 +159,7 @@ func (p *natsPubSub) Subscribe(ctx context.Context, event pubsub_datasource.Nats subscription, err := p.conn.ChanSubscribe(subject, msgChan) if err != nil { log.Error("error subscribing to NATS subject", zap.Error(err), zap.String("subscription_subject", subject)) - return pubsub.NewError(fmt.Sprintf(`failed to subscribe to NATS subject "%s"`, subject), err) + return datasource.NewError(fmt.Sprintf(`failed to subscribe to NATS subject "%s"`, subject), err) } subscriptions[i] = subscription } @@ -218,7 +213,7 @@ func (p *natsPubSub) Publish(_ context.Context, event pubsub_datasource.NatsPubl err := p.conn.Publish(event.Subject, event.Data) if err != nil { log.Error("publish error", zap.Error(err)) - return pubsub.NewError(fmt.Sprintf("error publishing to NATS subject %s", event.Subject), err) + return datasource.NewError(fmt.Sprintf("error publishing to NATS subject %s", event.Subject), err) } return nil @@ -236,7 +231,7 @@ func (p *natsPubSub) Request(ctx context.Context, event pubsub_datasource.NatsPu msg, err := p.conn.RequestWithContext(ctx, event.Subject, event.Data) if err != nil { log.Error("request error", zap.Error(err)) - return pubsub.NewError(fmt.Sprintf("error requesting from NATS subject %s", event.Subject), err) + return datasource.NewError(fmt.Sprintf("error requesting from NATS subject %s", event.Subject), err) } _, err = w.Write(msg.Data) diff --git a/router/pkg/pubsub/nats/nats_datasource.go b/router/pkg/pubsub/nats/nats_datasource.go index 8c893a0567..428175ec8d 100644 --- a/router/pkg/pubsub/nats/nats_datasource.go +++ b/router/pkg/pubsub/nats/nats_datasource.go @@ -3,11 +3,17 @@ package nats import ( "bytes" "context" + "errors" "fmt" + "time" "github.com/jensneuse/abstractlogger" + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "go.uber.org/zap" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" @@ -15,23 +21,126 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type Nats struct{} +func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { + if natsData := in.GetCustomEvents().GetNats(); natsData != nil { + k := NewPubSub() + err := k.PrepareProviders(ctx, in, dsMeta, config) + if err != nil { + return nil, err + } + factory := k.GetFactory(ctx, config, k.providers) + ds, err := plan.NewDataSourceConfiguration[Configuration]( + in.Id, + factory, + dsMeta, + Configuration{ + EventConfiguration: natsData, + Logger: logger, + }, + ) + + if err != nil { + return nil, err + } -func (n *Nats) VerifyConfig(in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { - providers := map[string]bool{} - for _, provider := range config.Providers.Nats { - providers[provider.ID] = true + return ds, nil + } + + return nil, nil +} + +func init() { + datasource.RegisterPubSub(GetDataSource) +} + +func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { + opts := []nats.Option{ + nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), + nats.ReconnectJitter(500*time.Millisecond, 2*time.Second), + nats.ClosedHandler(func(conn *nats.Conn) { + logger.Info("NATS connection closed", zap.String("provider_id", eventSource.ID), zap.Error(conn.LastError())) + }), + nats.ConnectHandler(func(nc *nats.Conn) { + logger.Info("NATS connection established", zap.String("provider_id", eventSource.ID), zap.String("url", nc.ConnectedUrlRedacted())) + }), + nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { + if err != nil { + logger.Error("NATS disconnected; will attempt to reconnect", zap.Error(err), zap.String("provider_id", eventSource.ID)) + } else { + logger.Info("NATS disconnected", zap.String("provider_id", eventSource.ID)) + } + }), + nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { + if errors.Is(err, nats.ErrSlowConsumer) { + logger.Warn( + "NATS slow consumer detected. Events are being dropped. Please consider increasing the buffer size or reducing the number of messages being sent.", + zap.Error(err), + zap.String("provider_id", eventSource.ID), + ) + } else { + logger.Error("NATS error", zap.Error(err)) + } + }), + nats.ReconnectHandler(func(conn *nats.Conn) { + logger.Info("NATS reconnected", zap.String("provider_id", eventSource.ID), zap.String("url", conn.ConnectedUrlRedacted())) + }), + } + + if eventSource.Authentication != nil { + if eventSource.Authentication.Token != nil { + opts = append(opts, nats.Token(*eventSource.Authentication.Token)) + } else if eventSource.Authentication.UserInfo.Username != nil && eventSource.Authentication.UserInfo.Password != nil { + opts = append(opts, nats.UserInfo(*eventSource.Authentication.UserInfo.Username, *eventSource.Authentication.UserInfo.Password)) + } + } + + return opts, nil +} + +type Nats struct { + providers map[string]*natsPubSub + logger *zap.Logger + hostName string // How to get it here? + routerListenAddr string // How to get it here? +} + +func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { + definedProviders := make(map[string]bool) + for _, provider := range config.Providers.Kafka { + definedProviders[provider.ID] = true } + usedProviders := make(map[string]bool) for _, event := range in.CustomEvents.GetNats() { - if !providers[event.EngineEventConfiguration.ProviderId] { + if _, found := definedProviders[event.EngineEventConfiguration.ProviderId]; !found { return fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) } + usedProviders[event.EngineEventConfiguration.ProviderId] = true + } + for _, provider := range config.Providers.Nats { + if !usedProviders[provider.ID] { + continue + } + options, err := buildNatsOptions(provider, n.logger) + if err != nil { + return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", provider.ID, err) + } + natsConnection, err := nats.Connect(provider.URL, options...) + if err != nil { + return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", provider.ID, err) + } + js, err := jetstream.New(natsConnection) + if err != nil { + return err + } + + n.providers[provider.ID] = NewConnector(n.logger, natsConnection, js, n.hostName, n.routerListenAddr).New(ctx) + } return nil } -func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { - return NewFactory(executionContext, config) +func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*natsPubSub) *Factory { + return NewFactory(executionContext, config, providers) } func NewPubSub() Nats { @@ -39,12 +148,15 @@ func NewPubSub() Nats { } type Configuration struct { - Data string `json:"data"` + Data string `json:"data"` + EventConfiguration []*nodev1.NatsEventConfiguration + Logger *zap.Logger } type Planner struct { - id int - config Configuration + id int + config Configuration + providers map[string]*natsPubSub } func (p *Planner) SetID(id int) { @@ -96,10 +208,11 @@ func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclie panic("not implemented") } -func NewFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { +func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*natsPubSub) *Factory { return &Factory{ executionContext: executionContext, eventsConfiguration: config, + providers: providers, } } @@ -107,11 +220,13 @@ type Factory struct { config Configuration eventsConfiguration config.EventsConfiguration executionContext context.Context + providers map[string]*natsPubSub } func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { return &Planner{ - config: f.config, + config: f.config, + providers: f.providers, } } diff --git a/router/pkg/pubsub/utils/utils.go b/router/pkg/pubsub/utils/utils.go new file mode 100644 index 0000000000..92410b641e --- /dev/null +++ b/router/pkg/pubsub/utils/utils.go @@ -0,0 +1,38 @@ +package utils + +import ( + "bytes" + "encoding/json" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +func BuildEventDataBytes(ref int, visitor *plan.Visitor, variables resolve.Variables) ([]byte, error) { + // Collect the field arguments for fetch based operations + fieldArgs := visitor.Operation.FieldArguments(ref) + var dataBuffer bytes.Buffer + dataBuffer.WriteByte('{') + for i, arg := range fieldArgs { + if i > 0 { + dataBuffer.WriteByte(',') + } + argValue := visitor.Operation.ArgumentValue(arg) + variableName := visitor.Operation.VariableValueNameBytes(argValue.Ref) + contextVariable := &resolve.ContextVariable{ + Path: []string{string(variableName)}, + Renderer: resolve.NewPlainVariableRenderer(), + } + variablePlaceHolder, _ := variables.AddVariable(contextVariable) + argumentName := visitor.Operation.ArgumentNameString(arg) + escapedKey, err := json.Marshal(argumentName) + if err != nil { + return nil, err + } + dataBuffer.Write(escapedKey) + dataBuffer.WriteByte(':') + dataBuffer.WriteString(variablePlaceHolder) + } + dataBuffer.WriteByte('}') + return dataBuffer.Bytes(), nil +} From 5df3613983d3e3cb33374d3423f7362c7b09bf45 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 18 Mar 2025 16:49:12 +0100 Subject: [PATCH 03/49] feat(plan-generator): add logger to router initialization --- router/core/plan_generator.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index a9e73ff777..a2628f5a54 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -8,6 +8,7 @@ import ( "os" log "github.com/jensneuse/abstractlogger" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/astnormalization" "github.com/wundergraph/graphql-go-tools/v2/pkg/astparser" @@ -174,7 +175,7 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo streamingClient: http.DefaultClient, subscriptionClient: subscriptionClient, transportOptions: &TransportOptions{SubgraphTransportOptions: NewSubgraphTransportOptions(config.TrafficShapingRules{})}, - }) + }, logger) // this generates the plan configuration using the data source factories from the config package planConfig, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &RouterEngineConfiguration{}) From efaa122cd7c052ca0583cf7e844cf86fcffa1f6f Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Thu, 20 Mar 2025 16:54:50 +0100 Subject: [PATCH 04/49] refactor(pubsub): nats is now working --- .../subgraphs/availability/availability.go | 4 +- .../availability/subgraph/resolver.go | 4 +- .../availability/subgraph/schema.resolvers.go | 6 +- demo/pkg/subgraphs/countries/countries.go | 4 +- .../subgraphs/countries/subgraph/resolver.go | 5 +- demo/pkg/subgraphs/employees/employees.go | 5 +- .../subgraphs/employees/subgraph/resolver.go | 4 +- demo/pkg/subgraphs/family/family.go | 4 +- .../pkg/subgraphs/family/subgraph/resolver.go | 4 +- demo/pkg/subgraphs/hobbies/hobbies.go | 4 +- .../subgraphs/hobbies/subgraph/resolver.go | 4 +- demo/pkg/subgraphs/mood/mood.go | 4 +- demo/pkg/subgraphs/mood/subgraph/resolver.go | 4 +- .../mood/subgraph/schema.resolvers.go | 6 +- demo/pkg/subgraphs/products/products.go | 4 +- .../subgraphs/products/subgraph/resolver.go | 7 +- demo/pkg/subgraphs/products_fg/products.go | 4 +- .../products_fg/subgraph/resolver.go | 7 +- demo/pkg/subgraphs/subgraphs.go | 5 +- demo/pkg/subgraphs/test1/subgraph/resolver.go | 4 +- demo/pkg/subgraphs/test1/test1.go | 4 +- router-tests/events/nats_events_test.go | 2 +- router-tests/testenv/testenv.go | 6 +- .../{kafka_datasource.go => datasource.go} | 4 +- router/pkg/pubsub/nats/datasource.go | 516 ++++++++++++++++++ router/pkg/pubsub/nats/nats.go | 21 +- router/pkg/pubsub/nats/nats_datasource.go | 239 -------- .../pubsub/nats/subscription_datasource.go | 122 +++++ router/pkg/pubsub/utils/utils.go | 2 +- 29 files changed, 703 insertions(+), 306 deletions(-) rename router/pkg/pubsub/kafka/{kafka_datasource.go => datasource.go} (99%) create mode 100644 router/pkg/pubsub/nats/datasource.go delete mode 100644 router/pkg/pubsub/nats/nats_datasource.go create mode 100644 router/pkg/pubsub/nats/subscription_datasource.go diff --git a/demo/pkg/subgraphs/availability/availability.go b/demo/pkg/subgraphs/availability/availability.go index f37ce470a5..ec693711cb 100644 --- a/demo/pkg/subgraphs/availability/availability.go +++ b/demo/pkg/subgraphs/availability/availability.go @@ -2,13 +2,13 @@ package availability import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/availability/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/availability/subgraph/generated" ) -func NewSchema(pubSubBySourceName map[string]pubsub_datasource.NatsPubSub, pubSubName func(string) string) graphql.ExecutableSchema { +func NewSchema(pubSubBySourceName map[string]*nats.NatsPubSub, pubSubName func(string) string) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: pubSubBySourceName, GetPubSubName: pubSubName, diff --git a/demo/pkg/subgraphs/availability/subgraph/resolver.go b/demo/pkg/subgraphs/availability/subgraph/resolver.go index 9ac6f82f89..554b11af76 100644 --- a/demo/pkg/subgraphs/availability/subgraph/resolver.go +++ b/demo/pkg/subgraphs/availability/subgraph/resolver.go @@ -1,7 +1,7 @@ package subgraph import ( - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -9,6 +9,6 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub GetPubSubName func(string) string } diff --git a/demo/pkg/subgraphs/availability/subgraph/schema.resolvers.go b/demo/pkg/subgraphs/availability/subgraph/schema.resolvers.go index 43797ed3dc..7067d88bf1 100644 --- a/demo/pkg/subgraphs/availability/subgraph/schema.resolvers.go +++ b/demo/pkg/subgraphs/availability/subgraph/schema.resolvers.go @@ -10,13 +10,13 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/availability/subgraph/generated" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/availability/subgraph/model" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // UpdateAvailability is the resolver for the updateAvailability field. func (r *mutationResolver) UpdateAvailability(ctx context.Context, employeeID int, isAvailable bool) (*model.Employee, error) { storage.Set(employeeID, isAvailable) - err := r.NatsPubSubByProviderID["default"].Publish(ctx, pubsub_datasource.NatsPublishAndRequestEventConfiguration{ + err := r.NatsPubSubByProviderID["default"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ Subject: r.GetPubSubName(fmt.Sprintf("employeeUpdated.%d", employeeID)), Data: []byte(fmt.Sprintf(`{"id":%d,"__typename": "Employee"}`, employeeID)), }) @@ -24,7 +24,7 @@ func (r *mutationResolver) UpdateAvailability(ctx context.Context, employeeID in if err != nil { return nil, err } - err = r.NatsPubSubByProviderID["my-nats"].Publish(ctx, pubsub_datasource.NatsPublishAndRequestEventConfiguration{ + err = r.NatsPubSubByProviderID["my-nats"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ Subject: r.GetPubSubName(fmt.Sprintf("employeeUpdatedMyNats.%d", employeeID)), Data: []byte(fmt.Sprintf(`{"id":%d,"__typename": "Employee"}`, employeeID)), }) diff --git a/demo/pkg/subgraphs/countries/countries.go b/demo/pkg/subgraphs/countries/countries.go index 3a1ccb7427..3d5940c1e6 100644 --- a/demo/pkg/subgraphs/countries/countries.go +++ b/demo/pkg/subgraphs/countries/countries.go @@ -2,13 +2,13 @@ package countries import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/countries/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/countries/subgraph/generated" ) -func NewSchema(pubSubBySourceName map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(pubSubBySourceName map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: pubSubBySourceName, }}) diff --git a/demo/pkg/subgraphs/countries/subgraph/resolver.go b/demo/pkg/subgraphs/countries/subgraph/resolver.go index 4b235fdec9..c1f3775ab1 100644 --- a/demo/pkg/subgraphs/countries/subgraph/resolver.go +++ b/demo/pkg/subgraphs/countries/subgraph/resolver.go @@ -1,8 +1,9 @@ package subgraph import ( - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "sync" + + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -11,5 +12,5 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub } diff --git a/demo/pkg/subgraphs/employees/employees.go b/demo/pkg/subgraphs/employees/employees.go index 408737da15..5b6b6eae19 100644 --- a/demo/pkg/subgraphs/employees/employees.go +++ b/demo/pkg/subgraphs/employees/employees.go @@ -2,13 +2,12 @@ package employees import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" - "github.com/wundergraph/cosmo/demo/pkg/subgraphs/employees/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/employees/subgraph/generated" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, EmployeesData: subgraph.Employees, diff --git a/demo/pkg/subgraphs/employees/subgraph/resolver.go b/demo/pkg/subgraphs/employees/subgraph/resolver.go index f71624ab76..1f75150132 100644 --- a/demo/pkg/subgraphs/employees/subgraph/resolver.go +++ b/demo/pkg/subgraphs/employees/subgraph/resolver.go @@ -6,7 +6,7 @@ import ( "sync" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/employees/subgraph/model" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -15,7 +15,7 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub EmployeesData []*model.Employee } diff --git a/demo/pkg/subgraphs/family/family.go b/demo/pkg/subgraphs/family/family.go index c55eae3fe4..4a682b538d 100644 --- a/demo/pkg/subgraphs/family/family.go +++ b/demo/pkg/subgraphs/family/family.go @@ -2,13 +2,13 @@ package family import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/family/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/family/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/family/subgraph/resolver.go b/demo/pkg/subgraphs/family/subgraph/resolver.go index f4678ba12e..7c404a867c 100644 --- a/demo/pkg/subgraphs/family/subgraph/resolver.go +++ b/demo/pkg/subgraphs/family/subgraph/resolver.go @@ -1,7 +1,7 @@ package subgraph import ( - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -9,5 +9,5 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub } diff --git a/demo/pkg/subgraphs/hobbies/hobbies.go b/demo/pkg/subgraphs/hobbies/hobbies.go index 103e8bb43a..926b397c6f 100644 --- a/demo/pkg/subgraphs/hobbies/hobbies.go +++ b/demo/pkg/subgraphs/hobbies/hobbies.go @@ -2,13 +2,13 @@ package hobbies import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/hobbies/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/hobbies/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/hobbies/subgraph/resolver.go b/demo/pkg/subgraphs/hobbies/subgraph/resolver.go index dc972ffd31..696412a06a 100644 --- a/demo/pkg/subgraphs/hobbies/subgraph/resolver.go +++ b/demo/pkg/subgraphs/hobbies/subgraph/resolver.go @@ -1,10 +1,10 @@ package subgraph import ( - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "reflect" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/hobbies/subgraph/model" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -12,7 +12,7 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub } func (r *Resolver) Employees(hobby model.Hobby) ([]*model.Employee, error) { diff --git a/demo/pkg/subgraphs/mood/mood.go b/demo/pkg/subgraphs/mood/mood.go index 7083a607e7..72b6c7118b 100644 --- a/demo/pkg/subgraphs/mood/mood.go +++ b/demo/pkg/subgraphs/mood/mood.go @@ -2,13 +2,13 @@ package mood import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/mood/subgraph/resolver.go b/demo/pkg/subgraphs/mood/subgraph/resolver.go index 9ac6f82f89..554b11af76 100644 --- a/demo/pkg/subgraphs/mood/subgraph/resolver.go +++ b/demo/pkg/subgraphs/mood/subgraph/resolver.go @@ -1,7 +1,7 @@ package subgraph import ( - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -9,6 +9,6 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub GetPubSubName func(string) string } diff --git a/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go b/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go index 1067b65909..3508b7d36d 100644 --- a/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go +++ b/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go @@ -10,7 +10,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph/generated" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph/model" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // UpdateMood is the resolver for the updateMood field. @@ -18,7 +18,7 @@ func (r *mutationResolver) UpdateMood(ctx context.Context, employeeID int, mood storage.Set(employeeID, mood) myNatsTopic := r.GetPubSubName(fmt.Sprintf("employeeUpdated.%d", employeeID)) payload := fmt.Sprintf(`{"id":%d,"__typename": "Employee"}`, employeeID) - err := r.NatsPubSubByProviderID["default"].Publish(ctx, pubsub_datasource.NatsPublishAndRequestEventConfiguration{ + err := r.NatsPubSubByProviderID["default"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ Subject: myNatsTopic, Data: []byte(payload), }) @@ -27,7 +27,7 @@ func (r *mutationResolver) UpdateMood(ctx context.Context, employeeID int, mood } defaultTopic := r.GetPubSubName(fmt.Sprintf("employeeUpdatedMyNats.%d", employeeID)) - err = r.NatsPubSubByProviderID["my-nats"].Publish(ctx, pubsub_datasource.NatsPublishAndRequestEventConfiguration{ + err = r.NatsPubSubByProviderID["my-nats"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ Subject: defaultTopic, Data: []byte(payload), }) diff --git a/demo/pkg/subgraphs/products/products.go b/demo/pkg/subgraphs/products/products.go index f14cc97813..6a19c4baba 100644 --- a/demo/pkg/subgraphs/products/products.go +++ b/demo/pkg/subgraphs/products/products.go @@ -2,13 +2,13 @@ package products import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, TopSecretFederationFactsData: subgraph.TopSecretFederationFacts, diff --git a/demo/pkg/subgraphs/products/subgraph/resolver.go b/demo/pkg/subgraphs/products/subgraph/resolver.go index c9d610cb04..64543f0a6d 100644 --- a/demo/pkg/subgraphs/products/subgraph/resolver.go +++ b/demo/pkg/subgraphs/products/subgraph/resolver.go @@ -1,9 +1,10 @@ package subgraph import ( - "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products/subgraph/model" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "sync" + + "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products/subgraph/model" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -12,6 +13,6 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub TopSecretFederationFactsData []model.TopSecretFact } diff --git a/demo/pkg/subgraphs/products_fg/products.go b/demo/pkg/subgraphs/products_fg/products.go index 155cffe417..436c34582b 100644 --- a/demo/pkg/subgraphs/products_fg/products.go +++ b/demo/pkg/subgraphs/products_fg/products.go @@ -2,13 +2,13 @@ package products_fg import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products_fg/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products_fg/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, TopSecretFederationFactsData: subgraph.TopSecretFederationFacts, diff --git a/demo/pkg/subgraphs/products_fg/subgraph/resolver.go b/demo/pkg/subgraphs/products_fg/subgraph/resolver.go index f0f8d7059a..c6dca70e73 100644 --- a/demo/pkg/subgraphs/products_fg/subgraph/resolver.go +++ b/demo/pkg/subgraphs/products_fg/subgraph/resolver.go @@ -1,9 +1,10 @@ package subgraph import ( - "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products_fg/subgraph/model" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "sync" + + "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products_fg/subgraph/model" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -12,6 +13,6 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub TopSecretFederationFactsData []model.TopSecretFact } diff --git a/demo/pkg/subgraphs/subgraphs.go b/demo/pkg/subgraphs/subgraphs.go index abcaf75973..ab3820d02c 100644 --- a/demo/pkg/subgraphs/subgraphs.go +++ b/demo/pkg/subgraphs/subgraphs.go @@ -22,7 +22,6 @@ import ( "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" natsPubsub "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "golang.org/x/sync/errgroup" "github.com/wundergraph/cosmo/demo/pkg/injector" @@ -162,7 +161,7 @@ func subgraphHandler(schema graphql.ExecutableSchema) http.Handler { } type SubgraphOptions struct { - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*natsPubsub.NatsPubSub GetPubSubName func(string) string } @@ -227,7 +226,7 @@ func New(ctx context.Context, config *Config) (*Subgraphs, error) { return nil, err } - natsPubSubByProviderID := map[string]pubsub_datasource.NatsPubSub{ + natsPubSubByProviderID := map[string]*natsPubsub.NatsPubSub{ "default": natsPubsub.NewConnector(zap.NewNop(), defaultConnection, defaultJetStream, "hostname", "test").New(ctx), "my-nats": natsPubsub.NewConnector(zap.NewNop(), myNatsConnection, myNatsJetStream, "hostname", "test").New(ctx), } diff --git a/demo/pkg/subgraphs/test1/subgraph/resolver.go b/demo/pkg/subgraphs/test1/subgraph/resolver.go index f4678ba12e..7c404a867c 100644 --- a/demo/pkg/subgraphs/test1/subgraph/resolver.go +++ b/demo/pkg/subgraphs/test1/subgraph/resolver.go @@ -1,7 +1,7 @@ package subgraph import ( - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) // This file will not be regenerated automatically. @@ -9,5 +9,5 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub + NatsPubSubByProviderID map[string]*nats.NatsPubSub } diff --git a/demo/pkg/subgraphs/test1/test1.go b/demo/pkg/subgraphs/test1/test1.go index 25f00b8ec7..8b74990bc4 100644 --- a/demo/pkg/subgraphs/test1/test1.go +++ b/demo/pkg/subgraphs/test1/test1.go @@ -2,13 +2,13 @@ package test1 import ( "github.com/99designs/gqlgen/graphql" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/test1/subgraph" "github.com/wundergraph/cosmo/demo/pkg/subgraphs/test1/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]pubsub_datasource.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/router-tests/events/nats_events_test.go b/router-tests/events/nats_events_test.go index 0f0fac6325..fdd068a083 100644 --- a/router-tests/events/nats_events_test.go +++ b/router-tests/events/nats_events_test.go @@ -136,7 +136,7 @@ func TestNatsEvents(t *testing.T) { natsLogs := xEnv.Observer().FilterMessageSnippet("Nats").All() require.Len(t, natsLogs, 4) providerIDFields := xEnv.Observer().FilterField(zap.String("provider_id", "my-nats")).All() - require.Len(t, providerIDFields, 2) + require.Len(t, providerIDFields, 3) }) }) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index daa023f9dd..2e0de276b1 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -51,8 +51,6 @@ import ( "go.uber.org/zap/zaptest/observer" "google.golang.org/protobuf/encoding/protojson" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" - "github.com/wundergraph/cosmo/demo/pkg/subgraphs" "github.com/wundergraph/cosmo/router/core" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" @@ -2032,11 +2030,11 @@ func DeflakeWSWriteJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err func subgraphOptions(ctx context.Context, t testing.TB, logger *zap.Logger, natsData *NatsData, pubSubName func(string) string) *subgraphs.SubgraphOptions { if natsData == nil { return &subgraphs.SubgraphOptions{ - NatsPubSubByProviderID: map[string]pubsub_datasource.NatsPubSub{}, + NatsPubSubByProviderID: map[string]*pubsubNats.NatsPubSub{}, GetPubSubName: pubSubName, } } - natsPubSubByProviderID := make(map[string]pubsub_datasource.NatsPubSub, len(demoNatsProviders)) + natsPubSubByProviderID := make(map[string]*pubsubNats.NatsPubSub, len(demoNatsProviders)) for _, sourceName := range demoNatsProviders { js, err := jetstream.New(natsData.Connections[0]) require.NoError(t, err) diff --git a/router/pkg/pubsub/kafka/kafka_datasource.go b/router/pkg/pubsub/kafka/datasource.go similarity index 99% rename from router/pkg/pubsub/kafka/kafka_datasource.go rename to router/pkg/pubsub/kafka/datasource.go index 780bc13c50..3a6fe3d8ee 100644 --- a/router/pkg/pubsub/kafka/kafka_datasource.go +++ b/router/pkg/pubsub/kafka/datasource.go @@ -31,7 +31,7 @@ func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMe return nil, err } factory := k.GetFactory(ctx, config) - ds, err := plan.NewDataSourceConfiguration[Configuration]( + ds, err := plan.NewDataSourceConfiguration( in.Id, factory, dsMeta, @@ -206,7 +206,7 @@ func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { topic := topics[0] - event, eventErr := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, p.variables) + event, eventErr := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) if eventErr != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to build event data bytes: %w", eventErr)) return resolve.FetchConfiguration{} diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go new file mode 100644 index 0000000000..b85deed14f --- /dev/null +++ b/router/pkg/pubsub/nats/datasource.go @@ -0,0 +1,516 @@ +package nats + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "regexp" + "slices" + "strings" + "time" + + "github.com/jensneuse/abstractlogger" + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" + "go.uber.org/zap" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +const ( + fwc = '>' + tsep = "." +) + +// A variable template has form $$number$$ where the number can range from one to multiple digits +var ( + variableTemplateRegex = regexp.MustCompile(`\$\$\d+\$\$`) +) + +func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { + if natsData := in.GetCustomEvents().GetNats(); natsData != nil { + k := NewPubSub(logger) + err := k.PrepareProviders(ctx, in, dsMeta, config) + if err != nil { + return nil, err + } + factory := k.GetFactory(ctx, config, k.providers) + ds, err := plan.NewDataSourceConfiguration[Configuration]( + in.Id, + factory, + dsMeta, + Configuration{ + EventConfiguration: natsData, + Logger: logger, + }, + ) + + if err != nil { + return nil, err + } + + return ds, nil + } + + return nil, nil +} + +func init() { + datasource.RegisterPubSub(GetDataSource) +} + +func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { + opts := []nats.Option{ + nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), + nats.ReconnectJitter(500*time.Millisecond, 2*time.Second), + nats.ClosedHandler(func(conn *nats.Conn) { + logger.Info("NATS connection closed", zap.String("provider_id", eventSource.ID), zap.Error(conn.LastError())) + }), + nats.ConnectHandler(func(nc *nats.Conn) { + logger.Info("NATS connection established", zap.String("provider_id", eventSource.ID), zap.String("url", nc.ConnectedUrlRedacted())) + }), + nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { + if err != nil { + logger.Error("NATS disconnected; will attempt to reconnect", zap.Error(err), zap.String("provider_id", eventSource.ID)) + } else { + logger.Info("NATS disconnected", zap.String("provider_id", eventSource.ID)) + } + }), + nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { + if errors.Is(err, nats.ErrSlowConsumer) { + logger.Warn( + "NATS slow consumer detected. Events are being dropped. Please consider increasing the buffer size or reducing the number of messages being sent.", + zap.Error(err), + zap.String("provider_id", eventSource.ID), + ) + } else { + logger.Error("NATS error", zap.Error(err)) + } + }), + nats.ReconnectHandler(func(conn *nats.Conn) { + logger.Info("NATS reconnected", zap.String("provider_id", eventSource.ID), zap.String("url", conn.ConnectedUrlRedacted())) + }), + } + + if eventSource.Authentication != nil { + if eventSource.Authentication.Token != nil { + opts = append(opts, nats.Token(*eventSource.Authentication.Token)) + } else if eventSource.Authentication.UserInfo.Username != nil && eventSource.Authentication.UserInfo.Password != nil { + opts = append(opts, nats.UserInfo(*eventSource.Authentication.UserInfo.Username, *eventSource.Authentication.UserInfo.Password)) + } + } + + return opts, nil +} + +type Nats struct { + providers map[string]*NatsPubSub + logger *zap.Logger + hostName string // How to get it here? + routerListenAddr string // How to get it here? +} + +func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { + definedProviders := make(map[string]bool) + for _, provider := range config.Providers.Nats { + definedProviders[provider.ID] = true + } + usedProviders := make(map[string]bool) + for _, event := range in.CustomEvents.GetNats() { + if _, found := definedProviders[event.EngineEventConfiguration.ProviderId]; !found { + return fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) + } + usedProviders[event.EngineEventConfiguration.ProviderId] = true + } + n.providers = map[string]*NatsPubSub{} + for _, provider := range config.Providers.Nats { + if !usedProviders[provider.ID] { + continue + } + options, err := buildNatsOptions(provider, n.logger) + if err != nil { + return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", provider.ID, err) + } + natsConnection, err := nats.Connect(provider.URL, options...) + if err != nil { + return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", provider.ID, err) + } + js, err := jetstream.New(natsConnection) + if err != nil { + return err + } + + n.providers[provider.ID] = NewConnector(n.logger, natsConnection, js, n.hostName, n.routerListenAddr).New(ctx) + } + return nil +} + +func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *Factory { + return NewFactory(executionContext, config, providers) +} + +func NewPubSub(logger *zap.Logger) Nats { + return Nats{ + logger: logger, + } +} + +type Configuration struct { + Data string `json:"data"` + EventConfiguration []*nodev1.NatsEventConfiguration + Logger *zap.Logger +} + +type Planner struct { + id int + config Configuration + providers map[string]*NatsPubSub + rootFieldRef int + variables resolve.Variables + visitor *plan.Visitor + eventConfig *nodev1.NatsEventConfiguration + publishConfig *PublishAndRequestEventConfiguration + requestConfig *PublishAndRequestEventConfiguration + subscribeConfig *SubscriptionEventConfiguration +} + +func (p *Planner) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { + variablePath, err := p.visitor.Operation.VariablePathByArgumentRefAndArgumentPath(argumentRef, argumentPath, p.visitor.Walker.Ancestors[0].Ref) + if err != nil { + return "", err + } + /* The definition is passed as both definition and operation below because getJSONRootType resolves the type + * from the first argument, but finalInputValueTypeRef comes from the definition + */ + contextVariable := &resolve.ContextVariable{ + Path: variablePath, + Renderer: resolve.NewPlainVariableRenderer(), + } + variablePlaceHolder, _ := p.variables.AddVariable(contextVariable) + return variablePlaceHolder, nil +} + +func (p *Planner) extractEventSubject(fieldRef int, subject string) (string, error) { + matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(subject, -1) + // If no argument templates are defined, there are only static values + if len(matches) < 1 { + if isValidNatsSubject(subject) { + return subject, nil + } + return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) + } + fieldNameBytes := p.visitor.Operation.FieldNameBytes(fieldRef) + // TODO: handling for interfaces and unions + fieldDefinitionRef, ok := p.visitor.Definition.ObjectTypeDefinitionFieldWithName(p.visitor.Walker.EnclosingTypeDefinition.Ref, fieldNameBytes) + if !ok { + return "", fmt.Errorf(`expected field definition to exist for field "%s"`, fieldNameBytes) + } + subjectWithVariableTemplateReplacements := subject + for templateNumber, groups := range matches { + // The first group is the whole template; the second is the period delimited argument path + if len(groups) != 2 { + return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: expected 2 matching groups but received %d`, templateNumber+1, fieldNameBytes, len(groups)-1) + } + validationResult, err := argument_templates.ValidateArgumentPath(p.visitor.Definition, groups[1], fieldDefinitionRef) + if err != nil { + return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: %w`, templateNumber+1, fieldNameBytes, err) + } + argumentNameBytes := []byte(validationResult.ArgumentPath[0]) + argumentRef, ok := p.visitor.Operation.FieldArgument(fieldRef, argumentNameBytes) + if !ok { + return "", fmt.Errorf(`operation field "%s" does not define argument "%s"`, fieldNameBytes, argumentNameBytes) + } + // variablePlaceholder has the form $$0$$, $$1$$, etc. + variablePlaceholder, err := p.addContextVariableByArgumentRef(argumentRef, validationResult.ArgumentPath) + if err != nil { + return "", fmt.Errorf(`failed to retrieve variable placeholder for argument ""%s" defined on operation field "%s": %w`, argumentNameBytes, fieldNameBytes, err) + } + // Replace the template literal with the variable placeholder (and reuse the variable if it already exists) + subjectWithVariableTemplateReplacements = strings.ReplaceAll(subjectWithVariableTemplateReplacements, groups[0], variablePlaceholder) + } + // Substitute the variable templates for dummy values to check naïvely that the string is a valid NATS subject + if isValidNatsSubject(variableTemplateRegex.ReplaceAllLiteralString(subjectWithVariableTemplateReplacements, "a")) { + return subjectWithVariableTemplateReplacements, nil + } + return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) +} + +func (p *Planner) SetID(id int) { + p.id = id +} + +func (p *Planner) ID() (id int) { + return p.id +} + +func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { + // skip, not required + return +} + +func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { + return plan.DataSourcePlanningBehavior{ + MergeAliasedRootNodes: false, + OverrideFieldPathFromAlias: false, + } +} + +func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { + p.visitor = visitor + visitor.Walker.RegisterEnterFieldVisitor(p) + visitor.Walker.RegisterEnterDocumentVisitor(p) + p.config = configuration.CustomConfiguration() + return nil +} + +func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { + var evtCfg PublishEventConfiguration + var dataSource resolve.DataSource + + event, eventErr := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) + if eventErr != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to build event data bytes: %w", eventErr)) + return resolve.FetchConfiguration{} + } + + if p.publishConfig != nil { + pubsub, ok := p.providers[p.publishConfig.ProviderID] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", p.subscribeConfig.ProviderID)) + return resolve.FetchConfiguration{} + } + evtCfg = PublishEventConfiguration{ + ProviderID: p.publishConfig.ProviderID, + Subject: p.publishConfig.Subject, + Data: event, + } + dataSource = &NatsPublishDataSource{ + pubSub: pubsub, + } + } else if p.requestConfig != nil { + pubsub, ok := p.providers[p.requestConfig.ProviderID] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", p.requestConfig.ProviderID)) + return resolve.FetchConfiguration{} + } + dataSource = &NatsRequestDataSource{ + pubSub: pubsub, + } + evtCfg = PublishEventConfiguration{ + ProviderID: p.requestConfig.ProviderID, + Subject: p.requestConfig.Subject, + Data: event, + } + } else { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) + return resolve.FetchConfiguration{} + } + + return resolve.FetchConfiguration{ + Input: evtCfg.MarshalJSONTemplate(), + Variables: p.variables, + DataSource: dataSource, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { + if p.subscribeConfig == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) + return plan.SubscriptionConfiguration{} + } + pubsub, ok := p.providers[p.subscribeConfig.ProviderID] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", p.subscribeConfig.ProviderID)) + return plan.SubscriptionConfiguration{} + } + evtCfg := SubscriptionEventConfiguration{ + ProviderID: p.subscribeConfig.ProviderID, + Subjects: p.subscribeConfig.Subjects, + StreamConfiguration: p.subscribeConfig.StreamConfiguration, + } + object, err := json.Marshal(evtCfg) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) + return plan.SubscriptionConfiguration{} + } + + return plan.SubscriptionConfiguration{ + Input: string(object), + Variables: p.variables, + DataSource: &SubscriptionSource{ + pubSub: pubsub, + }, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner) EnterDocument(_, _ *ast.Document) { + p.rootFieldRef = -1 + p.eventConfig = nil +} + +func (p *Planner) EnterField(ref int) { + if p.rootFieldRef != -1 { + // This is a nested field; nothing needs to be done + return + } + p.rootFieldRef = ref + + fieldName := p.visitor.Operation.FieldNameString(ref) + typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + + var eventConfig *nodev1.NatsEventConfiguration + for _, cfg := range p.config.EventConfiguration { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + eventConfig = cfg + break + } + } + + if eventConfig == nil { + return + } + + p.eventConfig = eventConfig + + providerId := eventConfig.GetEngineEventConfiguration().GetProviderId() + + switch v := eventConfig.GetEngineEventConfiguration().GetType(); v { + case nodev1.EventType_PUBLISH, nodev1.EventType_REQUEST: + if len(p.eventConfig.GetSubjects()) != 1 { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("publish events should define one subject but received %d", len(p.eventConfig.GetSubjects()))) + return + } + _, found := p.providers[providerId] + if !found { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("unable to publish events of provider with id %s", providerId)) + } + extractedSubject, err := p.extractEventSubject(ref, eventConfig.GetSubjects()[0]) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("unable to parse subject with id %s", eventConfig.GetSubjects()[0])) + } + cfg := &PublishAndRequestEventConfiguration{ + ProviderID: providerId, + Subject: extractedSubject, + Data: json.RawMessage("[]"), + } + if v == nodev1.EventType_REQUEST { + p.requestConfig = cfg + } else { + p.publishConfig = cfg + } + case nodev1.EventType_SUBSCRIBE: + if len(p.eventConfig.Subjects) == 0 { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("expected at least one subscription subject but received %d", len(p.eventConfig.Subjects))) + return + } + extractedSubjects := make([]string, 0, len(p.eventConfig.Subjects)) + for _, rawSubject := range p.eventConfig.Subjects { + extractedSubject, err := p.extractEventSubject(ref, rawSubject) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("could not extract subscription event subjects: %w", err)) + return + } + extractedSubjects = append(extractedSubjects, extractedSubject) + } + var streamConf *StreamConfiguration + if p.eventConfig.StreamConfiguration != nil { + streamConf = &StreamConfiguration{} + streamConf.Consumer = p.eventConfig.StreamConfiguration.ConsumerName + streamConf.ConsumerInactiveThreshold = p.eventConfig.StreamConfiguration.ConsumerInactiveThreshold + streamConf.StreamName = p.eventConfig.StreamConfiguration.StreamName + } + + slices.Sort(extractedSubjects) + p.subscribeConfig = &SubscriptionEventConfiguration{ + ProviderID: providerId, + Subjects: extractedSubjects, + StreamConfiguration: streamConf, + } + default: + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Kafka", eventConfig.GetEngineEventConfiguration().GetType())) + } +} + +type Source struct{} + +func (Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { + _, err = out.Write(input) + return +} + +func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { + panic("not implemented") +} + +func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *Factory { + return &Factory{ + executionContext: executionContext, + eventsConfiguration: config, + providers: providers, + } +} + +type Factory struct { + config Configuration + eventsConfiguration config.EventsConfiguration + executionContext context.Context + providers map[string]*NatsPubSub +} + +func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { + return &Planner{ + config: f.config, + providers: f.providers, + } +} + +func (f *Factory) Context() context.Context { + return f.executionContext +} + +func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { + return nil, false +} + +func isValidNatsSubject(subject string) bool { + if subject == "" { + return false + } + sfwc := false + tokens := strings.Split(subject, tsep) + for _, t := range tokens { + length := len(t) + if length == 0 || sfwc { + return false + } + if length > 1 { + if strings.ContainsAny(t, "\t\n\f\r ") { + return false + } + continue + } + switch t[0] { + case fwc: + sfwc = true + case ' ', '\t', '\n', '\r', '\f': + return false + } + } + return true +} diff --git a/router/pkg/pubsub/nats/nats.go b/router/pkg/pubsub/nats/nats.go index c609b55d3f..3528803379 100644 --- a/router/pkg/pubsub/nats/nats.go +++ b/router/pkg/pubsub/nats/nats.go @@ -12,7 +12,6 @@ import ( "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" ) @@ -35,8 +34,8 @@ func NewConnector(logger *zap.Logger, conn *nats.Conn, js jetstream.JetStream, h } } -func (c *connector) New(ctx context.Context) *natsPubSub { - return &natsPubSub{ +func (c *connector) New(ctx context.Context) *NatsPubSub { + return &NatsPubSub{ ctx: ctx, conn: c.conn, js: c.js, @@ -47,7 +46,7 @@ func (c *connector) New(ctx context.Context) *natsPubSub { } } -type natsPubSub struct { +type NatsPubSub struct { ctx context.Context conn *nats.Conn logger *zap.Logger @@ -61,7 +60,7 @@ type natsPubSub struct { // We use the hostname and the address the router is listening on, which should provide a good representation // of what a unique instance is from the perspective of the client that has started a subscription to this instance // and want to restart the subscription after a failure on the client or router side. -func (p *natsPubSub) getInstanceIdentifier() string { +func (p *NatsPubSub) getInstanceIdentifier() string { return fmt.Sprintf("%s-%s", p.hostName, p.routerListenAddr) } @@ -69,7 +68,7 @@ func (p *natsPubSub) getInstanceIdentifier() string { // we need to make sure that the durable consumer name is unique for each instance and subjects to prevent // multiple routers from changing the same consumer, which would lead to message loss and wrong messages delivered // to the subscribers -func (p *natsPubSub) getDurableConsumerName(durableName string, subjects []string) (string, error) { +func (p *NatsPubSub) getDurableConsumerName(durableName string, subjects []string) (string, error) { subjHash := xxhash.New() _, err := subjHash.WriteString(p.getInstanceIdentifier()) if err != nil { @@ -85,7 +84,7 @@ func (p *natsPubSub) getDurableConsumerName(durableName string, subjects []strin return fmt.Sprintf("%s-%x", durableName, subjHash.Sum64()), nil } -func (p *natsPubSub) Subscribe(ctx context.Context, event pubsub_datasource.NatsSubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { +func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "subscribe"), @@ -201,7 +200,7 @@ func (p *natsPubSub) Subscribe(ctx context.Context, event pubsub_datasource.Nats return nil } -func (p *natsPubSub) Publish(_ context.Context, event pubsub_datasource.NatsPublishAndRequestEventConfiguration) error { +func (p *NatsPubSub) Publish(_ context.Context, event PublishAndRequestEventConfiguration) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "publish"), @@ -219,7 +218,7 @@ func (p *natsPubSub) Publish(_ context.Context, event pubsub_datasource.NatsPubl return nil } -func (p *natsPubSub) Request(ctx context.Context, event pubsub_datasource.NatsPublishAndRequestEventConfiguration, w io.Writer) error { +func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "request"), @@ -243,11 +242,11 @@ func (p *natsPubSub) Request(ctx context.Context, event pubsub_datasource.NatsPu return err } -func (p *natsPubSub) flush(ctx context.Context) error { +func (p *NatsPubSub) flush(ctx context.Context) error { return p.conn.FlushWithContext(ctx) } -func (p *natsPubSub) Shutdown(ctx context.Context) error { +func (p *NatsPubSub) Shutdown(ctx context.Context) error { if p.conn.IsClosed() { return nil diff --git a/router/pkg/pubsub/nats/nats_datasource.go b/router/pkg/pubsub/nats/nats_datasource.go deleted file mode 100644 index 428175ec8d..0000000000 --- a/router/pkg/pubsub/nats/nats_datasource.go +++ /dev/null @@ -1,239 +0,0 @@ -package nats - -import ( - "bytes" - "context" - "errors" - "fmt" - "time" - - "github.com/jensneuse/abstractlogger" - "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "go.uber.org/zap" - - "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" -) - -func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { - if natsData := in.GetCustomEvents().GetNats(); natsData != nil { - k := NewPubSub() - err := k.PrepareProviders(ctx, in, dsMeta, config) - if err != nil { - return nil, err - } - factory := k.GetFactory(ctx, config, k.providers) - ds, err := plan.NewDataSourceConfiguration[Configuration]( - in.Id, - factory, - dsMeta, - Configuration{ - EventConfiguration: natsData, - Logger: logger, - }, - ) - - if err != nil { - return nil, err - } - - return ds, nil - } - - return nil, nil -} - -func init() { - datasource.RegisterPubSub(GetDataSource) -} - -func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { - opts := []nats.Option{ - nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), - nats.ReconnectJitter(500*time.Millisecond, 2*time.Second), - nats.ClosedHandler(func(conn *nats.Conn) { - logger.Info("NATS connection closed", zap.String("provider_id", eventSource.ID), zap.Error(conn.LastError())) - }), - nats.ConnectHandler(func(nc *nats.Conn) { - logger.Info("NATS connection established", zap.String("provider_id", eventSource.ID), zap.String("url", nc.ConnectedUrlRedacted())) - }), - nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { - if err != nil { - logger.Error("NATS disconnected; will attempt to reconnect", zap.Error(err), zap.String("provider_id", eventSource.ID)) - } else { - logger.Info("NATS disconnected", zap.String("provider_id", eventSource.ID)) - } - }), - nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { - if errors.Is(err, nats.ErrSlowConsumer) { - logger.Warn( - "NATS slow consumer detected. Events are being dropped. Please consider increasing the buffer size or reducing the number of messages being sent.", - zap.Error(err), - zap.String("provider_id", eventSource.ID), - ) - } else { - logger.Error("NATS error", zap.Error(err)) - } - }), - nats.ReconnectHandler(func(conn *nats.Conn) { - logger.Info("NATS reconnected", zap.String("provider_id", eventSource.ID), zap.String("url", conn.ConnectedUrlRedacted())) - }), - } - - if eventSource.Authentication != nil { - if eventSource.Authentication.Token != nil { - opts = append(opts, nats.Token(*eventSource.Authentication.Token)) - } else if eventSource.Authentication.UserInfo.Username != nil && eventSource.Authentication.UserInfo.Password != nil { - opts = append(opts, nats.UserInfo(*eventSource.Authentication.UserInfo.Username, *eventSource.Authentication.UserInfo.Password)) - } - } - - return opts, nil -} - -type Nats struct { - providers map[string]*natsPubSub - logger *zap.Logger - hostName string // How to get it here? - routerListenAddr string // How to get it here? -} - -func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { - definedProviders := make(map[string]bool) - for _, provider := range config.Providers.Kafka { - definedProviders[provider.ID] = true - } - usedProviders := make(map[string]bool) - for _, event := range in.CustomEvents.GetNats() { - if _, found := definedProviders[event.EngineEventConfiguration.ProviderId]; !found { - return fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) - } - usedProviders[event.EngineEventConfiguration.ProviderId] = true - } - for _, provider := range config.Providers.Nats { - if !usedProviders[provider.ID] { - continue - } - options, err := buildNatsOptions(provider, n.logger) - if err != nil { - return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", provider.ID, err) - } - natsConnection, err := nats.Connect(provider.URL, options...) - if err != nil { - return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", provider.ID, err) - } - js, err := jetstream.New(natsConnection) - if err != nil { - return err - } - - n.providers[provider.ID] = NewConnector(n.logger, natsConnection, js, n.hostName, n.routerListenAddr).New(ctx) - - } - return nil -} - -func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*natsPubSub) *Factory { - return NewFactory(executionContext, config, providers) -} - -func NewPubSub() Nats { - return Nats{} -} - -type Configuration struct { - Data string `json:"data"` - EventConfiguration []*nodev1.NatsEventConfiguration - Logger *zap.Logger -} - -type Planner struct { - id int - config Configuration - providers map[string]*natsPubSub -} - -func (p *Planner) SetID(id int) { - p.id = id -} - -func (p *Planner) ID() (id int) { - return p.id -} - -func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { - // skip, not required - return -} - -func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { - return plan.DataSourcePlanningBehavior{ - MergeAliasedRootNodes: false, - OverrideFieldPathFromAlias: false, - } -} - -func (p *Planner) Register(_ *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { - p.config = configuration.CustomConfiguration() - return nil -} - -func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { - return resolve.FetchConfiguration{ - Input: p.config.Data, - DataSource: Source{}, - } -} - -func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { - return plan.SubscriptionConfiguration{ - Input: p.config.Data, - } -} - -type Source struct{} - -func (Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { - _, err = out.Write(input) - return -} - -func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { - panic("not implemented") -} - -func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*natsPubSub) *Factory { - return &Factory{ - executionContext: executionContext, - eventsConfiguration: config, - providers: providers, - } -} - -type Factory struct { - config Configuration - eventsConfiguration config.EventsConfiguration - executionContext context.Context - providers map[string]*natsPubSub -} - -func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { - return &Planner{ - config: f.config, - providers: f.providers, - } -} - -func (f *Factory) Context() context.Context { - return f.executionContext -} - -func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { - return nil, false -} diff --git a/router/pkg/pubsub/nats/subscription_datasource.go b/router/pkg/pubsub/nats/subscription_datasource.go new file mode 100644 index 0000000000..ef5e2e5fe2 --- /dev/null +++ b/router/pkg/pubsub/nats/subscription_datasource.go @@ -0,0 +1,122 @@ +package nats + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + + "github.com/buger/jsonparser" + "github.com/cespare/xxhash/v2" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type StreamConfiguration struct { + Consumer string `json:"consumer"` + ConsumerInactiveThreshold int32 `json:"consumerInactiveThreshold"` + StreamName string `json:"streamName"` +} + +type SubscriptionEventConfiguration struct { + ProviderID string `json:"providerId"` + Subjects []string `json:"subjects"` + StreamConfiguration *StreamConfiguration `json:"streamConfiguration,omitempty"` +} + +type PublishAndRequestEventConfiguration struct { + ProviderID string `json:"providerId"` + Subject string `json:"subject"` + Data json.RawMessage `json:"data"` +} + +func (s *PublishAndRequestEventConfiguration) MarshalJSONTemplate() string { + return fmt.Sprintf(`{"subject":"%s", "data": %s, "providerId":"%s"}`, s.Subject, s.Data, s.ProviderID) +} + +type PublishEventConfiguration struct { + ProviderID string `json:"providerId"` + Subject string `json:"subject"` + Data json.RawMessage `json:"data"` +} + +func (s *PublishEventConfiguration) MarshalJSONTemplate() string { + return fmt.Sprintf(`{"subject":"%s", "data": %s, "providerId":"%s"}`, s.Subject, s.Data, s.ProviderID) +} + +type SubscriptionSource struct { + pubSub *NatsPubSub +} + +func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { + + val, _, _, err := jsonparser.Get(input, "subjects") + if err != nil { + return err + } + + _, err = xxh.Write(val) + if err != nil { + return err + } + + val, _, _, err = jsonparser.Get(input, "providerId") + if err != nil { + return err + } + + _, err = xxh.Write(val) + return err +} + +func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater resolve.SubscriptionUpdater) error { + var subscriptionConfiguration SubscriptionEventConfiguration + err := json.Unmarshal(input, &subscriptionConfiguration) + if err != nil { + return err + } + + return s.pubSub.Subscribe(ctx.Context(), subscriptionConfiguration, updater) +} + +type NatsPublishDataSource struct { + pubSub *NatsPubSub +} + +func (s *NatsPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { + var publishConfiguration PublishAndRequestEventConfiguration + err := json.Unmarshal(input, &publishConfiguration) + if err != nil { + return err + } + + if err := s.pubSub.Publish(ctx, publishConfiguration); err != nil { + _, err = io.WriteString(out, `{"success": false}`) + return err + } + _, err = io.WriteString(out, `{"success": true}`) + return err +} + +func (s *NatsPublishDataSource) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { + panic("not implemented") +} + +type NatsRequestDataSource struct { + pubSub *NatsPubSub +} + +func (s *NatsRequestDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { + var subscriptionConfiguration PublishAndRequestEventConfiguration + err := json.Unmarshal(input, &subscriptionConfiguration) + if err != nil { + return err + } + + return s.pubSub.Request(ctx, subscriptionConfiguration, out) +} + +func (s *NatsRequestDataSource) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) error { + panic("not implemented") +} diff --git a/router/pkg/pubsub/utils/utils.go b/router/pkg/pubsub/utils/utils.go index 92410b641e..c51775a37f 100644 --- a/router/pkg/pubsub/utils/utils.go +++ b/router/pkg/pubsub/utils/utils.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -func BuildEventDataBytes(ref int, visitor *plan.Visitor, variables resolve.Variables) ([]byte, error) { +func BuildEventDataBytes(ref int, visitor *plan.Visitor, variables *resolve.Variables) ([]byte, error) { // Collect the field arguments for fetch based operations fieldArgs := visitor.Operation.FieldArguments(ref) var dataBuffer bytes.Buffer From 23eb5bb93c117a78358c5a6716505b786abd9604 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 21 Mar 2025 09:46:08 +0100 Subject: [PATCH 05/49] test(kafka): fix mutation kafka test --- router-tests/events/kafka_events_test.go | 69 ++++++++---------------- router-tests/testenv/testenv.go | 4 ++ 2 files changed, 27 insertions(+), 46 deletions(-) diff --git a/router-tests/events/kafka_events_test.go b/router-tests/events/kafka_events_test.go index 36dd03b086..ed386dd47e 100644 --- a/router-tests/events/kafka_events_test.go +++ b/router-tests/events/kafka_events_test.go @@ -1054,7 +1054,7 @@ func TestKafkaEvents(t *testing.T) { t.Run("mutate", func(t *testing.T) { t.Parallel() - topics := []string{"employeeUpdated", "employeeUpdatedTwo"} + topics := []string{"employeeUpdated"} testenv.Run(t, &testenv.Config{ RouterConfigJSONTemplate: testenv.ConfigWithEdfsKafkaJSONTemplate, @@ -1062,56 +1062,16 @@ func TestKafkaEvents(t *testing.T) { }, func(t *testing.T, xEnv *testenv.Environment) { ensureTopicExists(t, xEnv, topics...) - var subscriptionOne struct { - employeeUpdatedMyKafka struct { - ID float64 `graphql:"id"` - Details struct { - Forename string `graphql:"forename"` - Surname string `graphql:"surname"` - } `graphql:"details"` - } `graphql:"employeeUpdatedMyKafka(employeeID: 3)"` - } - - surl := xEnv.GraphQLWebSocketSubscriptionURL() - client := graphql.NewSubscriptionClient(surl) - t.Cleanup(func() { - _ = client.Close() - }) - - var counter atomic.Uint32 - - subscriptionOneID, err := client.Subscribe(&subscriptionOne, nil, func(dataValue []byte, errValue error) error { - defer counter.Add(1) - require.NoError(t, errValue) - require.JSONEq(t, `{"updateEmployeeMyKafka":{"id":1,"details":{"forename":"Jens","surname":"Neuse"}}}`, string(dataValue)) - return nil - }) - require.NoError(t, err) - require.NotEmpty(t, subscriptionOneID) - - go func() { - clientErr := client.Run() - require.NoError(t, clientErr) - }() - - go func() { - require.Eventually(t, func() bool { - return counter.Load() == 1 - }, KafkaWaitTimeout, time.Millisecond*100) - _ = client.Close() - }() - - xEnv.WaitForSubscriptionCount(1, KafkaWaitTimeout) - // Send a mutation to trigger the first subscription resOne := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{ Query: `mutation { updateEmployeeMyKafka(employeeID: 3, update: {name: "name test"}) { success } }`, }) - require.JSONEq(t, `{"data":{"updateEmployeeMyKafka":{"id":3}}}`, resOne.Body) + require.JSONEq(t, `{"data":{"updateEmployeeMyKafka":{"success":true}}}`, resOne.Body) - xEnv.WaitForMessagesSent(1, KafkaWaitTimeout) - xEnv.WaitForSubscriptionCount(0, KafkaWaitTimeout) - xEnv.WaitForConnectionCount(0, KafkaWaitTimeout) + records, err := readKafkaMessages(xEnv, topics[0], 1) + require.NoError(t, err) + require.Equal(t, 1, len(records)) + require.Equal(t, `{"employeeID":3,"update":{"name":"name test"}}`, string(records[0].Value)) }) }) } @@ -1294,3 +1254,20 @@ func produceKafkaMessage(t *testing.T, xEnv *testenv.Environment, topicName stri fErr := xEnv.KafkaClient.Flush(ctx) require.NoError(t, fErr) } + +func readKafkaMessages(xEnv *testenv.Environment, topicName string, msgs int) ([]*kgo.Record, error) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + client, err := kgo.NewClient( + kgo.SeedBrokers(xEnv.GetKafkaSeeds()...), + kgo.ConsumeTopics(xEnv.GetPubSubName(topicName)), + ) + if err != nil { + return nil, err + } + + fetchs := client.PollRecords(ctx, msgs) + + return fetchs.Records(), nil +} diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index 2e0de276b1..fa953fda57 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -1152,6 +1152,10 @@ func (e *Environment) GetPubSubName(name string) string { return e.getPubSubName(name) } +func (e *Environment) GetKafkaSeeds() []string { + return e.cfg.KafkaSeeds +} + func (e *Environment) RouterConfigVersionMain() string { return e.routerConfigVersionMain } From b0660d9e26d6ea16282c0b8d0531c2205046cf56 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Sat, 22 Mar 2025 17:31:53 +0100 Subject: [PATCH 06/49] chore: initial refactor to reduce code needed to start new provider --- router/pkg/pubsub/datasource/datasource.go | 272 ++++++++++++++ router/pkg/pubsub/datasource/old.go | 308 --------------- router/pkg/pubsub/kafka/datasource.go | 238 +++--------- .../pubsub/kafka/subscription_datasource.go | 1 - router/pkg/pubsub/nats/datasource.go | 351 ++++-------------- 5 files changed, 381 insertions(+), 789 deletions(-) delete mode 100644 router/pkg/pubsub/datasource/old.go diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index 20a7ee12ab..19b7259e10 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -1,11 +1,20 @@ package datasource import ( + "bytes" "context" + "fmt" + "strings" + "github.com/jensneuse/abstractlogger" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" ) @@ -21,7 +30,270 @@ func GetRegisteredPubSubs() []Getter { return pubSubs } +type ArgumentTemplateCallback func(tpl string) (string, error) + type PubSubImplementer[F any] interface { PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error GetFactory(executionContext context.Context, config config.EventsConfiguration) F } + +type EventConfigType interface { + GetEngineEventConfiguration() *nodev1.EngineEventConfiguration +} + +// Create an interface for Configuration +type Implementer[EC EventConfigType, P any] interface { + GetResolveDataSource(eventConfig EC, pubsub P) (resolve.DataSource, error) + GetResolveDataSourceSubscription(eventConfig EC, pubsub P) (resolve.SubscriptionDataSource, error) + GetResolveDataSourceSubscriptionInput(eventConfig EC, pubsub P) (string, error) + //GetResolveDataSourceInput(eventConfig EC, ref int, visitor *plan.Visitor, variables *resolve.Variables) (string, error) + GetResolveDataSourceInput(eventConfig EC, event []byte) (string, error) + GetProviderId(eventConfig EC) string + FindEventConfig(eventConfigs []EC, typeName string, fieldName string, fn ArgumentTemplateCallback) (EC, error) + GetEventsDataConfigurations() []EC +} + +type Planner[EC EventConfigType, P any] struct { + id int + config Implementer[EC, P] + eventsConfig config.EventsConfiguration + eventConfig EC + rootFieldRef int + variables resolve.Variables + visitor *plan.Visitor + providers map[string]P +} + +func (p *Planner[EC, P]) SetID(id int) { + p.id = id +} + +func (p *Planner[EC, P]) ID() (id int) { + return p.id +} + +func (p *Planner[EC, P]) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { + // skip, not required + return +} + +func (p *Planner[EC, P]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { + return plan.DataSourcePlanningBehavior{ + MergeAliasedRootNodes: false, + OverrideFieldPathFromAlias: false, + } +} + +func (p *Planner[EC, P]) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Implementer[EC, P]], _ plan.DataSourcePlannerConfiguration) error { + p.visitor = visitor + visitor.Walker.RegisterEnterFieldVisitor(p) + visitor.Walker.RegisterEnterDocumentVisitor(p) + p.config = configuration.CustomConfiguration() + return nil +} + +func (p *Planner[EC, P]) ConfigureFetch() resolve.FetchConfiguration { + if any(p.eventConfig) == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) + return resolve.FetchConfiguration{} + } + + var dataSource resolve.DataSource + providerId := p.config.GetProviderId(p.eventConfig) + pubsub, ok := p.providers[providerId] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) + return resolve.FetchConfiguration{} + } + + dataSource, err := p.config.GetResolveDataSource(p.eventConfig, pubsub) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get data source: %w", err)) + return resolve.FetchConfiguration{} + } + + event, err := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) + return resolve.FetchConfiguration{} + } + + input, err := p.config.GetResolveDataSourceInput(p.eventConfig, event) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) + return resolve.FetchConfiguration{} + } + + return resolve.FetchConfiguration{ + Input: input, + Variables: p.variables, + DataSource: dataSource, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner[EC, P]) ConfigureSubscription() plan.SubscriptionConfiguration { + if any(p.eventConfig) == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) + return plan.SubscriptionConfiguration{} + } + providerId := p.config.GetProviderId(p.eventConfig) + pubsub, ok := p.providers[providerId] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) + return plan.SubscriptionConfiguration{} + } + + dataSource, err := p.config.GetResolveDataSourceSubscription(p.eventConfig, pubsub) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription: %w", err)) + return plan.SubscriptionConfiguration{} + } + + input, err := p.config.GetResolveDataSourceSubscriptionInput(p.eventConfig, pubsub) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription input: %w", err)) + return plan.SubscriptionConfiguration{} + } + + return plan.SubscriptionConfiguration{ + Input: input, + Variables: p.variables, + DataSource: dataSource, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner[EC, P]) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { + variablePath, err := p.visitor.Operation.VariablePathByArgumentRefAndArgumentPath(argumentRef, argumentPath, p.visitor.Walker.Ancestors[0].Ref) + if err != nil { + return "", err + } + /* The definition is passed as both definition and operation below because getJSONRootType resolves the type + * from the first argument, but finalInputValueTypeRef comes from the definition + */ + contextVariable := &resolve.ContextVariable{ + Path: variablePath, + Renderer: resolve.NewPlainVariableRenderer(), + } + variablePlaceHolder, _ := p.variables.AddVariable(contextVariable) + return variablePlaceHolder, nil +} + +func StringParser(subject string) (string, error) { + matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(subject, -1) + if len(matches) < 1 { + return subject, nil + } + return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) +} + +func (p *Planner[EC, P]) extractArgumentTemplate(fieldRef int, template string) (string, error) { + matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(template, -1) + // If no argument templates are defined, there are only static values + if len(matches) < 1 { + return template, nil + } + fieldNameBytes := p.visitor.Operation.FieldNameBytes(fieldRef) + // TODO: handling for interfaces and unions + fieldDefinitionRef, ok := p.visitor.Definition.ObjectTypeDefinitionFieldWithName(p.visitor.Walker.EnclosingTypeDefinition.Ref, fieldNameBytes) + if !ok { + return "", fmt.Errorf(`expected field definition to exist for field "%s"`, fieldNameBytes) + } + templateWithVariableTemplateReplacements := template + for templateNumber, groups := range matches { + // The first group is the whole template; the second is the period delimited argument path + if len(groups) != 2 { + return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: expected 2 matching groups but received %d`, templateNumber+1, fieldNameBytes, len(groups)-1) + } + validationResult, err := argument_templates.ValidateArgumentPath(p.visitor.Definition, groups[1], fieldDefinitionRef) + if err != nil { + return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: %w`, templateNumber+1, fieldNameBytes, err) + } + argumentNameBytes := []byte(validationResult.ArgumentPath[0]) + argumentRef, ok := p.visitor.Operation.FieldArgument(fieldRef, argumentNameBytes) + if !ok { + return "", fmt.Errorf(`operation field "%s" does not define argument "%s"`, fieldNameBytes, argumentNameBytes) + } + // variablePlaceholder has the form $$0$$, $$1$$, etc. + variablePlaceholder, err := p.addContextVariableByArgumentRef(argumentRef, validationResult.ArgumentPath) + if err != nil { + return "", fmt.Errorf(`failed to retrieve variable placeholder for argument ""%s" defined on operation field "%s": %w`, argumentNameBytes, fieldNameBytes, err) + } + // Replace the template literal with the variable placeholder (and reuse the variable if it already exists) + templateWithVariableTemplateReplacements = strings.ReplaceAll(templateWithVariableTemplateReplacements, groups[0], variablePlaceholder) + } + + return templateWithVariableTemplateReplacements, nil +} + +func (p *Planner[EC, P]) EnterDocument(_, _ *ast.Document) { + p.rootFieldRef = -1 + var zero EC + p.eventConfig = zero +} + +func (p *Planner[EC, P]) EnterField(ref int) { + if p.rootFieldRef != -1 { + // This is a nested field; nothing needs to be done + return + } + p.rootFieldRef = ref + + fieldName := p.visitor.Operation.FieldNameString(ref) + typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + + extractFn := func(tpl string) (string, error) { + return p.extractArgumentTemplate(ref, tpl) + } + + eventConfigV, err := p.config.FindEventConfig(p.config.GetEventsDataConfigurations(), typeName, fieldName, extractFn) + if err != nil { + return + } + + p.eventConfig = eventConfigV +} + +type Source struct{} + +func (s *Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { + _, err = out.Write(input) + return +} + +func (s *Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { + panic("not implemented") +} + +func NewFactory[EC EventConfigType, P any](executionContext context.Context, config config.EventsConfiguration, providers map[string]P) *Factory[EC, P] { + return &Factory[EC, P]{ + providers: providers, + executionContext: executionContext, + config: config, + } +} + +type Factory[EC EventConfigType, P any] struct { + providers map[string]P + executionContext context.Context + config config.EventsConfiguration +} + +func (f *Factory[EC, P]) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Implementer[EC, P]] { + return &Planner[EC, P]{ + providers: f.providers, + } +} + +func (f *Factory[EC, P]) Context() context.Context { + return f.executionContext +} + +func (f *Factory[EC, P]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Implementer[EC, P]]) (*ast.Document, bool) { + return nil, false +} diff --git a/router/pkg/pubsub/datasource/old.go b/router/pkg/pubsub/datasource/old.go deleted file mode 100644 index b4528243f5..0000000000 --- a/router/pkg/pubsub/datasource/old.go +++ /dev/null @@ -1,308 +0,0 @@ -package datasource - -import ( - "context" - "fmt" - "regexp" - "strings" - - "github.com/jensneuse/abstractlogger" - - "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" -) - -type EventType string - -const ( - EventTypePublish EventType = "publish" - EventTypeRequest EventType = "request" - EventTypeSubscribe EventType = "subscribe" -) - -var eventSubjectRegex = regexp.MustCompile(`{{ args.([a-zA-Z0-9_]+) }}`) - -func EventTypeFromString(s string) (EventType, error) { - et := EventType(strings.ToLower(s)) - switch et { - case EventTypePublish, EventTypeRequest, EventTypeSubscribe: - return et, nil - default: - return "", fmt.Errorf("invalid event type: %q", s) - } -} - -type EventMetadata struct { - ProviderID string `json:"providerId"` - Type EventType `json:"type"` - TypeName string `json:"typeName"` - FieldName string `json:"fieldName"` -} - -type Configuration struct { - Metadata *EventMetadata `json:"metadata"` - pubSubs []plan.Planner -} - -type Planner struct { - id int - pubSubs []any - eventManager any - rootFieldRef int - variables resolve.Variables - visitor *plan.Visitor - config Configuration -} - -func (p *Planner) SetID(id int) { - p.id = id -} - -func (p *Planner) ID() (id int) { - return p.id -} - -func (p *Planner) EnterField(ref int) { - if p.rootFieldRef != -1 { - // This is a nested field; nothing needs to be done - return - } - p.rootFieldRef = ref - - fieldName := p.visitor.Operation.FieldNameString(ref) - typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) - - //var eventConfig PubSubber - //for _, cfg := range p.config.pubSubs { - // if pubSub, ok := cfg.(PubSubber); ok { - // if pubSub.MatchFieldNameAndType(fieldName, typeName) { - // eventConfig = &cfg - // break - // } - // } - //} - - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("nope fieldName %s and typeName %s", fieldName, typeName)) - - //switch v := eventConfig.Configuration.(type) { - //case *NatsEventConfiguration: - // em := &NatsEventManager{ - // visitor: p.visitor, - // variables: &p.variables, - // eventMetadata: *eventConfig.Metadata, - // eventConfiguration: v, - // } - // p.eventManager = em - // - // switch eventConfig.Metadata.Type { - // case EventTypePublish, EventTypeRequest: - // em.handlePublishAndRequestEvent(ref) - // case EventTypeSubscribe: - // em.handleSubscriptionEvent(ref) - // default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Nats", eventConfig.Metadata.Type)) - // } - //case *KafkaEventConfiguration: - // em := &KafkaEventManager{ - // visitor: p.visitor, - // variables: &p.variables, - // eventMetadata: *eventConfig.Metadata, - // eventConfiguration: v, - // } - // p.eventManager = em - // - // switch eventConfig.Metadata.Type { - // case EventTypePublish: - // em.handlePublishEvent(ref) - // case EventTypeSubscribe: - // em.handleSubscriptionEvent(ref) - // default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Kafka", eventConfig.Metadata.Type)) - // } - //default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid event configuration type: %T", v)) - //} -} - -func (p *Planner) EnterDocument(_, _ *ast.Document) { - p.rootFieldRef = -1 - p.eventManager = nil -} - -func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], dataSourcePlannerConfiguration plan.DataSourcePlannerConfiguration) error { - p.visitor = visitor - visitor.Walker.RegisterEnterFieldVisitor(p) - visitor.Walker.RegisterEnterDocumentVisitor(p) - p.config = configuration.CustomConfiguration() - return nil -} - -func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { - if p.eventManager == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event manager is nil")) - return resolve.FetchConfiguration{} - } - - //var dataSource resolve.DataSource - // - //switch v := p.eventManager.(type) { - //case *NatsEventManager: - // pubsub, ok := p.natsPubSubByProviderID[v.eventMetadata.ProviderID] - // if !ok { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) - // return resolve.FetchConfiguration{} - // } - // - // switch v.eventMetadata.Type { - // case EventTypePublish: - // dataSource = &NatsPublishDataSource{ - // pubSub: pubsub, - // } - // case EventTypeRequest: - // dataSource = &NatsRequestDataSource{ - // pubSub: pubsub, - // } - // default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Nats", v.eventMetadata.Type)) - // return resolve.FetchConfiguration{} - // } - // - // return resolve.FetchConfiguration{ - // Input: v.publishAndRequestEventConfiguration.MarshalJSONTemplate(), - // Variables: p.variables, - // DataSource: dataSource, - // PostProcessing: resolve.PostProcessingConfiguration{ - // MergePath: []string{v.eventMetadata.FieldName}, - // }, - // } - // - //case *KafkaEventManager: - // pubsub, ok := p.kafkaPubSubByProviderID[v.eventMetadata.ProviderID] - // if !ok { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) - // return resolve.FetchConfiguration{} - // } - // - // switch v.eventMetadata.Type { - // case EventTypePublish: - // dataSource = &KafkaPublishDataSource{ - // pubSub: pubsub, - // } - // case EventTypeRequest: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("event type \"%s\" is not supported for Kafka", v.eventMetadata.Type)) - // return resolve.FetchConfiguration{} - // default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", v.eventMetadata.Type)) - // return resolve.FetchConfiguration{} - // } - // - // return resolve.FetchConfiguration{ - // Input: v.publishEventConfiguration.MarshalJSONTemplate(), - // Variables: p.variables, - // DataSource: dataSource, - // PostProcessing: resolve.PostProcessingConfiguration{ - // MergePath: []string{v.eventMetadata.FieldName}, - // }, - // } - // - //default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event manager type: %T", p.eventManager)) - //} - - return resolve.FetchConfiguration{} -} - -func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { - if p.eventManager == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) - return plan.SubscriptionConfiguration{} - } - - //switch v := p.eventManager.(type) { - //case *NatsEventManager: - // pubsub, ok := p.natsPubSubByProviderID[v.eventMetadata.ProviderID] - // if !ok { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) - // return plan.SubscriptionConfiguration{} - // } - // object, err := json.Marshal(v.subscriptionEventConfiguration) - // if err != nil { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) - // return plan.SubscriptionConfiguration{} - // } - // return plan.SubscriptionConfiguration{ - // Input: string(object), - // Variables: p.variables, - // DataSource: &NatsSubscriptionSource{ - // pubSub: pubsub, - // }, - // PostProcessing: resolve.PostProcessingConfiguration{ - // MergePath: []string{v.eventMetadata.FieldName}, - // }, - // } - //case *KafkaEventManager: - // pubsub, ok := p.kafkaPubSubByProviderID[v.eventMetadata.ProviderID] - // if !ok { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", v.eventMetadata.ProviderID)) - // return plan.SubscriptionConfiguration{} - // } - // object, err := json.Marshal(v.subscriptionEventConfiguration) - // if err != nil { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) - // return plan.SubscriptionConfiguration{} - // } - // return plan.SubscriptionConfiguration{ - // Input: string(object), - // Variables: p.variables, - // DataSource: &KafkaSubscriptionSource{ - // pubSub: pubsub, - // }, - // PostProcessing: resolve.PostProcessingConfiguration{ - // MergePath: []string{v.eventMetadata.FieldName}, - // }, - // } - //default: - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: invalid event manager type: %T", p.eventManager)) - //} - - return plan.SubscriptionConfiguration{} -} - -func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { - return plan.DataSourcePlanningBehavior{ - MergeAliasedRootNodes: false, - OverrideFieldPathFromAlias: false, - IncludeTypeNameFields: true, - } -} - -func (p *Planner) DownstreamResponseFieldAlias(_ int) (alias string, exists bool) { - return "", false -} - -func NewFactory(executionContext context.Context, pubSubs []any) *Factory { - return &Factory{ - executionContext: executionContext, - pubSubs: pubSubs, - } -} - -type Factory struct { - executionContext context.Context - pubSubs []any -} - -func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { - return &Planner{ - pubSubs: f.pubSubs, - } -} - -func (f *Factory) Context() context.Context { - return f.executionContext -} - -func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { - return nil, false -} diff --git a/router/pkg/pubsub/kafka/datasource.go b/router/pkg/pubsub/kafka/datasource.go index 3a6fe3d8ee..e242828d2a 100644 --- a/router/pkg/pubsub/kafka/datasource.go +++ b/router/pkg/pubsub/kafka/datasource.go @@ -1,29 +1,24 @@ package kafka import ( - "bytes" "context" "crypto/tls" "encoding/json" "fmt" "time" - "github.com/jensneuse/abstractlogger" "github.com/twmb/franz-go/pkg/kgo" "github.com/twmb/franz-go/pkg/sasl/plain" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" "go.uber.org/zap" - "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { +func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { k := NewPubSub(logger) err := k.PrepareProviders(ctx, in, dsMeta, config) @@ -31,11 +26,11 @@ func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMe return nil, err } factory := k.GetFactory(ctx, config) - ds, err := plan.NewDataSourceConfiguration( + ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.KafkaEventConfiguration, *kafkaPubSub]]( in.Id, factory, dsMeta, - Configuration{ + &Configuration{ EventConfiguration: kafkaData, Logger: logger, }, @@ -52,7 +47,7 @@ func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMe } func init() { - datasource.RegisterPubSub(GetDataSource) + datasource.RegisterPubSub(GetPlanDataSource) } type Kafka struct { @@ -117,8 +112,8 @@ func (k *Kafka) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfi return nil } -func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *Factory { - return NewFactory(executionContext, config, k.providers) +func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *datasource.Factory[*nodev1.KafkaEventConfiguration, *kafkaPubSub] { + return datasource.NewFactory[*nodev1.KafkaEventConfiguration](executionContext, config, k.providers) } func NewPubSub(logger *zap.Logger) Kafka { @@ -134,223 +129,74 @@ type Configuration struct { Logger *zap.Logger } -type Planner struct { - id int - config Configuration - eventsConfig config.EventsConfiguration - eventConfig *nodev1.KafkaEventConfiguration - rootFieldRef int - variables resolve.Variables - visitor *plan.Visitor - providers map[string]*kafkaPubSub +func (c *Configuration) GetEventsDataConfigurations() []*nodev1.KafkaEventConfiguration { + return c.EventConfiguration } -func (p *Planner) SetID(id int) { - p.id = id -} - -func (p *Planner) ID() (id int) { - return p.id -} - -func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { - // skip, not required - return -} - -func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { - return plan.DataSourcePlanningBehavior{ - MergeAliasedRootNodes: false, - OverrideFieldPathFromAlias: false, - } -} - -func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { - p.visitor = visitor - visitor.Walker.RegisterEnterFieldVisitor(p) - visitor.Walker.RegisterEnterDocumentVisitor(p) - p.config = configuration.CustomConfiguration() - return nil -} - -func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { - if p.eventConfig == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) - return resolve.FetchConfiguration{} - } - +func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.KafkaEventConfiguration, pubsub *kafkaPubSub) (resolve.DataSource, error) { var dataSource resolve.DataSource - providerId := p.eventConfig.GetEngineEventConfiguration().GetProviderId() - typeName := p.eventConfig.GetEngineEventConfiguration().GetType() - topics := p.eventConfig.GetTopics() - pubsub, ok := p.providers[providerId] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) - return resolve.FetchConfiguration{} - } - switch p.eventConfig.GetEngineEventConfiguration().GetType() { + typeName := eventConfig.GetEngineEventConfiguration().GetType() + switch typeName { case nodev1.EventType_PUBLISH: dataSource = &KafkaPublishDataSource{ pubSub: pubsub, } default: - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String())) - return resolve.FetchConfiguration{} + return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String()) } - if len(topics) != 1 { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("publish and request events should define one subject but received %d", len(topics))) - return resolve.FetchConfiguration{} - } - - topic := topics[0] - - event, eventErr := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) - if eventErr != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to build event data bytes: %w", eventErr)) - return resolve.FetchConfiguration{} - } - - evtCfg := PublishEventConfiguration{ - ProviderID: providerId, - Topic: topic, - Data: event, - } + return dataSource, nil +} - return resolve.FetchConfiguration{ - Input: evtCfg.MarshalJSONTemplate(), - Variables: p.variables, - DataSource: dataSource, - PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, - }, - } +func (c *Configuration) GetResolveDataSourceSubscription(eventConfig *nodev1.KafkaEventConfiguration, pubsub *kafkaPubSub) (resolve.SubscriptionDataSource, error) { + return &SubscriptionSource{ + pubSub: pubsub, + }, nil } -func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { - if p.eventConfig == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) - return plan.SubscriptionConfiguration{} - } - providerId := p.eventConfig.GetEngineEventConfiguration().GetProviderId() - pubsub, ok := p.providers[providerId] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) - return plan.SubscriptionConfiguration{} - } +func (c *Configuration) GetResolveDataSourceSubscriptionInput(eventConfig *nodev1.KafkaEventConfiguration, pubsub *kafkaPubSub) (string, error) { + providerId := c.GetProviderId(eventConfig) evtCfg := SubscriptionEventConfiguration{ ProviderID: providerId, - Topics: p.eventConfig.GetTopics(), + Topics: eventConfig.GetTopics(), } object, err := json.Marshal(evtCfg) if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) - return plan.SubscriptionConfiguration{} - } - - return plan.SubscriptionConfiguration{ - Input: string(object), - Variables: p.variables, - DataSource: &SubscriptionSource{ - pubSub: pubsub, - }, - PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, - }, + return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") } + return string(object), nil } -func (p *Planner) EnterDocument(_, _ *ast.Document) { - p.rootFieldRef = -1 - p.eventConfig = nil -} - -func (p *Planner) EnterField(ref int) { - if p.rootFieldRef != -1 { - // This is a nested field; nothing needs to be done - return - } - p.rootFieldRef = ref - - fieldName := p.visitor.Operation.FieldNameString(ref) - typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) - - var eventConfig *nodev1.KafkaEventConfiguration - for _, cfg := range p.config.EventConfiguration { - if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - eventConfig = cfg - break - } - } +func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.KafkaEventConfiguration, event []byte) (string, error) { + topics := eventConfig.GetTopics() - if eventConfig == nil { - return + if len(topics) != 1 { + return "", fmt.Errorf("publish and request events should define one topic but received %d", len(topics)) } - p.eventConfig = eventConfig + topic := topics[0] - providerId := eventConfig.GetEngineEventConfiguration().GetProviderId() + providerId := c.GetProviderId(eventConfig) - switch eventConfig.GetEngineEventConfiguration().GetType() { - case nodev1.EventType_PUBLISH: - if len(p.eventConfig.GetTopics()) != 1 { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("publish events should define one subject but received %d", len(p.eventConfig.GetTopics()))) - return - } - _, found := p.providers[providerId] - if !found { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("unable to publish events of provider %d", len(p.eventConfig.GetTopics()))) - } - _ = PublishEventConfiguration{ - ProviderID: providerId, - Topic: eventConfig.GetTopics()[0], - Data: json.RawMessage("[]"), - } - p.config.Logger.Warn("Publishing!") - // provider.Publish(provider.ctx, pubCfg) - case nodev1.EventType_SUBSCRIBE: - p.config.Logger.Warn("Subscribing!") - default: - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Kafka", eventConfig.GetEngineEventConfiguration().GetType())) + evtCfg := PublishEventConfiguration{ + ProviderID: providerId, + Topic: topic, + Data: event, } -} - -type Source struct{} - -func (s *Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { - _, err = out.Write(input) - return -} - -func (s *Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { - panic("not implemented") -} -func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*kafkaPubSub) *Factory { - return &Factory{ - providers: providers, - executionContext: executionContext, - config: config, - } + return evtCfg.MarshalJSONTemplate(), nil } -type Factory struct { - providers map[string]*kafkaPubSub - executionContext context.Context - config config.EventsConfiguration +func (c *Configuration) GetProviderId(eventConfig *nodev1.KafkaEventConfiguration) string { + return eventConfig.GetEngineEventConfiguration().GetProviderId() } -func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { - return &Planner{ - providers: f.providers, +func (c *Configuration) FindEventConfig(eventConfigs []*nodev1.KafkaEventConfiguration, typeName string, fieldName string, fn datasource.ArgumentTemplateCallback) (*nodev1.KafkaEventConfiguration, error) { + for _, cfg := range eventConfigs { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + return cfg, nil + } } -} - -func (f *Factory) Context() context.Context { - return f.executionContext -} - -func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { - return nil, false + return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) } diff --git a/router/pkg/pubsub/kafka/subscription_datasource.go b/router/pkg/pubsub/kafka/subscription_datasource.go index 3109d3cf35..70623170cf 100644 --- a/router/pkg/pubsub/kafka/subscription_datasource.go +++ b/router/pkg/pubsub/kafka/subscription_datasource.go @@ -33,7 +33,6 @@ type SubscriptionSource struct { } func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { - val, _, _, err := jsonparser.Get(input, "topics") if err != nil { return err diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go index b85deed14f..f970e3be9b 100644 --- a/router/pkg/pubsub/nats/datasource.go +++ b/router/pkg/pubsub/nats/datasource.go @@ -1,7 +1,6 @@ package nats import ( - "bytes" "context" "encoding/json" "errors" @@ -11,18 +10,13 @@ import ( "strings" "time" - "github.com/jensneuse/abstractlogger" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" "go.uber.org/zap" - "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) @@ -45,11 +39,11 @@ func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMe return nil, err } factory := k.GetFactory(ctx, config, k.providers) - ds, err := plan.NewDataSourceConfiguration[Configuration]( + ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.NatsEventConfiguration, *NatsPubSub]]( in.Id, factory, dsMeta, - Configuration{ + &Configuration{ EventConfiguration: natsData, Logger: logger, }, @@ -155,8 +149,8 @@ func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfig return nil } -func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *Factory { - return NewFactory(executionContext, config, providers) +func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *datasource.Factory[*nodev1.NatsEventConfiguration, *NatsPubSub] { + return datasource.NewFactory[*nodev1.NatsEventConfiguration](executionContext, config, n.providers) } func NewPubSub(logger *zap.Logger) Nats { @@ -169,323 +163,112 @@ type Configuration struct { Data string `json:"data"` EventConfiguration []*nodev1.NatsEventConfiguration Logger *zap.Logger + requestConfig *PublishAndRequestEventConfiguration + publishConfig *PublishAndRequestEventConfiguration + subscribeConfig *SubscriptionEventConfiguration } -type Planner struct { - id int - config Configuration - providers map[string]*NatsPubSub - rootFieldRef int - variables resolve.Variables - visitor *plan.Visitor - eventConfig *nodev1.NatsEventConfiguration - publishConfig *PublishAndRequestEventConfiguration - requestConfig *PublishAndRequestEventConfiguration - subscribeConfig *SubscriptionEventConfiguration +func (c *Configuration) GetEventsDataConfigurations() []*nodev1.NatsEventConfiguration { + return c.EventConfiguration } -func (p *Planner) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { - variablePath, err := p.visitor.Operation.VariablePathByArgumentRefAndArgumentPath(argumentRef, argumentPath, p.visitor.Walker.Ancestors[0].Ref) - if err != nil { - return "", err - } - /* The definition is passed as both definition and operation below because getJSONRootType resolves the type - * from the first argument, but finalInputValueTypeRef comes from the definition - */ - contextVariable := &resolve.ContextVariable{ - Path: variablePath, - Renderer: resolve.NewPlainVariableRenderer(), - } - variablePlaceHolder, _ := p.variables.AddVariable(contextVariable) - return variablePlaceHolder, nil -} +func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.NatsEventConfiguration, pubsub *NatsPubSub) (resolve.DataSource, error) { + var dataSource resolve.DataSource -func (p *Planner) extractEventSubject(fieldRef int, subject string) (string, error) { - matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(subject, -1) - // If no argument templates are defined, there are only static values - if len(matches) < 1 { - if isValidNatsSubject(subject) { - return subject, nil - } - return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) - } - fieldNameBytes := p.visitor.Operation.FieldNameBytes(fieldRef) - // TODO: handling for interfaces and unions - fieldDefinitionRef, ok := p.visitor.Definition.ObjectTypeDefinitionFieldWithName(p.visitor.Walker.EnclosingTypeDefinition.Ref, fieldNameBytes) - if !ok { - return "", fmt.Errorf(`expected field definition to exist for field "%s"`, fieldNameBytes) - } - subjectWithVariableTemplateReplacements := subject - for templateNumber, groups := range matches { - // The first group is the whole template; the second is the period delimited argument path - if len(groups) != 2 { - return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: expected 2 matching groups but received %d`, templateNumber+1, fieldNameBytes, len(groups)-1) - } - validationResult, err := argument_templates.ValidateArgumentPath(p.visitor.Definition, groups[1], fieldDefinitionRef) - if err != nil { - return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: %w`, templateNumber+1, fieldNameBytes, err) - } - argumentNameBytes := []byte(validationResult.ArgumentPath[0]) - argumentRef, ok := p.visitor.Operation.FieldArgument(fieldRef, argumentNameBytes) - if !ok { - return "", fmt.Errorf(`operation field "%s" does not define argument "%s"`, fieldNameBytes, argumentNameBytes) - } - // variablePlaceholder has the form $$0$$, $$1$$, etc. - variablePlaceholder, err := p.addContextVariableByArgumentRef(argumentRef, validationResult.ArgumentPath) - if err != nil { - return "", fmt.Errorf(`failed to retrieve variable placeholder for argument ""%s" defined on operation field "%s": %w`, argumentNameBytes, fieldNameBytes, err) + typeName := eventConfig.GetEngineEventConfiguration().GetType() + switch typeName { + case nodev1.EventType_PUBLISH, nodev1.EventType_REQUEST: + dataSource = &NatsPublishDataSource{ + pubSub: pubsub, } - // Replace the template literal with the variable placeholder (and reuse the variable if it already exists) - subjectWithVariableTemplateReplacements = strings.ReplaceAll(subjectWithVariableTemplateReplacements, groups[0], variablePlaceholder) - } - // Substitute the variable templates for dummy values to check naïvely that the string is a valid NATS subject - if isValidNatsSubject(variableTemplateRegex.ReplaceAllLiteralString(subjectWithVariableTemplateReplacements, "a")) { - return subjectWithVariableTemplateReplacements, nil + default: + return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Nats", typeName.String()) } - return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) -} - -func (p *Planner) SetID(id int) { - p.id = id -} - -func (p *Planner) ID() (id int) { - return p.id -} - -func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { - // skip, not required - return -} -func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { - return plan.DataSourcePlanningBehavior{ - MergeAliasedRootNodes: false, - OverrideFieldPathFromAlias: false, - } + return dataSource, nil } -func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Configuration], _ plan.DataSourcePlannerConfiguration) error { - p.visitor = visitor - visitor.Walker.RegisterEnterFieldVisitor(p) - visitor.Walker.RegisterEnterDocumentVisitor(p) - p.config = configuration.CustomConfiguration() - return nil +func (c *Configuration) GetResolveDataSourceSubscription(eventConfig *nodev1.NatsEventConfiguration, pubsub *NatsPubSub) (resolve.SubscriptionDataSource, error) { + return &SubscriptionSource{ + pubSub: pubsub, + }, nil } -func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { - var evtCfg PublishEventConfiguration - var dataSource resolve.DataSource +func (c *Configuration) GetResolveDataSourceSubscriptionInput(eventConfig *nodev1.NatsEventConfiguration, pubsub *NatsPubSub) (string, error) { + providerId := c.GetProviderId(eventConfig) - event, eventErr := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) - if eventErr != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to build event data bytes: %w", eventErr)) - return resolve.FetchConfiguration{} + evtCfg := SubscriptionEventConfiguration{ + ProviderID: providerId, + Subjects: eventConfig.GetSubjects(), } - - if p.publishConfig != nil { - pubsub, ok := p.providers[p.publishConfig.ProviderID] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", p.subscribeConfig.ProviderID)) - return resolve.FetchConfiguration{} + if eventConfig.StreamConfiguration != nil { + evtCfg.StreamConfiguration = &StreamConfiguration{ + Consumer: eventConfig.StreamConfiguration.ConsumerName, + StreamName: eventConfig.StreamConfiguration.StreamName, + ConsumerInactiveThreshold: eventConfig.StreamConfiguration.ConsumerInactiveThreshold, } - evtCfg = PublishEventConfiguration{ - ProviderID: p.publishConfig.ProviderID, - Subject: p.publishConfig.Subject, - Data: event, - } - dataSource = &NatsPublishDataSource{ - pubSub: pubsub, - } - } else if p.requestConfig != nil { - pubsub, ok := p.providers[p.requestConfig.ProviderID] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", p.requestConfig.ProviderID)) - return resolve.FetchConfiguration{} - } - dataSource = &NatsRequestDataSource{ - pubSub: pubsub, - } - evtCfg = PublishEventConfiguration{ - ProviderID: p.requestConfig.ProviderID, - Subject: p.requestConfig.Subject, - Data: event, - } - } else { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) - return resolve.FetchConfiguration{} - } - - return resolve.FetchConfiguration{ - Input: evtCfg.MarshalJSONTemplate(), - Variables: p.variables, - DataSource: dataSource, - PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, - }, - } -} - -func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { - if p.subscribeConfig == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) - return plan.SubscriptionConfiguration{} - } - pubsub, ok := p.providers[p.subscribeConfig.ProviderID] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", p.subscribeConfig.ProviderID)) - return plan.SubscriptionConfiguration{} - } - evtCfg := SubscriptionEventConfiguration{ - ProviderID: p.subscribeConfig.ProviderID, - Subjects: p.subscribeConfig.Subjects, - StreamConfiguration: p.subscribeConfig.StreamConfiguration, } object, err := json.Marshal(evtCfg) if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to marshal event subscription streamConfiguration")) - return plan.SubscriptionConfiguration{} - } - - return plan.SubscriptionConfiguration{ - Input: string(object), - Variables: p.variables, - DataSource: &SubscriptionSource{ - pubSub: pubsub, - }, - PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, - }, + return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") } + return string(object), nil } -func (p *Planner) EnterDocument(_, _ *ast.Document) { - p.rootFieldRef = -1 - p.eventConfig = nil -} +func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.NatsEventConfiguration, event []byte) (string, error) { + subjects := eventConfig.GetSubjects() -func (p *Planner) EnterField(ref int) { - if p.rootFieldRef != -1 { - // This is a nested field; nothing needs to be done - return + if len(subjects) != 1 { + return "", fmt.Errorf("publish and request events should define one subject but received %d", len(subjects)) } - p.rootFieldRef = ref - fieldName := p.visitor.Operation.FieldNameString(ref) - typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + subject := subjects[0] - var eventConfig *nodev1.NatsEventConfiguration - for _, cfg := range p.config.EventConfiguration { - if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - eventConfig = cfg - break - } - } + providerId := c.GetProviderId(eventConfig) - if eventConfig == nil { - return + evtCfg := PublishEventConfiguration{ + ProviderID: providerId, + Subject: subject, + Data: event, } - p.eventConfig = eventConfig + return evtCfg.MarshalJSONTemplate(), nil +} - providerId := eventConfig.GetEngineEventConfiguration().GetProviderId() +func (c *Configuration) GetProviderId(eventConfig *nodev1.NatsEventConfiguration) string { + return eventConfig.GetEngineEventConfiguration().GetProviderId() +} - switch v := eventConfig.GetEngineEventConfiguration().GetType(); v { +func (c *Configuration) transformEventConfig(cfg *nodev1.NatsEventConfiguration, fn datasource.ArgumentTemplateCallback) (*nodev1.NatsEventConfiguration, error) { + switch v := cfg.GetEngineEventConfiguration().GetType(); v { case nodev1.EventType_PUBLISH, nodev1.EventType_REQUEST: - if len(p.eventConfig.GetSubjects()) != 1 { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("publish events should define one subject but received %d", len(p.eventConfig.GetSubjects()))) - return - } - _, found := p.providers[providerId] - if !found { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("unable to publish events of provider with id %s", providerId)) - } - extractedSubject, err := p.extractEventSubject(ref, eventConfig.GetSubjects()[0]) + extractedSubject, err := fn(cfg.GetSubjects()[0]) if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("unable to parse subject with id %s", eventConfig.GetSubjects()[0])) - } - cfg := &PublishAndRequestEventConfiguration{ - ProviderID: providerId, - Subject: extractedSubject, - Data: json.RawMessage("[]"), - } - if v == nodev1.EventType_REQUEST { - p.requestConfig = cfg - } else { - p.publishConfig = cfg + return cfg, fmt.Errorf("unable to parse subject with id %s", cfg.GetSubjects()[0]) } + cfg.Subjects = []string{extractedSubject} case nodev1.EventType_SUBSCRIBE: - if len(p.eventConfig.Subjects) == 0 { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("expected at least one subscription subject but received %d", len(p.eventConfig.Subjects))) - return - } - extractedSubjects := make([]string, 0, len(p.eventConfig.Subjects)) - for _, rawSubject := range p.eventConfig.Subjects { - extractedSubject, err := p.extractEventSubject(ref, rawSubject) + extractedSubjects := make([]string, 0, len(cfg.Subjects)) + for _, rawSubject := range cfg.Subjects { + extractedSubject, err := fn(rawSubject) if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("could not extract subscription event subjects: %w", err)) - return + return cfg, nil } extractedSubjects = append(extractedSubjects, extractedSubject) } - var streamConf *StreamConfiguration - if p.eventConfig.StreamConfiguration != nil { - streamConf = &StreamConfiguration{} - streamConf.Consumer = p.eventConfig.StreamConfiguration.ConsumerName - streamConf.ConsumerInactiveThreshold = p.eventConfig.StreamConfiguration.ConsumerInactiveThreshold - streamConf.StreamName = p.eventConfig.StreamConfiguration.StreamName - } - slices.Sort(extractedSubjects) - p.subscribeConfig = &SubscriptionEventConfiguration{ - ProviderID: providerId, - Subjects: extractedSubjects, - StreamConfiguration: streamConf, - } - default: - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("invalid EventType \"%s\" for Kafka", eventConfig.GetEngineEventConfiguration().GetType())) - } -} - -type Source struct{} - -func (Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { - _, err = out.Write(input) - return -} - -func (Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { - panic("not implemented") -} - -func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *Factory { - return &Factory{ - executionContext: executionContext, - eventsConfiguration: config, - providers: providers, + cfg.Subjects = extractedSubjects } + return cfg, nil } -type Factory struct { - config Configuration - eventsConfiguration config.EventsConfiguration - executionContext context.Context - providers map[string]*NatsPubSub -} - -func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Configuration] { - return &Planner{ - config: f.config, - providers: f.providers, +func (c *Configuration) FindEventConfig(eventConfigs []*nodev1.NatsEventConfiguration, typeName string, fieldName string, fn datasource.ArgumentTemplateCallback) (*nodev1.NatsEventConfiguration, error) { + for _, cfg := range eventConfigs { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + return c.transformEventConfig(cfg, fn) + } } -} - -func (f *Factory) Context() context.Context { - return f.executionContext -} - -func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Configuration]) (*ast.Document, bool) { - return nil, false + return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) } func isValidNatsSubject(subject string) bool { From 163aa5b6f35ca954dab266276e96110d2be349ca Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 24 Mar 2025 18:49:55 +0100 Subject: [PATCH 07/49] refactor(pubsub): implement lazy connection handling for NATS and Kafka, streamline provider initialization --- demo/pkg/subgraphs/subgraphs.go | 22 +- router-tests/events/nats_events_test.go | 2 +- router-tests/structured_logging_test.go | 7 +- router-tests/testenv/pubsub.go | 19 ++ router-tests/testenv/testenv.go | 6 +- router/core/plan_generator.go | 31 ++- router/pkg/plan_generator/plan_generator.go | 1 + .../pkg/plan_generator/plan_generator_test.go | 2 + router/pkg/pubsub/datasource/datasource.go | 254 +----------------- router/pkg/pubsub/datasource/factory.go | 38 +++ router/pkg/pubsub/datasource/planner.go | 219 +++++++++++++++ router/pkg/pubsub/kafka/datasource.go | 102 ++----- router/pkg/pubsub/kafka/datasource_impl.go | 89 ++++++ router/pkg/pubsub/kafka/kafka.go | 68 +++-- .../pubsub/kafka/subscription_datasource.go | 4 +- router/pkg/pubsub/nats/datasource.go | 59 +++- router/pkg/pubsub/nats/nats.go | 32 +-- 17 files changed, 545 insertions(+), 410 deletions(-) create mode 100644 router/pkg/pubsub/datasource/factory.go create mode 100644 router/pkg/pubsub/datasource/planner.go create mode 100644 router/pkg/pubsub/kafka/datasource_impl.go diff --git a/demo/pkg/subgraphs/subgraphs.go b/demo/pkg/subgraphs/subgraphs.go index ab3820d02c..0b43b17aff 100644 --- a/demo/pkg/subgraphs/subgraphs.go +++ b/demo/pkg/subgraphs/subgraphs.go @@ -206,31 +206,21 @@ func New(ctx context.Context, config *Config) (*Subgraphs, error) { if defaultSourceNameURL := os.Getenv("NATS_URL"); defaultSourceNameURL != "" { url = defaultSourceNameURL } - defaultConnection, err := nats.Connect(url) - if err != nil { - log.Printf("failed to connect to nats source \"nats\": %v", err) - } - myNatsConnection, err := nats.Connect(url) - if err != nil { - log.Printf("failed to connect to nats source \"my-nats\": %v", err) + natsPubSubByProviderID := map[string]*natsPubsub.NatsPubSub{ + "default": natsPubsub.NewConnector(zap.NewNop(), url, []nats.Option{}, "hostname", "test").New(ctx), + "my-nats": natsPubsub.NewConnector(zap.NewNop(), url, []nats.Option{}, "hostname", "test").New(ctx), } - defaultJetStream, err := jetstream.New(defaultConnection) + defaultConnection, err := nats.Connect(url) if err != nil { - return nil, err + log.Printf("failed to connect to nats source \"nats\": %v", err) } - - myNatsJetStream, err := jetstream.New(myNatsConnection) + defaultJetStream, err := jetstream.New(defaultConnection) if err != nil { return nil, err } - natsPubSubByProviderID := map[string]*natsPubsub.NatsPubSub{ - "default": natsPubsub.NewConnector(zap.NewNop(), defaultConnection, defaultJetStream, "hostname", "test").New(ctx), - "my-nats": natsPubsub.NewConnector(zap.NewNop(), myNatsConnection, myNatsJetStream, "hostname", "test").New(ctx), - } - _, err = defaultJetStream.CreateOrUpdateStream(ctx, jetstream.StreamConfig{ Name: "streamName", Subjects: []string{"employeeUpdated.>"}, diff --git a/router-tests/events/nats_events_test.go b/router-tests/events/nats_events_test.go index fdd068a083..6659b8f234 100644 --- a/router-tests/events/nats_events_test.go +++ b/router-tests/events/nats_events_test.go @@ -136,7 +136,7 @@ func TestNatsEvents(t *testing.T) { natsLogs := xEnv.Observer().FilterMessageSnippet("Nats").All() require.Len(t, natsLogs, 4) providerIDFields := xEnv.Observer().FilterField(zap.String("provider_id", "my-nats")).All() - require.Len(t, providerIDFields, 3) + require.Len(t, providerIDFields, 1) }) }) diff --git a/router-tests/structured_logging_test.go b/router-tests/structured_logging_test.go index ba6a2e4a96..257558b992 100644 --- a/router-tests/structured_logging_test.go +++ b/router-tests/structured_logging_test.go @@ -4,13 +4,14 @@ import ( "bytes" "encoding/json" "fmt" - "github.com/stretchr/testify/assert" "math" "net/http" "os" "path/filepath" "testing" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "go.uber.org/zap" "go.uber.org/zap/zapcore" @@ -162,11 +163,11 @@ func TestRouterStartLogs(t *testing.T) { }, }, func(t *testing.T, xEnv *testenv.Environment) { logEntries := xEnv.Observer().All() - require.Len(t, logEntries, 13) + require.Len(t, logEntries, 11) natsLogs := xEnv.Observer().FilterMessageSnippet("Nats Event source enabled").All() require.Len(t, natsLogs, 4) providerIDFields := xEnv.Observer().FilterField(zap.String("provider_id", "default")).All() - require.Len(t, providerIDFields, 2) + require.Len(t, providerIDFields, 1) kafkaLogs := xEnv.Observer().FilterMessageSnippet("Kafka Event source enabled").All() require.Len(t, kafkaLogs, 2) playgroundLog := xEnv.Observer().FilterMessage("Serving GraphQL playground") diff --git a/router-tests/testenv/pubsub.go b/router-tests/testenv/pubsub.go index 6842344d1e..25048e1981 100644 --- a/router-tests/testenv/pubsub.go +++ b/router-tests/testenv/pubsub.go @@ -137,8 +137,14 @@ var ( natsServer *natsserver.Server ) +type NatParams struct { + Opts []nats.Option + Url string +} + type NatsData struct { Connections []*nats.Conn + Params []*NatParams Server *natsserver.Server } @@ -148,6 +154,17 @@ func setupNatsData(t testing.TB) (*NatsData, error) { } natsData.Server = natsServer for range demoNatsProviders { + param := &NatParams{ + Url: natsData.Server.ClientURL(), + Opts: []nats.Option{ + nats.MaxReconnects(10), + nats.ReconnectWait(1 * time.Second), + nats.Timeout(5 * time.Second), + nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { + t.Log(err) + }), + }, + } natsConnection, err := nats.Connect( natsData.Server.ClientURL(), nats.MaxReconnects(10), @@ -160,6 +177,8 @@ func setupNatsData(t testing.TB) (*NatsData, error) { if err != nil { return nil, err } + + natsData.Params = append(natsData.Params, param) natsData.Connections = append(natsData.Connections, natsConnection) } return natsData, nil diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index fa953fda57..1c38a48a29 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -37,7 +37,6 @@ import ( "github.com/hashicorp/go-cleanhttp" "github.com/hashicorp/go-retryablehttp" "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/require" "github.com/twmb/franz-go/pkg/kadm" @@ -2040,10 +2039,7 @@ func subgraphOptions(ctx context.Context, t testing.TB, logger *zap.Logger, nats } natsPubSubByProviderID := make(map[string]*pubsubNats.NatsPubSub, len(demoNatsProviders)) for _, sourceName := range demoNatsProviders { - js, err := jetstream.New(natsData.Connections[0]) - require.NoError(t, err) - - natsPubSubByProviderID[sourceName] = pubsubNats.NewConnector(logger, natsData.Connections[0], js, "hostname", "listenaddr").New(ctx) + natsPubSubByProviderID[sourceName] = pubsubNats.NewConnector(logger, natsData.Params[0].Url, natsData.Params[0].Opts, "hostname", "listenaddr").New(ctx) } return &subgraphs.SubgraphOptions{ diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index a2628f5a54..6b59a564f6 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -22,6 +22,8 @@ import ( "go.uber.org/zap" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/cosmo/router/pkg/execution_config" ) @@ -158,6 +160,33 @@ func (pg *PlanGenerator) buildRouterConfig(configFilePath string) (*nodev1.Route } func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, logger *zap.Logger, maxDataSourceCollectorsConcurrency uint) error { + routerEngineConfig := RouterEngineConfiguration{} + natSources := map[string]*nats.NatsPubSub{} + kafkaSources := map[string]*kafka.KafkaPubSub{} + for _, ds := range routerConfig.GetEngineConfig().GetDatasourceConfigurations() { + if ds.GetKind() != nodev1.DataSourceKind_PUBSUB || ds.GetCustomEvents() == nil { + continue + } + for _, natConfig := range ds.GetCustomEvents().GetNats() { + providerId := natConfig.GetEngineEventConfiguration().GetProviderId() + if _, ok := natSources[providerId]; !ok { + natSources[providerId] = nil + routerEngineConfig.Events.Providers.Nats = append(routerEngineConfig.Events.Providers.Nats, config.NatsEventSource{ + ID: providerId, + }) + } + } + for _, kafkaConfig := range ds.GetCustomEvents().GetKafka() { + providerId := kafkaConfig.GetEngineEventConfiguration().GetProviderId() + if _, ok := kafkaSources[providerId]; !ok { + kafkaSources[providerId] = nil + routerEngineConfig.Events.Providers.Kafka = append(routerEngineConfig.Events.Providers.Kafka, config.KafkaEventSource{ + ID: providerId, + }) + } + } + } + var netPollConfig graphql_datasource.NetPollConfiguration netPollConfig.ApplyDefaults() @@ -178,7 +207,7 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo }, logger) // this generates the plan configuration using the data source factories from the config package - planConfig, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &RouterEngineConfiguration{}) + planConfig, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &routerEngineConfig) if err != nil { return fmt.Errorf("failed to load configuration: %w", err) } diff --git a/router/pkg/plan_generator/plan_generator.go b/router/pkg/plan_generator/plan_generator.go index de572d4a08..f0afcc81a2 100644 --- a/router/pkg/plan_generator/plan_generator.go +++ b/router/pkg/plan_generator/plan_generator.go @@ -116,6 +116,7 @@ func PlanGenerator(ctx context.Context, cfg QueryPlanConfig) error { planner, err := pg.GetPlanner() if err != nil { cancelError(fmt.Errorf("failed to get planner: %v", err)) + return } for { select { diff --git a/router/pkg/plan_generator/plan_generator_test.go b/router/pkg/plan_generator/plan_generator_test.go index 3157fa710c..0f6a574522 100644 --- a/router/pkg/plan_generator/plan_generator_test.go +++ b/router/pkg/plan_generator/plan_generator_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/zap" ) func getTestDataDir() string { @@ -130,6 +131,7 @@ func TestPlanGenerator(t *testing.T) { ExecutionConfig: path.Join(getTestDataDir(), "execution_config", "base.json"), Timeout: "30s", OutputFiles: true, + Logger: zap.NewNop(), } err = PlanGenerator(context.Background(), cfg) diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index 19b7259e10..f76994b1bd 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -1,18 +1,10 @@ package datasource import ( - "bytes" "context" - "fmt" - "strings" - "github.com/jensneuse/abstractlogger" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" - "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" @@ -34,6 +26,7 @@ type ArgumentTemplateCallback func(tpl string) (string, error) type PubSubImplementer[F any] interface { PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error + // ConnectProviders(ctx context.Context) error GetFactory(executionContext context.Context, config config.EventsConfiguration) F } @@ -52,248 +45,3 @@ type Implementer[EC EventConfigType, P any] interface { FindEventConfig(eventConfigs []EC, typeName string, fieldName string, fn ArgumentTemplateCallback) (EC, error) GetEventsDataConfigurations() []EC } - -type Planner[EC EventConfigType, P any] struct { - id int - config Implementer[EC, P] - eventsConfig config.EventsConfiguration - eventConfig EC - rootFieldRef int - variables resolve.Variables - visitor *plan.Visitor - providers map[string]P -} - -func (p *Planner[EC, P]) SetID(id int) { - p.id = id -} - -func (p *Planner[EC, P]) ID() (id int) { - return p.id -} - -func (p *Planner[EC, P]) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { - // skip, not required - return -} - -func (p *Planner[EC, P]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { - return plan.DataSourcePlanningBehavior{ - MergeAliasedRootNodes: false, - OverrideFieldPathFromAlias: false, - } -} - -func (p *Planner[EC, P]) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Implementer[EC, P]], _ plan.DataSourcePlannerConfiguration) error { - p.visitor = visitor - visitor.Walker.RegisterEnterFieldVisitor(p) - visitor.Walker.RegisterEnterDocumentVisitor(p) - p.config = configuration.CustomConfiguration() - return nil -} - -func (p *Planner[EC, P]) ConfigureFetch() resolve.FetchConfiguration { - if any(p.eventConfig) == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) - return resolve.FetchConfiguration{} - } - - var dataSource resolve.DataSource - providerId := p.config.GetProviderId(p.eventConfig) - pubsub, ok := p.providers[providerId] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) - return resolve.FetchConfiguration{} - } - - dataSource, err := p.config.GetResolveDataSource(p.eventConfig, pubsub) - if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get data source: %w", err)) - return resolve.FetchConfiguration{} - } - - event, err := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) - if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) - return resolve.FetchConfiguration{} - } - - input, err := p.config.GetResolveDataSourceInput(p.eventConfig, event) - if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) - return resolve.FetchConfiguration{} - } - - return resolve.FetchConfiguration{ - Input: input, - Variables: p.variables, - DataSource: dataSource, - PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, - }, - } -} - -func (p *Planner[EC, P]) ConfigureSubscription() plan.SubscriptionConfiguration { - if any(p.eventConfig) == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) - return plan.SubscriptionConfiguration{} - } - providerId := p.config.GetProviderId(p.eventConfig) - pubsub, ok := p.providers[providerId] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) - return plan.SubscriptionConfiguration{} - } - - dataSource, err := p.config.GetResolveDataSourceSubscription(p.eventConfig, pubsub) - if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription: %w", err)) - return plan.SubscriptionConfiguration{} - } - - input, err := p.config.GetResolveDataSourceSubscriptionInput(p.eventConfig, pubsub) - if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription input: %w", err)) - return plan.SubscriptionConfiguration{} - } - - return plan.SubscriptionConfiguration{ - Input: input, - Variables: p.variables, - DataSource: dataSource, - PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, - }, - } -} - -func (p *Planner[EC, P]) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { - variablePath, err := p.visitor.Operation.VariablePathByArgumentRefAndArgumentPath(argumentRef, argumentPath, p.visitor.Walker.Ancestors[0].Ref) - if err != nil { - return "", err - } - /* The definition is passed as both definition and operation below because getJSONRootType resolves the type - * from the first argument, but finalInputValueTypeRef comes from the definition - */ - contextVariable := &resolve.ContextVariable{ - Path: variablePath, - Renderer: resolve.NewPlainVariableRenderer(), - } - variablePlaceHolder, _ := p.variables.AddVariable(contextVariable) - return variablePlaceHolder, nil -} - -func StringParser(subject string) (string, error) { - matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(subject, -1) - if len(matches) < 1 { - return subject, nil - } - return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) -} - -func (p *Planner[EC, P]) extractArgumentTemplate(fieldRef int, template string) (string, error) { - matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(template, -1) - // If no argument templates are defined, there are only static values - if len(matches) < 1 { - return template, nil - } - fieldNameBytes := p.visitor.Operation.FieldNameBytes(fieldRef) - // TODO: handling for interfaces and unions - fieldDefinitionRef, ok := p.visitor.Definition.ObjectTypeDefinitionFieldWithName(p.visitor.Walker.EnclosingTypeDefinition.Ref, fieldNameBytes) - if !ok { - return "", fmt.Errorf(`expected field definition to exist for field "%s"`, fieldNameBytes) - } - templateWithVariableTemplateReplacements := template - for templateNumber, groups := range matches { - // The first group is the whole template; the second is the period delimited argument path - if len(groups) != 2 { - return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: expected 2 matching groups but received %d`, templateNumber+1, fieldNameBytes, len(groups)-1) - } - validationResult, err := argument_templates.ValidateArgumentPath(p.visitor.Definition, groups[1], fieldDefinitionRef) - if err != nil { - return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: %w`, templateNumber+1, fieldNameBytes, err) - } - argumentNameBytes := []byte(validationResult.ArgumentPath[0]) - argumentRef, ok := p.visitor.Operation.FieldArgument(fieldRef, argumentNameBytes) - if !ok { - return "", fmt.Errorf(`operation field "%s" does not define argument "%s"`, fieldNameBytes, argumentNameBytes) - } - // variablePlaceholder has the form $$0$$, $$1$$, etc. - variablePlaceholder, err := p.addContextVariableByArgumentRef(argumentRef, validationResult.ArgumentPath) - if err != nil { - return "", fmt.Errorf(`failed to retrieve variable placeholder for argument ""%s" defined on operation field "%s": %w`, argumentNameBytes, fieldNameBytes, err) - } - // Replace the template literal with the variable placeholder (and reuse the variable if it already exists) - templateWithVariableTemplateReplacements = strings.ReplaceAll(templateWithVariableTemplateReplacements, groups[0], variablePlaceholder) - } - - return templateWithVariableTemplateReplacements, nil -} - -func (p *Planner[EC, P]) EnterDocument(_, _ *ast.Document) { - p.rootFieldRef = -1 - var zero EC - p.eventConfig = zero -} - -func (p *Planner[EC, P]) EnterField(ref int) { - if p.rootFieldRef != -1 { - // This is a nested field; nothing needs to be done - return - } - p.rootFieldRef = ref - - fieldName := p.visitor.Operation.FieldNameString(ref) - typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) - - extractFn := func(tpl string) (string, error) { - return p.extractArgumentTemplate(ref, tpl) - } - - eventConfigV, err := p.config.FindEventConfig(p.config.GetEventsDataConfigurations(), typeName, fieldName, extractFn) - if err != nil { - return - } - - p.eventConfig = eventConfigV -} - -type Source struct{} - -func (s *Source) Load(ctx context.Context, input []byte, out *bytes.Buffer) (err error) { - _, err = out.Write(input) - return -} - -func (s *Source) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { - panic("not implemented") -} - -func NewFactory[EC EventConfigType, P any](executionContext context.Context, config config.EventsConfiguration, providers map[string]P) *Factory[EC, P] { - return &Factory[EC, P]{ - providers: providers, - executionContext: executionContext, - config: config, - } -} - -type Factory[EC EventConfigType, P any] struct { - providers map[string]P - executionContext context.Context - config config.EventsConfiguration -} - -func (f *Factory[EC, P]) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Implementer[EC, P]] { - return &Planner[EC, P]{ - providers: f.providers, - } -} - -func (f *Factory[EC, P]) Context() context.Context { - return f.executionContext -} - -func (f *Factory[EC, P]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Implementer[EC, P]]) (*ast.Document, bool) { - return nil, false -} diff --git a/router/pkg/pubsub/datasource/factory.go b/router/pkg/pubsub/datasource/factory.go new file mode 100644 index 0000000000..498ef93510 --- /dev/null +++ b/router/pkg/pubsub/datasource/factory.go @@ -0,0 +1,38 @@ +package datasource + +import ( + "context" + + "github.com/jensneuse/abstractlogger" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" +) + +func NewFactory[EC EventConfigType, P any](executionContext context.Context, config config.EventsConfiguration, providers map[string]P) *Factory[EC, P] { + return &Factory[EC, P]{ + providers: providers, + executionContext: executionContext, + config: config, + } +} + +type Factory[EC EventConfigType, P any] struct { + providers map[string]P + executionContext context.Context + config config.EventsConfiguration +} + +func (f *Factory[EC, P]) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Implementer[EC, P]] { + return &Planner[EC, P]{ + providers: f.providers, + } +} + +func (f *Factory[EC, P]) Context() context.Context { + return f.executionContext +} + +func (f *Factory[EC, P]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Implementer[EC, P]]) (*ast.Document, bool) { + return nil, false +} diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go new file mode 100644 index 0000000000..2929d44813 --- /dev/null +++ b/router/pkg/pubsub/datasource/planner.go @@ -0,0 +1,219 @@ +package datasource + +import ( + "fmt" + "strings" + + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type Planner[EC EventConfigType, P any] struct { + id int + config Implementer[EC, P] + eventsConfig config.EventsConfiguration + eventConfig EC + rootFieldRef int + variables resolve.Variables + visitor *plan.Visitor + providers map[string]P +} + +func (p *Planner[EC, P]) SetID(id int) { + p.id = id +} + +func (p *Planner[EC, P]) ID() (id int) { + return p.id +} + +func (p *Planner[EC, P]) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { + // skip, not required + return +} + +func (p *Planner[EC, P]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { + return plan.DataSourcePlanningBehavior{ + MergeAliasedRootNodes: false, + OverrideFieldPathFromAlias: false, + } +} + +func (p *Planner[EC, P]) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Implementer[EC, P]], _ plan.DataSourcePlannerConfiguration) error { + p.visitor = visitor + visitor.Walker.RegisterEnterFieldVisitor(p) + visitor.Walker.RegisterEnterDocumentVisitor(p) + p.config = configuration.CustomConfiguration() + return nil +} + +func (p *Planner[EC, P]) ConfigureFetch() resolve.FetchConfiguration { + if any(p.eventConfig) == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) + return resolve.FetchConfiguration{} + } + + var dataSource resolve.DataSource + providerId := p.config.GetProviderId(p.eventConfig) + pubsub, ok := p.providers[providerId] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) + return resolve.FetchConfiguration{} + } + + dataSource, err := p.config.GetResolveDataSource(p.eventConfig, pubsub) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get data source: %w", err)) + return resolve.FetchConfiguration{} + } + + event, err := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) + return resolve.FetchConfiguration{} + } + + input, err := p.config.GetResolveDataSourceInput(p.eventConfig, event) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) + return resolve.FetchConfiguration{} + } + + return resolve.FetchConfiguration{ + Input: input, + Variables: p.variables, + DataSource: dataSource, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner[EC, P]) ConfigureSubscription() plan.SubscriptionConfiguration { + if any(p.eventConfig) == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) + return plan.SubscriptionConfiguration{} + } + providerId := p.config.GetProviderId(p.eventConfig) + pubsub, ok := p.providers[providerId] + if !ok { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) + return plan.SubscriptionConfiguration{} + } + + dataSource, err := p.config.GetResolveDataSourceSubscription(p.eventConfig, pubsub) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription: %w", err)) + return plan.SubscriptionConfiguration{} + } + + input, err := p.config.GetResolveDataSourceSubscriptionInput(p.eventConfig, pubsub) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription input: %w", err)) + return plan.SubscriptionConfiguration{} + } + + return plan.SubscriptionConfiguration{ + Input: input, + Variables: p.variables, + DataSource: dataSource, + PostProcessing: resolve.PostProcessingConfiguration{ + MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + }, + } +} + +func (p *Planner[EC, P]) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { + variablePath, err := p.visitor.Operation.VariablePathByArgumentRefAndArgumentPath(argumentRef, argumentPath, p.visitor.Walker.Ancestors[0].Ref) + if err != nil { + return "", err + } + /* The definition is passed as both definition and operation below because getJSONRootType resolves the type + * from the first argument, but finalInputValueTypeRef comes from the definition + */ + contextVariable := &resolve.ContextVariable{ + Path: variablePath, + Renderer: resolve.NewPlainVariableRenderer(), + } + variablePlaceHolder, _ := p.variables.AddVariable(contextVariable) + return variablePlaceHolder, nil +} + +func StringParser(subject string) (string, error) { + matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(subject, -1) + if len(matches) < 1 { + return subject, nil + } + return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) +} + +func (p *Planner[EC, P]) extractArgumentTemplate(fieldRef int, template string) (string, error) { + matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(template, -1) + // If no argument templates are defined, there are only static values + if len(matches) < 1 { + return template, nil + } + fieldNameBytes := p.visitor.Operation.FieldNameBytes(fieldRef) + // TODO: handling for interfaces and unions + fieldDefinitionRef, ok := p.visitor.Definition.ObjectTypeDefinitionFieldWithName(p.visitor.Walker.EnclosingTypeDefinition.Ref, fieldNameBytes) + if !ok { + return "", fmt.Errorf(`expected field definition to exist for field "%s"`, fieldNameBytes) + } + templateWithVariableTemplateReplacements := template + for templateNumber, groups := range matches { + // The first group is the whole template; the second is the period delimited argument path + if len(groups) != 2 { + return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: expected 2 matching groups but received %d`, templateNumber+1, fieldNameBytes, len(groups)-1) + } + validationResult, err := argument_templates.ValidateArgumentPath(p.visitor.Definition, groups[1], fieldDefinitionRef) + if err != nil { + return "", fmt.Errorf(`argument template #%d defined on field "%s" is invalid: %w`, templateNumber+1, fieldNameBytes, err) + } + argumentNameBytes := []byte(validationResult.ArgumentPath[0]) + argumentRef, ok := p.visitor.Operation.FieldArgument(fieldRef, argumentNameBytes) + if !ok { + return "", fmt.Errorf(`operation field "%s" does not define argument "%s"`, fieldNameBytes, argumentNameBytes) + } + // variablePlaceholder has the form $$0$$, $$1$$, etc. + variablePlaceholder, err := p.addContextVariableByArgumentRef(argumentRef, validationResult.ArgumentPath) + if err != nil { + return "", fmt.Errorf(`failed to retrieve variable placeholder for argument ""%s" defined on operation field "%s": %w`, argumentNameBytes, fieldNameBytes, err) + } + // Replace the template literal with the variable placeholder (and reuse the variable if it already exists) + templateWithVariableTemplateReplacements = strings.ReplaceAll(templateWithVariableTemplateReplacements, groups[0], variablePlaceholder) + } + + return templateWithVariableTemplateReplacements, nil +} + +func (p *Planner[EC, P]) EnterDocument(_, _ *ast.Document) { + p.rootFieldRef = -1 + var zero EC + p.eventConfig = zero +} + +func (p *Planner[EC, P]) EnterField(ref int) { + if p.rootFieldRef != -1 { + // This is a nested field; nothing needs to be done + return + } + p.rootFieldRef = ref + + fieldName := p.visitor.Operation.FieldNameString(ref) + typeName := p.visitor.Walker.EnclosingTypeDefinition.NameString(p.visitor.Definition) + + extractFn := func(tpl string) (string, error) { + return p.extractArgumentTemplate(ref, tpl) + } + + eventConfigV, err := p.config.FindEventConfig(p.config.GetEventsDataConfigurations(), typeName, fieldName, extractFn) + if err != nil { + return + } + + p.eventConfig = eventConfigV +} diff --git a/router/pkg/pubsub/kafka/datasource.go b/router/pkg/pubsub/kafka/datasource.go index e242828d2a..2f789b12bb 100644 --- a/router/pkg/pubsub/kafka/datasource.go +++ b/router/pkg/pubsub/kafka/datasource.go @@ -3,7 +3,6 @@ package kafka import ( "context" "crypto/tls" - "encoding/json" "fmt" "time" @@ -15,7 +14,6 @@ import ( "go.uber.org/zap" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { @@ -26,8 +24,8 @@ func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, return nil, err } factory := k.GetFactory(ctx, config) - ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.KafkaEventConfiguration, *kafkaPubSub]]( - in.Id, + ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.KafkaEventConfiguration, *KafkaPubSub]]( + in.Id+"-kafka", factory, dsMeta, &Configuration{ @@ -50,9 +48,11 @@ func init() { datasource.RegisterPubSub(GetPlanDataSource) } +// var _ datasource.PubSubImplementer[*KafkaPubSub] = &Kafka{} + type Kafka struct { logger *zap.Logger - providers map[string]*kafkaPubSub + providers map[string]*KafkaPubSub } // buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. @@ -112,91 +112,23 @@ func (k *Kafka) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfi return nil } -func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *datasource.Factory[*nodev1.KafkaEventConfiguration, *kafkaPubSub] { +// func (k *Kafka) ConnectProviders(ctx context.Context) error { +// for _, provider := range k.providers { +// err := provider.Connect() +// if err != nil { +// return fmt.Errorf("failed to connect to Kafka provider with ID \"%s\": %w", provider.ID, err) +// } +// } +// return nil +// } + +func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *datasource.Factory[*nodev1.KafkaEventConfiguration, *KafkaPubSub] { return datasource.NewFactory[*nodev1.KafkaEventConfiguration](executionContext, config, k.providers) } func NewPubSub(logger *zap.Logger) Kafka { return Kafka{ - providers: map[string]*kafkaPubSub{}, + providers: map[string]*KafkaPubSub{}, logger: logger, } } - -type Configuration struct { - Data string `json:"data"` - EventConfiguration []*nodev1.KafkaEventConfiguration - Logger *zap.Logger -} - -func (c *Configuration) GetEventsDataConfigurations() []*nodev1.KafkaEventConfiguration { - return c.EventConfiguration -} - -func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.KafkaEventConfiguration, pubsub *kafkaPubSub) (resolve.DataSource, error) { - var dataSource resolve.DataSource - - typeName := eventConfig.GetEngineEventConfiguration().GetType() - switch typeName { - case nodev1.EventType_PUBLISH: - dataSource = &KafkaPublishDataSource{ - pubSub: pubsub, - } - default: - return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String()) - } - - return dataSource, nil -} - -func (c *Configuration) GetResolveDataSourceSubscription(eventConfig *nodev1.KafkaEventConfiguration, pubsub *kafkaPubSub) (resolve.SubscriptionDataSource, error) { - return &SubscriptionSource{ - pubSub: pubsub, - }, nil -} - -func (c *Configuration) GetResolveDataSourceSubscriptionInput(eventConfig *nodev1.KafkaEventConfiguration, pubsub *kafkaPubSub) (string, error) { - providerId := c.GetProviderId(eventConfig) - evtCfg := SubscriptionEventConfiguration{ - ProviderID: providerId, - Topics: eventConfig.GetTopics(), - } - object, err := json.Marshal(evtCfg) - if err != nil { - return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") - } - return string(object), nil -} - -func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.KafkaEventConfiguration, event []byte) (string, error) { - topics := eventConfig.GetTopics() - - if len(topics) != 1 { - return "", fmt.Errorf("publish and request events should define one topic but received %d", len(topics)) - } - - topic := topics[0] - - providerId := c.GetProviderId(eventConfig) - - evtCfg := PublishEventConfiguration{ - ProviderID: providerId, - Topic: topic, - Data: event, - } - - return evtCfg.MarshalJSONTemplate(), nil -} - -func (c *Configuration) GetProviderId(eventConfig *nodev1.KafkaEventConfiguration) string { - return eventConfig.GetEngineEventConfiguration().GetProviderId() -} - -func (c *Configuration) FindEventConfig(eventConfigs []*nodev1.KafkaEventConfiguration, typeName string, fieldName string, fn datasource.ArgumentTemplateCallback) (*nodev1.KafkaEventConfiguration, error) { - for _, cfg := range eventConfigs { - if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - return cfg, nil - } - } - return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) -} diff --git a/router/pkg/pubsub/kafka/datasource_impl.go b/router/pkg/pubsub/kafka/datasource_impl.go new file mode 100644 index 0000000000..2828987f97 --- /dev/null +++ b/router/pkg/pubsub/kafka/datasource_impl.go @@ -0,0 +1,89 @@ +package kafka + +import ( + "encoding/json" + "fmt" + + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" + "go.uber.org/zap" +) + +type Configuration struct { + Data string `json:"data"` + EventConfiguration []*nodev1.KafkaEventConfiguration + Logger *zap.Logger +} + +func (c *Configuration) GetEventsDataConfigurations() []*nodev1.KafkaEventConfiguration { + return c.EventConfiguration +} + +func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.KafkaEventConfiguration, pubsub *KafkaPubSub) (resolve.DataSource, error) { + var dataSource resolve.DataSource + + typeName := eventConfig.GetEngineEventConfiguration().GetType() + switch typeName { + case nodev1.EventType_PUBLISH: + dataSource = &KafkaPublishDataSource{ + pubSub: pubsub, + } + default: + return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String()) + } + + return dataSource, nil +} + +func (c *Configuration) GetResolveDataSourceSubscription(eventConfig *nodev1.KafkaEventConfiguration, pubsub *KafkaPubSub) (resolve.SubscriptionDataSource, error) { + return &SubscriptionSource{ + pubSub: pubsub, + }, nil +} + +func (c *Configuration) GetResolveDataSourceSubscriptionInput(eventConfig *nodev1.KafkaEventConfiguration, pubsub *KafkaPubSub) (string, error) { + providerId := c.GetProviderId(eventConfig) + evtCfg := SubscriptionEventConfiguration{ + ProviderID: providerId, + Topics: eventConfig.GetTopics(), + } + object, err := json.Marshal(evtCfg) + if err != nil { + return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") + } + return string(object), nil +} + +func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.KafkaEventConfiguration, event []byte) (string, error) { + topics := eventConfig.GetTopics() + + if len(topics) != 1 { + return "", fmt.Errorf("publish and request events should define one topic but received %d", len(topics)) + } + + topic := topics[0] + + providerId := c.GetProviderId(eventConfig) + + evtCfg := PublishEventConfiguration{ + ProviderID: providerId, + Topic: topic, + Data: event, + } + + return evtCfg.MarshalJSONTemplate(), nil +} + +func (c *Configuration) GetProviderId(eventConfig *nodev1.KafkaEventConfiguration) string { + return eventConfig.GetEngineEventConfiguration().GetProviderId() +} + +func (c *Configuration) FindEventConfig(eventConfigs []*nodev1.KafkaEventConfiguration, typeName string, fieldName string, fn datasource.ArgumentTemplateCallback) (*nodev1.KafkaEventConfiguration, error) { + for _, cfg := range eventConfigs { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + return cfg, nil + } + } + return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) +} diff --git a/router/pkg/pubsub/kafka/kafka.go b/router/pkg/pubsub/kafka/kafka.go index 6a9260342a..a897ac245b 100644 --- a/router/pkg/pubsub/kafka/kafka.go +++ b/router/pkg/pubsub/kafka/kafka.go @@ -20,19 +20,46 @@ var ( ) type connector struct { - writeClient *kgo.Client + writeClient *LazyClient opts []kgo.Opt logger *zap.Logger } +type LazyClient struct { + once sync.Once + client *kgo.Client + opts []kgo.Opt +} + +func (c *LazyClient) Connect() (err error) { + c.once.Do(func() { + c.client, err = kgo.NewClient(append(c.opts, + // For observability, we set the client ID to "router" + kgo.ClientID("cosmo.router.producer"))..., + ) + }) + + return +} + +func (c *LazyClient) GetClient() *kgo.Client { + if c.client == nil { + c.Connect() + } + return c.client +} + +func NewLazyClient(opts ...kgo.Opt) *LazyClient { + return &LazyClient{ + opts: opts, + } +} + func NewConnector(logger *zap.Logger, opts []kgo.Opt) (*connector, error) { - writeClient, err := kgo.NewClient(append(opts, + writeClient := NewLazyClient(append(opts, // For observability, we set the client ID to "router" kgo.ClientID("cosmo.router.producer"))..., ) - if err != nil { - return nil, fmt.Errorf("failed to create write client for Kafka: %w", err) - } return &connector{ writeClient: writeClient, @@ -41,11 +68,15 @@ func NewConnector(logger *zap.Logger, opts []kgo.Opt) (*connector, error) { }, nil } -func (c *connector) New(ctx context.Context) *kafkaPubSub { +func (c *connector) New(ctx context.Context) *KafkaPubSub { ctx, cancel := context.WithCancel(ctx) - ps := &kafkaPubSub{ + if c.logger == nil { + c.logger = zap.NewNop() + } + + ps := &KafkaPubSub{ ctx: ctx, logger: c.logger.With(zap.String("pubsub", "kafka")), opts: c.opts, @@ -57,22 +88,22 @@ func (c *connector) New(ctx context.Context) *kafkaPubSub { return ps } -// kafkaPubSub is a Kafka pubsub implementation. +// KafkaPubSub is a Kafka pubsub implementation. // It uses the franz-go Kafka client to consume and produce messages. // The pubsub is stateless and does not store any messages. // It uses a single write client to produce messages and a client per topic to consume messages. // Each client polls the Kafka topic for new records and updates the subscriptions with the new data. -type kafkaPubSub struct { +type KafkaPubSub struct { ctx context.Context opts []kgo.Opt logger *zap.Logger - writeClient *kgo.Client + writeClient *LazyClient closeWg sync.WaitGroup cancel context.CancelFunc } // topicPoller polls the Kafka topic for new records and calls the updateTriggers function. -func (p *kafkaPubSub) topicPoller(ctx context.Context, client *kgo.Client, updater resolve.SubscriptionUpdater) error { +func (p *KafkaPubSub) topicPoller(ctx context.Context, client *kgo.Client, updater resolve.SubscriptionUpdater) error { for { select { case <-p.ctx.Done(): // Close the poller if the application context was canceled @@ -126,7 +157,7 @@ func (p *kafkaPubSub) topicPoller(ctx context.Context, client *kgo.Client, updat // Subscribe subscribes to the given topics and updates the subscription updater. // The engine already deduplicates subscriptions with the same topics, stream configuration, extensions, headers, etc. -func (p *kafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { +func (p *KafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), @@ -143,6 +174,9 @@ func (p *kafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConf kgo.ConsumeResetOffset(kgo.NewOffset().AfterMilli(time.Now().UnixMilli())), // For observability, we set the client ID to "router" kgo.ClientID(fmt.Sprintf("cosmo.router.consumer.%s", strings.Join(event.Topics, "-"))), + // FIXME: the client id should have some unique identifier, like in nats + // What if we have multiple subscriptions for the same topics? + // What if we have more router instances? )...) if err != nil { log.Error("failed to create client", zap.Error(err)) @@ -173,7 +207,7 @@ func (p *kafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConf // Publish publishes the given event to the Kafka topic in a non-blocking way. // Publish errors are logged and returned as a pubsub error. // The event is written with a dedicated write client. -func (p *kafkaPubSub) Publish(ctx context.Context, event PublishEventConfiguration) error { +func (p *KafkaPubSub) Publish(ctx context.Context, event PublishEventConfiguration) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "publish"), @@ -187,7 +221,7 @@ func (p *kafkaPubSub) Publish(ctx context.Context, event PublishEventConfigurati var pErr error - p.writeClient.Produce(ctx, &kgo.Record{ + p.writeClient.GetClient().Produce(ctx, &kgo.Record{ Topic: event.Topic, Value: event.Data, }, func(record *kgo.Record, err error) { @@ -207,14 +241,14 @@ func (p *kafkaPubSub) Publish(ctx context.Context, event PublishEventConfigurati return nil } -func (p *kafkaPubSub) Shutdown(ctx context.Context) error { +func (p *KafkaPubSub) Shutdown(ctx context.Context) error { - err := p.writeClient.Flush(ctx) + err := p.writeClient.GetClient().Flush(ctx) if err != nil { p.logger.Error("error flushing write client", zap.Error(err)) } - p.writeClient.Close() + p.writeClient.GetClient().Close() // Cancel the context to stop all pollers p.cancel() diff --git a/router/pkg/pubsub/kafka/subscription_datasource.go b/router/pkg/pubsub/kafka/subscription_datasource.go index 70623170cf..0cdc230e87 100644 --- a/router/pkg/pubsub/kafka/subscription_datasource.go +++ b/router/pkg/pubsub/kafka/subscription_datasource.go @@ -29,7 +29,7 @@ func (s *PublishEventConfiguration) MarshalJSONTemplate() string { } type SubscriptionSource struct { - pubSub *kafkaPubSub + pubSub *KafkaPubSub } func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { @@ -63,7 +63,7 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r } type KafkaPublishDataSource struct { - pubSub *kafkaPubSub + pubSub *KafkaPubSub } func (s *KafkaPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go index f970e3be9b..b8de491ae7 100644 --- a/router/pkg/pubsub/nats/datasource.go +++ b/router/pkg/pubsub/nats/datasource.go @@ -8,6 +8,7 @@ import ( "regexp" "slices" "strings" + "sync" "time" "github.com/nats-io/nats.go" @@ -40,7 +41,7 @@ func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMe } factory := k.GetFactory(ctx, config, k.providers) ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.NatsEventConfiguration, *NatsPubSub]]( - in.Id, + in.Id+"-nats", factory, dsMeta, &Configuration{ @@ -114,6 +115,46 @@ type Nats struct { routerListenAddr string // How to get it here? } +type LazyClient struct { + once sync.Once + url string + opts []nats.Option + client *nats.Conn + js jetstream.JetStream +} + +func (c *LazyClient) Connect(opts ...nats.Option) (err error) { + c.once.Do(func() { + c.client, err = nats.Connect(c.url, opts...) + if err != nil { + return + } + c.js, err = jetstream.New(c.client) + }) + return +} + +func (c *LazyClient) GetClient() *nats.Conn { + if c.client == nil { + c.Connect() + } + return c.client +} + +func (c *LazyClient) GetJetStream() jetstream.JetStream { + if c.js == nil { + c.Connect() + } + return c.js +} + +func NewLazyClient(url string, opts ...nats.Option) *LazyClient { + return &LazyClient{ + url: url, + opts: opts, + } +} + func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { definedProviders := make(map[string]bool) for _, provider := range config.Providers.Nats { @@ -135,16 +176,8 @@ func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfig if err != nil { return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", provider.ID, err) } - natsConnection, err := nats.Connect(provider.URL, options...) - if err != nil { - return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", provider.ID, err) - } - js, err := jetstream.New(natsConnection) - if err != nil { - return err - } - n.providers[provider.ID] = NewConnector(n.logger, natsConnection, js, n.hostName, n.routerListenAddr).New(ctx) + n.providers[provider.ID] = NewConnector(n.logger, provider.URL, options, n.hostName, n.routerListenAddr).New(ctx) } return nil } @@ -177,10 +210,14 @@ func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.NatsEventConfig typeName := eventConfig.GetEngineEventConfiguration().GetType() switch typeName { - case nodev1.EventType_PUBLISH, nodev1.EventType_REQUEST: + case nodev1.EventType_PUBLISH: dataSource = &NatsPublishDataSource{ pubSub: pubsub, } + case nodev1.EventType_REQUEST: + dataSource = &NatsRequestDataSource{ + pubSub: pubsub, + } default: return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Nats", typeName.String()) } diff --git a/router/pkg/pubsub/nats/nats.go b/router/pkg/pubsub/nats/nats.go index 3528803379..e57df681ad 100644 --- a/router/pkg/pubsub/nats/nats.go +++ b/router/pkg/pubsub/nats/nats.go @@ -17,28 +17,29 @@ import ( ) type connector struct { - conn *nats.Conn + client *LazyClient logger *zap.Logger - js jetstream.JetStream hostName string routerListenAddr string } -func NewConnector(logger *zap.Logger, conn *nats.Conn, js jetstream.JetStream, hostName string, routerListenAddr string) *connector { +func NewConnector(logger *zap.Logger, url string, opts []nats.Option, hostName string, routerListenAddr string) *connector { + client := NewLazyClient(url, opts...) return &connector{ - conn: conn, + client: client, logger: logger, - js: js, hostName: hostName, routerListenAddr: routerListenAddr, } } func (c *connector) New(ctx context.Context) *NatsPubSub { + if c.logger == nil { + c.logger = zap.NewNop() + } return &NatsPubSub{ ctx: ctx, - conn: c.conn, - js: c.js, + client: c.client, logger: c.logger.With(zap.String("pubsub", "nats")), closeWg: sync.WaitGroup{}, hostName: c.hostName, @@ -48,9 +49,8 @@ func (c *connector) New(ctx context.Context) *NatsPubSub { type NatsPubSub struct { ctx context.Context - conn *nats.Conn + client *LazyClient logger *zap.Logger - js jetstream.JetStream closeWg sync.WaitGroup hostName string routerListenAddr string @@ -104,7 +104,7 @@ func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfi if event.StreamConfiguration.ConsumerInactiveThreshold > 0 { consumerConfig.InactiveThreshold = time.Duration(event.StreamConfiguration.ConsumerInactiveThreshold) * time.Second } - consumer, err := p.js.CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) + consumer, err := p.client.GetJetStream().CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) if err != nil { log.Error("error creating or updating consumer", zap.Error(err)) return datasource.NewError(fmt.Sprintf(`failed to create or update consumer for stream "%s"`, event.StreamConfiguration.StreamName), err) @@ -155,7 +155,7 @@ func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfi msgChan := make(chan *nats.Msg) subscriptions := make([]*nats.Subscription, len(event.Subjects)) for i, subject := range event.Subjects { - subscription, err := p.conn.ChanSubscribe(subject, msgChan) + subscription, err := p.client.GetClient().ChanSubscribe(subject, msgChan) if err != nil { log.Error("error subscribing to NATS subject", zap.Error(err), zap.String("subscription_subject", subject)) return datasource.NewError(fmt.Sprintf(`failed to subscribe to NATS subject "%s"`, subject), err) @@ -209,7 +209,7 @@ func (p *NatsPubSub) Publish(_ context.Context, event PublishAndRequestEventConf log.Debug("publish", zap.ByteString("data", event.Data)) - err := p.conn.Publish(event.Subject, event.Data) + err := p.client.GetClient().Publish(event.Subject, event.Data) if err != nil { log.Error("publish error", zap.Error(err)) return datasource.NewError(fmt.Sprintf("error publishing to NATS subject %s", event.Subject), err) @@ -227,7 +227,7 @@ func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventCo log.Debug("request", zap.ByteString("data", event.Data)) - msg, err := p.conn.RequestWithContext(ctx, event.Subject, event.Data) + msg, err := p.client.GetClient().RequestWithContext(ctx, event.Subject, event.Data) if err != nil { log.Error("request error", zap.Error(err)) return datasource.NewError(fmt.Sprintf("error requesting from NATS subject %s", event.Subject), err) @@ -243,12 +243,12 @@ func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventCo } func (p *NatsPubSub) flush(ctx context.Context) error { - return p.conn.FlushWithContext(ctx) + return p.client.GetClient().FlushWithContext(ctx) } func (p *NatsPubSub) Shutdown(ctx context.Context) error { - if p.conn.IsClosed() { + if p.client.GetClient().IsClosed() { return nil } @@ -260,7 +260,7 @@ func (p *NatsPubSub) Shutdown(ctx context.Context) error { err = errors.Join(err, fErr) } - drainErr := p.conn.Drain() + drainErr := p.client.GetClient().Drain() if drainErr != nil { p.logger.Error("error draining NATS connection", zap.Error(drainErr)) } From 33603e88b3c57dfbad778699d892d345e6f7dee7 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 24 Mar 2025 19:05:46 +0100 Subject: [PATCH 08/49] refactor(pubsub): remove unused variables and comments in datasource and planner files, enhance NATS subject validation --- router/pkg/pubsub/datasource/datasource.go | 2 -- router/pkg/pubsub/datasource/planner.go | 2 -- router/pkg/pubsub/nats/datasource.go | 15 ++++++--------- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index f76994b1bd..cc355dfeae 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -26,7 +26,6 @@ type ArgumentTemplateCallback func(tpl string) (string, error) type PubSubImplementer[F any] interface { PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error - // ConnectProviders(ctx context.Context) error GetFactory(executionContext context.Context, config config.EventsConfiguration) F } @@ -39,7 +38,6 @@ type Implementer[EC EventConfigType, P any] interface { GetResolveDataSource(eventConfig EC, pubsub P) (resolve.DataSource, error) GetResolveDataSourceSubscription(eventConfig EC, pubsub P) (resolve.SubscriptionDataSource, error) GetResolveDataSourceSubscriptionInput(eventConfig EC, pubsub P) (string, error) - //GetResolveDataSourceInput(eventConfig EC, ref int, visitor *plan.Visitor, variables *resolve.Variables) (string, error) GetResolveDataSourceInput(eventConfig EC, event []byte) (string, error) GetProviderId(eventConfig EC) string FindEventConfig(eventConfigs []EC, typeName string, fieldName string, fn ArgumentTemplateCallback) (EC, error) diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index 2929d44813..756ba996ce 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -4,7 +4,6 @@ import ( "fmt" "strings" - "github.com/wundergraph/cosmo/router/pkg/config" "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" @@ -15,7 +14,6 @@ import ( type Planner[EC EventConfigType, P any] struct { id int config Implementer[EC, P] - eventsConfig config.EventsConfiguration eventConfig EC rootFieldRef int variables resolve.Variables diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go index b8de491ae7..701271ae67 100644 --- a/router/pkg/pubsub/nats/datasource.go +++ b/router/pkg/pubsub/nats/datasource.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "regexp" "slices" "strings" "sync" @@ -27,11 +26,6 @@ const ( tsep = "." ) -// A variable template has form $$number$$ where the number can range from one to multiple digits -var ( - variableTemplateRegex = regexp.MustCompile(`\$\$\d+\$\$`) -) - func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { if natsData := in.GetCustomEvents().GetNats(); natsData != nil { k := NewPubSub(logger) @@ -196,9 +190,6 @@ type Configuration struct { Data string `json:"data"` EventConfiguration []*nodev1.NatsEventConfiguration Logger *zap.Logger - requestConfig *PublishAndRequestEventConfiguration - publishConfig *PublishAndRequestEventConfiguration - subscribeConfig *SubscriptionEventConfiguration } func (c *Configuration) GetEventsDataConfigurations() []*nodev1.NatsEventConfiguration { @@ -283,6 +274,9 @@ func (c *Configuration) transformEventConfig(cfg *nodev1.NatsEventConfiguration, if err != nil { return cfg, fmt.Errorf("unable to parse subject with id %s", cfg.GetSubjects()[0]) } + if !isValidNatsSubject(extractedSubject) { + return cfg, fmt.Errorf("invalid subject: %s", extractedSubject) + } cfg.Subjects = []string{extractedSubject} case nodev1.EventType_SUBSCRIBE: extractedSubjects := make([]string, 0, len(cfg.Subjects)) @@ -291,6 +285,9 @@ func (c *Configuration) transformEventConfig(cfg *nodev1.NatsEventConfiguration, if err != nil { return cfg, nil } + if !isValidNatsSubject(extractedSubject) { + return cfg, fmt.Errorf("invalid subject: %s", extractedSubject) + } extractedSubjects = append(extractedSubjects, extractedSubject) } slices.Sort(extractedSubjects) From a473fddc521412f035f3a084e49db99addd07a50 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 7 Apr 2025 23:12:27 +0200 Subject: [PATCH 09/49] make it works with different providers in the same datasource --- router/pkg/pubsub/datasource.go | 23 +++-- router/pkg/pubsub/datasource/datasource.go | 4 +- router/pkg/pubsub/datasource/factory.go | 18 ++-- router/pkg/pubsub/datasource/planner.go | 93 ++++++++++-------- router/pkg/pubsub/kafka/datasource.go | 35 +++---- router/pkg/pubsub/kafka/datasource_impl.go | 43 +++++++- router/pkg/pubsub/nats/datasource.go | 109 ++++++++++++++------- router/pkg/pubsub/nats/nats.go | 56 ++++++++--- 8 files changed, 252 insertions(+), 129 deletions(-) diff --git a/router/pkg/pubsub/datasource.go b/router/pkg/pubsub/datasource.go index c5bf21a575..0d8b7131f7 100644 --- a/router/pkg/pubsub/datasource.go +++ b/router/pkg/pubsub/datasource.go @@ -16,20 +16,31 @@ import ( ) func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) ([]plan.DataSource, error) { - var dataSources []plan.DataSource + var pubSubs datasource.PubSubGeneralImplementerList + for _, pubSub := range datasource.GetRegisteredPubSubs() { - ds, err := pubSub(ctx, in, dsMeta, config, logger) + pubSub, err := pubSub(ctx, in, dsMeta, config, logger) if err != nil { return nil, err } - if ds != nil { - dataSources = append(dataSources, ds) + if pubSub != nil { + pubSubs = append(pubSubs, pubSub) } } - if len(dataSources) == 0 { + if len(pubSubs) == 0 { return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) } - return dataSources, nil + ds, err := plan.NewDataSourceConfiguration( + in.Id, + datasource.NewFactory(ctx, config, pubSubs), + dsMeta, + pubSubs, + ) + if err != nil { + return nil, err + } + + return []plan.DataSource{ds}, nil } diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index cc355dfeae..ed3ab2b664 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -10,7 +10,7 @@ import ( "go.uber.org/zap" ) -type Getter func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) +type Getter func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (PubSubGeneralImplementer, error) var pubSubs []Getter @@ -26,7 +26,7 @@ type ArgumentTemplateCallback func(tpl string) (string, error) type PubSubImplementer[F any] interface { PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error - GetFactory(executionContext context.Context, config config.EventsConfiguration) F + //GetFactory(executionContext context.Context, config config.EventsConfiguration) F } type EventConfigType interface { diff --git a/router/pkg/pubsub/datasource/factory.go b/router/pkg/pubsub/datasource/factory.go index 498ef93510..25a1564ae2 100644 --- a/router/pkg/pubsub/datasource/factory.go +++ b/router/pkg/pubsub/datasource/factory.go @@ -9,30 +9,30 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" ) -func NewFactory[EC EventConfigType, P any](executionContext context.Context, config config.EventsConfiguration, providers map[string]P) *Factory[EC, P] { - return &Factory[EC, P]{ +func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers PubSubGeneralImplementerList) *Factory { + return &Factory{ providers: providers, executionContext: executionContext, config: config, } } -type Factory[EC EventConfigType, P any] struct { - providers map[string]P +type Factory struct { + providers PubSubGeneralImplementerList executionContext context.Context config config.EventsConfiguration } -func (f *Factory[EC, P]) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[Implementer[EC, P]] { - return &Planner[EC, P]{ - providers: f.providers, +func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[PubSubGeneralImplementerList] { + return &Planner{ + pubSubs: f.providers, } } -func (f *Factory[EC, P]) Context() context.Context { +func (f *Factory) Context() context.Context { return f.executionContext } -func (f *Factory[EC, P]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[Implementer[EC, P]]) (*ast.Document, bool) { +func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[PubSubGeneralImplementerList]) (*ast.Document, bool) { return nil, false } diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index 756ba996ce..8d22875b53 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -4,6 +4,7 @@ import ( "fmt" "strings" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" @@ -11,59 +12,79 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type Planner[EC EventConfigType, P any] struct { +type EventConfigType2 interface { + GetResolveDataSource() (resolve.DataSource, error) + GetResolveDataSourceInput(event []byte) (string, error) + GetEngineEventConfiguration() *nodev1.EngineEventConfiguration + GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) + GetResolveDataSourceSubscriptionInput() (string, error) +} + +type PubSubGeneralImplementerList []PubSubGeneralImplementer + +func (p *PubSubGeneralImplementerList) FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType2, error) { + for _, pubSub := range *p { + eventConfigV, err := pubSub.FindEventConfig2(typeName, fieldName, extractFn) + if err != nil { + return nil, err + } + if eventConfigV != nil { + return eventConfigV, nil + } + } + return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) +} + +type PubSubGeneralImplementer interface { + FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType2, error) +} + +type Planner struct { id int - config Implementer[EC, P] - eventConfig EC + pubSubs PubSubGeneralImplementerList + eventConfig EventConfigType2 rootFieldRef int variables resolve.Variables visitor *plan.Visitor - providers map[string]P } -func (p *Planner[EC, P]) SetID(id int) { +func (p *Planner) SetID(id int) { p.id = id } -func (p *Planner[EC, P]) ID() (id int) { +func (p *Planner) ID() (id int) { return p.id } -func (p *Planner[EC, P]) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { +func (p *Planner) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias string, exists bool) { // skip, not required return } -func (p *Planner[EC, P]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { +func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { return plan.DataSourcePlanningBehavior{ MergeAliasedRootNodes: false, OverrideFieldPathFromAlias: false, } } -func (p *Planner[EC, P]) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[Implementer[EC, P]], _ plan.DataSourcePlannerConfiguration) error { +func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[PubSubGeneralImplementerList], _ plan.DataSourcePlannerConfiguration) error { p.visitor = visitor visitor.Walker.RegisterEnterFieldVisitor(p) visitor.Walker.RegisterEnterDocumentVisitor(p) - p.config = configuration.CustomConfiguration() + p.pubSubs = configuration.CustomConfiguration() return nil } -func (p *Planner[EC, P]) ConfigureFetch() resolve.FetchConfiguration { - if any(p.eventConfig) == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) +func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { + if p.eventConfig == nil { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) return resolve.FetchConfiguration{} } var dataSource resolve.DataSource - providerId := p.config.GetProviderId(p.eventConfig) - pubsub, ok := p.providers[providerId] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) - return resolve.FetchConfiguration{} - } - dataSource, err := p.config.GetResolveDataSource(p.eventConfig, pubsub) + dataSource, err := p.eventConfig.GetResolveDataSource() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get data source: %w", err)) return resolve.FetchConfiguration{} @@ -75,7 +96,7 @@ func (p *Planner[EC, P]) ConfigureFetch() resolve.FetchConfiguration { return resolve.FetchConfiguration{} } - input, err := p.config.GetResolveDataSourceInput(p.eventConfig, event) + input, err := p.eventConfig.GetResolveDataSourceInput(event) if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) return resolve.FetchConfiguration{} @@ -91,25 +112,19 @@ func (p *Planner[EC, P]) ConfigureFetch() resolve.FetchConfiguration { } } -func (p *Planner[EC, P]) ConfigureSubscription() plan.SubscriptionConfiguration { - if any(p.eventConfig) == nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) - return plan.SubscriptionConfiguration{} - } - providerId := p.config.GetProviderId(p.eventConfig) - pubsub, ok := p.providers[providerId] - if !ok { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("no pubsub connection exists with provider id \"%s\"", providerId)) +func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { + if p.eventConfig == nil { + // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) return plan.SubscriptionConfiguration{} } - dataSource, err := p.config.GetResolveDataSourceSubscription(p.eventConfig, pubsub) + dataSource, err := p.eventConfig.GetResolveDataSourceSubscription() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription: %w", err)) return plan.SubscriptionConfiguration{} } - input, err := p.config.GetResolveDataSourceSubscriptionInput(p.eventConfig, pubsub) + input, err := p.eventConfig.GetResolveDataSourceSubscriptionInput() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription input: %w", err)) return plan.SubscriptionConfiguration{} @@ -125,7 +140,7 @@ func (p *Planner[EC, P]) ConfigureSubscription() plan.SubscriptionConfiguration } } -func (p *Planner[EC, P]) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { +func (p *Planner) addContextVariableByArgumentRef(argumentRef int, argumentPath []string) (string, error) { variablePath, err := p.visitor.Operation.VariablePathByArgumentRefAndArgumentPath(argumentRef, argumentPath, p.visitor.Walker.Ancestors[0].Ref) if err != nil { return "", err @@ -149,7 +164,7 @@ func StringParser(subject string) (string, error) { return "", fmt.Errorf(`subject "%s" is not a valid NATS subject`, subject) } -func (p *Planner[EC, P]) extractArgumentTemplate(fieldRef int, template string) (string, error) { +func (p *Planner) extractArgumentTemplate(fieldRef int, template string) (string, error) { matches := argument_templates.ArgumentTemplateRegex.FindAllStringSubmatch(template, -1) // If no argument templates are defined, there are only static values if len(matches) < 1 { @@ -188,13 +203,11 @@ func (p *Planner[EC, P]) extractArgumentTemplate(fieldRef int, template string) return templateWithVariableTemplateReplacements, nil } -func (p *Planner[EC, P]) EnterDocument(_, _ *ast.Document) { +func (p *Planner) EnterDocument(_, _ *ast.Document) { p.rootFieldRef = -1 - var zero EC - p.eventConfig = zero } -func (p *Planner[EC, P]) EnterField(ref int) { +func (p *Planner) EnterField(ref int) { if p.rootFieldRef != -1 { // This is a nested field; nothing needs to be done return @@ -208,10 +221,10 @@ func (p *Planner[EC, P]) EnterField(ref int) { return p.extractArgumentTemplate(ref, tpl) } - eventConfigV, err := p.config.FindEventConfig(p.config.GetEventsDataConfigurations(), typeName, fieldName, extractFn) + eventConfigV, err := p.pubSubs.FindEventConfig2(typeName, fieldName, extractFn) if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\": %w", typeName, fieldName, err)) return } - p.eventConfig = eventConfigV } diff --git a/router/pkg/pubsub/kafka/datasource.go b/router/pkg/pubsub/kafka/datasource.go index 2f789b12bb..4a2087dabc 100644 --- a/router/pkg/pubsub/kafka/datasource.go +++ b/router/pkg/pubsub/kafka/datasource.go @@ -16,29 +16,14 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" ) -func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { +func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (datasource.PubSubGeneralImplementer, error) { if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { k := NewPubSub(logger) err := k.PrepareProviders(ctx, in, dsMeta, config) if err != nil { return nil, err } - factory := k.GetFactory(ctx, config) - ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.KafkaEventConfiguration, *KafkaPubSub]]( - in.Id+"-kafka", - factory, - dsMeta, - &Configuration{ - EventConfiguration: kafkaData, - Logger: logger, - }, - ) - - if err != nil { - return nil, err - } - - return ds, nil + return k.config, nil } return nil, nil @@ -53,6 +38,7 @@ func init() { type Kafka struct { logger *zap.Logger providers map[string]*KafkaPubSub + config *Configuration } // buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. @@ -109,9 +95,18 @@ func (k *Kafka) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfi } k.providers[provider.ID] = ps.New(ctx) } + k.config = &Configuration{ + EventConfiguration: in.CustomEvents.GetKafka(), + Logger: k.logger, + Providers: k.providers, + } return nil } +func (k *Kafka) GetPubSubGeneralImplementerList() datasource.PubSubGeneralImplementer { + return k.config +} + // func (k *Kafka) ConnectProviders(ctx context.Context) error { // for _, provider := range k.providers { // err := provider.Connect() @@ -122,9 +117,9 @@ func (k *Kafka) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfi // return nil // } -func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *datasource.Factory[*nodev1.KafkaEventConfiguration, *KafkaPubSub] { - return datasource.NewFactory[*nodev1.KafkaEventConfiguration](executionContext, config, k.providers) -} +// func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *datasource.Factory { +// return datasource.NewFactory(executionContext, config, k.config) +// } func NewPubSub(logger *zap.Logger) Kafka { return Kafka{ diff --git a/router/pkg/pubsub/kafka/datasource_impl.go b/router/pkg/pubsub/kafka/datasource_impl.go index 2828987f97..3dfd122bd2 100644 --- a/router/pkg/pubsub/kafka/datasource_impl.go +++ b/router/pkg/pubsub/kafka/datasource_impl.go @@ -14,6 +14,7 @@ type Configuration struct { Data string `json:"data"` EventConfiguration []*nodev1.KafkaEventConfiguration Logger *zap.Logger + Providers map[string]*KafkaPubSub } func (c *Configuration) GetEventsDataConfigurations() []*nodev1.KafkaEventConfiguration { @@ -78,12 +79,44 @@ func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.KafkaEvent func (c *Configuration) GetProviderId(eventConfig *nodev1.KafkaEventConfiguration) string { return eventConfig.GetEngineEventConfiguration().GetProviderId() } - -func (c *Configuration) FindEventConfig(eventConfigs []*nodev1.KafkaEventConfiguration, typeName string, fieldName string, fn datasource.ArgumentTemplateCallback) (*nodev1.KafkaEventConfiguration, error) { - for _, cfg := range eventConfigs { +func (c *Configuration) FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType2, error) { + for _, cfg := range c.EventConfiguration { if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - return cfg, nil + return &SelectedConfiguration{ + Provider: c.Providers[c.GetProviderId(cfg)], + EventConfiguration: cfg, + }, nil } } - return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) + return nil, nil +} + +type SelectedConfiguration struct { + Config *Configuration + EventConfiguration *nodev1.KafkaEventConfiguration + Provider *KafkaPubSub +} + +func (c *SelectedConfiguration) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { + return c.EventConfiguration.GetEngineEventConfiguration() +} + +func (c *SelectedConfiguration) GetResolveDataSource() (resolve.DataSource, error) { + return c.Config.GetResolveDataSource(c.EventConfiguration, c.Provider) +} + +func (c *SelectedConfiguration) GetResolveDataSourceInput(event []byte) (string, error) { + return c.Config.GetResolveDataSourceInput(c.EventConfiguration, event) +} + +func (c *SelectedConfiguration) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { + return c.Config.GetResolveDataSourceSubscription(c.EventConfiguration, c.Provider) +} + +func (c *SelectedConfiguration) GetResolveDataSourceSubscriptionInput() (string, error) { + return c.Config.GetResolveDataSourceSubscriptionInput(c.EventConfiguration, c.Provider) +} + +func (c *SelectedConfiguration) GetProviderId() string { + return c.Config.GetProviderId(c.EventConfiguration) } diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go index 701271ae67..ab2e56d295 100644 --- a/router/pkg/pubsub/nats/datasource.go +++ b/router/pkg/pubsub/nats/datasource.go @@ -26,36 +26,21 @@ const ( tsep = "." ) -func GetDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (plan.DataSource, error) { +func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (datasource.PubSubGeneralImplementer, error) { if natsData := in.GetCustomEvents().GetNats(); natsData != nil { k := NewPubSub(logger) err := k.PrepareProviders(ctx, in, dsMeta, config) if err != nil { return nil, err } - factory := k.GetFactory(ctx, config, k.providers) - ds, err := plan.NewDataSourceConfiguration[datasource.Implementer[*nodev1.NatsEventConfiguration, *NatsPubSub]]( - in.Id+"-nats", - factory, - dsMeta, - &Configuration{ - EventConfiguration: natsData, - Logger: logger, - }, - ) - - if err != nil { - return nil, err - } - - return ds, nil + return k.config, nil } return nil, nil } func init() { - datasource.RegisterPubSub(GetDataSource) + datasource.RegisterPubSub(GetPlanDataSource) } func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { @@ -107,6 +92,7 @@ type Nats struct { logger *zap.Logger hostName string // How to get it here? routerListenAddr string // How to get it here? + config *Configuration } type LazyClient struct { @@ -115,31 +101,36 @@ type LazyClient struct { opts []nats.Option client *nats.Conn js jetstream.JetStream + err error } -func (c *LazyClient) Connect(opts ...nats.Option) (err error) { +func (c *LazyClient) Connect(opts ...nats.Option) error { c.once.Do(func() { - c.client, err = nats.Connect(c.url, opts...) - if err != nil { + c.client, c.err = nats.Connect(c.url, opts...) + if c.err != nil { return } - c.js, err = jetstream.New(c.client) + c.js, c.err = jetstream.New(c.client) }) - return + return c.err } -func (c *LazyClient) GetClient() *nats.Conn { +func (c *LazyClient) GetClient() (*nats.Conn, error) { if c.client == nil { - c.Connect() + if err := c.Connect(c.opts...); err != nil { + return nil, err + } } - return c.client + return c.client, c.err } -func (c *LazyClient) GetJetStream() jetstream.JetStream { +func (c *LazyClient) GetJetStream() (jetstream.JetStream, error) { if c.js == nil { - c.Connect() + if err := c.Connect(c.opts...); err != nil { + return nil, err + } } - return c.js + return c.js, c.err } func NewLazyClient(url string, opts ...nats.Option) *LazyClient { @@ -173,13 +164,22 @@ func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfig n.providers[provider.ID] = NewConnector(n.logger, provider.URL, options, n.hostName, n.routerListenAddr).New(ctx) } + n.config = &Configuration{ + EventConfiguration: in.CustomEvents.GetNats(), + Logger: n.logger, + Providers: n.providers, + } return nil } -func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *datasource.Factory[*nodev1.NatsEventConfiguration, *NatsPubSub] { - return datasource.NewFactory[*nodev1.NatsEventConfiguration](executionContext, config, n.providers) +func (n *Nats) GetPubSubGeneralImplementerList() datasource.PubSubGeneralImplementer { + return n.config } +// func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *datasource.Factory[*nodev1.NatsEventConfiguration, *NatsPubSub] { +// return datasource.NewFactory[*nodev1.NatsEventConfiguration](executionContext, config, n.providers) +// } + func NewPubSub(logger *zap.Logger) Nats { return Nats{ logger: logger, @@ -190,6 +190,7 @@ type Configuration struct { Data string `json:"data"` EventConfiguration []*nodev1.NatsEventConfiguration Logger *zap.Logger + Providers map[string]*NatsPubSub } func (c *Configuration) GetEventsDataConfigurations() []*nodev1.NatsEventConfiguration { @@ -296,13 +297,51 @@ func (c *Configuration) transformEventConfig(cfg *nodev1.NatsEventConfiguration, return cfg, nil } -func (c *Configuration) FindEventConfig(eventConfigs []*nodev1.NatsEventConfiguration, typeName string, fieldName string, fn datasource.ArgumentTemplateCallback) (*nodev1.NatsEventConfiguration, error) { - for _, cfg := range eventConfigs { +func (c *Configuration) FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType2, error) { + for _, cfg := range c.EventConfiguration { if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - return c.transformEventConfig(cfg, fn) + transformedCfg, err := c.transformEventConfig(cfg, extractFn) + if err != nil { + return nil, err + } + return &SelectedConfiguration{ + Config: c, + EventConfiguration: transformedCfg, + Provider: c.Providers[c.GetProviderId(transformedCfg)], + }, nil } } - return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) + return nil, nil +} + +type SelectedConfiguration struct { + Config *Configuration + EventConfiguration *nodev1.NatsEventConfiguration + Provider *NatsPubSub +} + +func (c *SelectedConfiguration) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { + return c.EventConfiguration.GetEngineEventConfiguration() +} + +func (c *SelectedConfiguration) GetResolveDataSource() (resolve.DataSource, error) { + return c.Config.GetResolveDataSource(c.EventConfiguration, c.Provider) +} + +func (c *SelectedConfiguration) GetResolveDataSourceInput(event []byte) (string, error) { + return c.Config.GetResolveDataSourceInput(c.EventConfiguration, event) +} + +func (c *SelectedConfiguration) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { + return c.Config.GetResolveDataSourceSubscription(c.EventConfiguration, c.Provider) +} + +func (c *SelectedConfiguration) GetResolveDataSourceSubscriptionInput() (string, error) { + return c.Config.GetResolveDataSourceSubscriptionInput(c.EventConfiguration, c.Provider) +} + +func (c *SelectedConfiguration) GetProviderId() string { + return c.Config.GetProviderId(c.EventConfiguration) } func isValidNatsSubject(subject string) bool { diff --git a/router/pkg/pubsub/nats/nats.go b/router/pkg/pubsub/nats/nats.go index 505ab45fb6..496e23bc1c 100644 --- a/router/pkg/pubsub/nats/nats.go +++ b/router/pkg/pubsub/nats/nats.go @@ -104,7 +104,13 @@ func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfi if event.StreamConfiguration.ConsumerInactiveThreshold > 0 { consumerConfig.InactiveThreshold = time.Duration(event.StreamConfiguration.ConsumerInactiveThreshold) * time.Second } - consumer, err := p.client.GetJetStream().CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) + js, err := p.client.GetJetStream() + if err != nil { + log.Error("getting jetstream client", zap.Error(err)) + return datasource.NewError("failed to get jetstream client", err) + } + + consumer, err := js.CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) if err != nil { log.Error("creating or updating consumer", zap.Error(err)) return datasource.NewError(fmt.Sprintf(`failed to create or update consumer for stream "%s"`, event.StreamConfiguration.StreamName), err) @@ -152,10 +158,16 @@ func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfi return nil } + nc, err := p.client.GetClient() + if err != nil { + log.Error("getting nats client", zap.Error(err)) + return datasource.NewError("failed to get nats client", err) + } + msgChan := make(chan *nats.Msg) subscriptions := make([]*nats.Subscription, len(event.Subjects)) for i, subject := range event.Subjects { - subscription, err := p.client.GetClient().ChanSubscribe(subject, msgChan) + subscription, err := nc.ChanSubscribe(subject, msgChan) if err != nil { log.Error("subscribing to NATS subject", zap.Error(err), zap.String("subscription_subject", subject)) return datasource.NewError(fmt.Sprintf(`failed to subscribe to NATS subject "%s"`, subject), err) @@ -209,7 +221,13 @@ func (p *NatsPubSub) Publish(_ context.Context, event PublishAndRequestEventConf log.Debug("publish", zap.ByteString("data", event.Data)) - err := p.client.GetClient().Publish(event.Subject, event.Data) + nc, err := p.client.GetClient() + if err != nil { + log.Error("getting nats client", zap.Error(err)) + return datasource.NewError("failed to get nats client", err) + } + + err = nc.Publish(event.Subject, event.Data) if err != nil { log.Error("publish error", zap.Error(err)) return datasource.NewError(fmt.Sprintf("error publishing to NATS subject %s", event.Subject), err) @@ -227,7 +245,13 @@ func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventCo log.Debug("request", zap.ByteString("data", event.Data)) - msg, err := p.client.GetClient().RequestWithContext(ctx, event.Subject, event.Data) + nc, err := p.client.GetClient() + if err != nil { + log.Error("getting nats client", zap.Error(err)) + return datasource.NewError("failed to get nats client", err) + } + + msg, err := nc.RequestWithContext(ctx, event.Subject, event.Data) if err != nil { log.Error("request error", zap.Error(err)) return datasource.NewError(fmt.Sprintf("error requesting from NATS subject %s", event.Subject), err) @@ -243,32 +267,40 @@ func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventCo } func (p *NatsPubSub) flush(ctx context.Context) error { - return p.client.GetClient().FlushWithContext(ctx) + nc, err := p.client.GetClient() + if err != nil { + return err + } + return nc.FlushWithContext(ctx) } func (p *NatsPubSub) Shutdown(ctx context.Context) error { + nc, err := p.client.GetClient() + if err != nil { + return nil // Already disconnected or failed to connect + } - if p.client.GetClient().IsClosed() { + if nc.IsClosed() { return nil } - var err error + var shutdownErr error fErr := p.flush(ctx) if fErr != nil { - err = errors.Join(err, fErr) + shutdownErr = errors.Join(shutdownErr, fErr) } - drainErr := p.client.GetClient().Drain() + drainErr := nc.Drain() if drainErr != nil { - err = errors.Join(err, drainErr) + shutdownErr = errors.Join(shutdownErr, drainErr) } // Wait for all subscriptions to be closed p.closeWg.Wait() - if err != nil { - return fmt.Errorf("nats pubsub shutdown: %w", err) + if shutdownErr != nil { + return fmt.Errorf("nats pubsub shutdown: %w", shutdownErr) } return nil From e6645c89fdd891721eef498932daae4666595e6a Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 8 Apr 2025 09:40:39 +0200 Subject: [PATCH 10/49] chore: cleaning up a little, reintroduced host and listener vars --- router/core/factoryresolver.go | 2 ++ router/pkg/pubsub/datasource.go | 4 ++-- router/pkg/pubsub/datasource/datasource.go | 21 ++------------------- router/pkg/pubsub/datasource/planner.go | 12 ++++++------ router/pkg/pubsub/kafka/datasource.go | 2 +- router/pkg/pubsub/kafka/datasource_impl.go | 3 +-- router/pkg/pubsub/nats/datasource.go | 17 +++++++++-------- 7 files changed, 23 insertions(+), 38 deletions(-) diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 73b064ef21..f2198a1c8c 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -410,6 +410,8 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod l.dataSourceMetaData(in), routerEngineConfig.Events, l.logger, + "localhost", + "localhost:8080", ) if err != nil { return nil, err diff --git a/router/pkg/pubsub/datasource.go b/router/pkg/pubsub/datasource.go index 0d8b7131f7..b763cd38e3 100644 --- a/router/pkg/pubsub/datasource.go +++ b/router/pkg/pubsub/datasource.go @@ -15,11 +15,11 @@ import ( _ "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) -func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) ([]plan.DataSource, error) { +func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) ([]plan.DataSource, error) { var pubSubs datasource.PubSubGeneralImplementerList for _, pubSub := range datasource.GetRegisteredPubSubs() { - pubSub, err := pubSub(ctx, in, dsMeta, config, logger) + pubSub, err := pubSub(ctx, in, dsMeta, config, logger, hostName, routerListenAddr) if err != nil { return nil, err } diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index ed3ab2b664..04b5ab0552 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -6,11 +6,10 @@ import ( nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" "go.uber.org/zap" ) -type Getter func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (PubSubGeneralImplementer, error) +type Getter func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (PubSubGeneralImplementer, error) var pubSubs []Getter @@ -24,22 +23,6 @@ func GetRegisteredPubSubs() []Getter { type ArgumentTemplateCallback func(tpl string) (string, error) -type PubSubImplementer[F any] interface { +type PubSubImplementer interface { PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error - //GetFactory(executionContext context.Context, config config.EventsConfiguration) F -} - -type EventConfigType interface { - GetEngineEventConfiguration() *nodev1.EngineEventConfiguration -} - -// Create an interface for Configuration -type Implementer[EC EventConfigType, P any] interface { - GetResolveDataSource(eventConfig EC, pubsub P) (resolve.DataSource, error) - GetResolveDataSourceSubscription(eventConfig EC, pubsub P) (resolve.SubscriptionDataSource, error) - GetResolveDataSourceSubscriptionInput(eventConfig EC, pubsub P) (string, error) - GetResolveDataSourceInput(eventConfig EC, event []byte) (string, error) - GetProviderId(eventConfig EC) string - FindEventConfig(eventConfigs []EC, typeName string, fieldName string, fn ArgumentTemplateCallback) (EC, error) - GetEventsDataConfigurations() []EC } diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index 8d22875b53..21f0244c65 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -12,7 +12,7 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type EventConfigType2 interface { +type EventConfigType interface { GetResolveDataSource() (resolve.DataSource, error) GetResolveDataSourceInput(event []byte) (string, error) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration @@ -22,9 +22,9 @@ type EventConfigType2 interface { type PubSubGeneralImplementerList []PubSubGeneralImplementer -func (p *PubSubGeneralImplementerList) FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType2, error) { +func (p *PubSubGeneralImplementerList) FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType, error) { for _, pubSub := range *p { - eventConfigV, err := pubSub.FindEventConfig2(typeName, fieldName, extractFn) + eventConfigV, err := pubSub.FindEventConfig(typeName, fieldName, extractFn) if err != nil { return nil, err } @@ -36,13 +36,13 @@ func (p *PubSubGeneralImplementerList) FindEventConfig2(typeName string, fieldNa } type PubSubGeneralImplementer interface { - FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType2, error) + FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType, error) } type Planner struct { id int pubSubs PubSubGeneralImplementerList - eventConfig EventConfigType2 + eventConfig EventConfigType rootFieldRef int variables resolve.Variables visitor *plan.Visitor @@ -221,7 +221,7 @@ func (p *Planner) EnterField(ref int) { return p.extractArgumentTemplate(ref, tpl) } - eventConfigV, err := p.pubSubs.FindEventConfig2(typeName, fieldName, extractFn) + eventConfigV, err := p.pubSubs.FindEventConfig(typeName, fieldName, extractFn) if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\": %w", typeName, fieldName, err)) return diff --git a/router/pkg/pubsub/kafka/datasource.go b/router/pkg/pubsub/kafka/datasource.go index 4a2087dabc..e40b5f4f2c 100644 --- a/router/pkg/pubsub/kafka/datasource.go +++ b/router/pkg/pubsub/kafka/datasource.go @@ -16,7 +16,7 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" ) -func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (datasource.PubSubGeneralImplementer, error) { +func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubGeneralImplementer, error) { if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { k := NewPubSub(logger) err := k.PrepareProviders(ctx, in, dsMeta, config) diff --git a/router/pkg/pubsub/kafka/datasource_impl.go b/router/pkg/pubsub/kafka/datasource_impl.go index 3dfd122bd2..a3766b2994 100644 --- a/router/pkg/pubsub/kafka/datasource_impl.go +++ b/router/pkg/pubsub/kafka/datasource_impl.go @@ -11,7 +11,6 @@ import ( ) type Configuration struct { - Data string `json:"data"` EventConfiguration []*nodev1.KafkaEventConfiguration Logger *zap.Logger Providers map[string]*KafkaPubSub @@ -79,7 +78,7 @@ func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.KafkaEvent func (c *Configuration) GetProviderId(eventConfig *nodev1.KafkaEventConfiguration) string { return eventConfig.GetEngineEventConfiguration().GetProviderId() } -func (c *Configuration) FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType2, error) { +func (c *Configuration) FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType, error) { for _, cfg := range c.EventConfiguration { if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { return &SelectedConfiguration{ diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go index ab2e56d295..0b4fa8d4b3 100644 --- a/router/pkg/pubsub/nats/datasource.go +++ b/router/pkg/pubsub/nats/datasource.go @@ -26,9 +26,9 @@ const ( tsep = "." ) -func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger) (datasource.PubSubGeneralImplementer, error) { +func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubGeneralImplementer, error) { if natsData := in.GetCustomEvents().GetNats(); natsData != nil { - k := NewPubSub(logger) + k := NewPubSub(logger, hostName, routerListenAddr) err := k.PrepareProviders(ctx, in, dsMeta, config) if err != nil { return nil, err @@ -90,8 +90,8 @@ func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([ type Nats struct { providers map[string]*NatsPubSub logger *zap.Logger - hostName string // How to get it here? - routerListenAddr string // How to get it here? + hostName string + routerListenAddr string config *Configuration } @@ -180,14 +180,15 @@ func (n *Nats) GetPubSubGeneralImplementerList() datasource.PubSubGeneralImpleme // return datasource.NewFactory[*nodev1.NatsEventConfiguration](executionContext, config, n.providers) // } -func NewPubSub(logger *zap.Logger) Nats { +func NewPubSub(logger *zap.Logger, hostName string, routerListenAddr string) Nats { return Nats{ - logger: logger, + logger: logger, + hostName: hostName, + routerListenAddr: routerListenAddr, } } type Configuration struct { - Data string `json:"data"` EventConfiguration []*nodev1.NatsEventConfiguration Logger *zap.Logger Providers map[string]*NatsPubSub @@ -297,7 +298,7 @@ func (c *Configuration) transformEventConfig(cfg *nodev1.NatsEventConfiguration, return cfg, nil } -func (c *Configuration) FindEventConfig2(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType2, error) { +func (c *Configuration) FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType, error) { for _, cfg := range c.EventConfiguration { if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { transformedCfg, err := c.transformEventConfig(cfg, extractFn) From 5e900cb12d9e8e62dd8c8fce44c6d61ccf28d522 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 8 Apr 2025 10:50:31 +0200 Subject: [PATCH 11/49] feat: introduce InstanceData struct for improved configuration management - Added InstanceData struct to encapsulate hostName and listenAddress. - Updated ExecutorConfigurationBuilder and ExecutorBuildOptions to include InstanceData. - Modified Loader and DefaultFactoryResolver to utilize InstanceData for better clarity and maintainability. - Refactored graphServer to initialize with InstanceData, enhancing the overall structure of the configuration. --- router/core/executor.go | 3 +++ router/core/factoryresolver.go | 21 +++++++++++++++++---- router/core/graph_server.go | 10 ++++++---- 3 files changed, 26 insertions(+), 8 deletions(-) diff --git a/router/core/executor.go b/router/core/executor.go index fbe217cd8a..13c8d350ee 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -29,6 +29,7 @@ type ExecutorConfigurationBuilder struct { transportOptions *TransportOptions subscriptionClientOptions *SubscriptionClientOptions + instanceData InstanceData } type Executor struct { @@ -51,6 +52,7 @@ type ExecutorBuildOptions struct { ApolloCompatibilityFlags config.ApolloCompatibilityFlags ApolloRouterCompatibilityFlags config.ApolloRouterCompatibilityFlags HeartbeatInterval time.Duration + InstanceData InstanceData } func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *ExecutorBuildOptions) (*Executor, error) { @@ -204,6 +206,7 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con b.logger, routerEngineCfg.Execution.EnableSingleFlight, routerEngineCfg.Execution.EnableNetPoll, + b.instanceData, ), b.logger) // this generates the plan config using the data source factories from the config package diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index f2198a1c8c..8d44ecf25d 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -29,13 +29,20 @@ import ( type Loader struct { resolver FactoryResolver // includeInfo controls whether additional information like type usage and field usage is included in the plan de - includeInfo bool - logger *zap.Logger + includeInfo bool + logger *zap.Logger + instanceData InstanceData +} + +type InstanceData struct { + hostName string + listenAddress string } type FactoryResolver interface { ResolveGraphqlFactory(subgraphName string) (plan.PlannerFactory[graphql_datasource.Configuration], error) ResolveStaticFactory() (plan.PlannerFactory[staticdatasource.Configuration], error) + InstanceData() InstanceData } type ApiTransportFactory interface { @@ -58,6 +65,7 @@ type DefaultFactoryResolver struct { subscriptionClient graphql_datasource.GraphQLSubscriptionClient factoryLogger abstractlogger.Logger + instanceData InstanceData } func NewDefaultFactoryResolver( @@ -68,6 +76,7 @@ func NewDefaultFactoryResolver( log *zap.Logger, enableSingleFlight bool, enableNetPoll bool, + instanceData InstanceData, ) *DefaultFactoryResolver { transportFactory := NewTransport(transportOptions) @@ -146,6 +155,10 @@ func (d *DefaultFactoryResolver) ResolveStaticFactory() (factory plan.PlannerFac return d.static, nil } +func (d *DefaultFactoryResolver) InstanceData() InstanceData { + return d.instanceData +} + func NewLoader(includeInfo bool, resolver FactoryResolver, logger *zap.Logger) *Loader { return &Loader{ resolver: resolver, @@ -410,8 +423,8 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod l.dataSourceMetaData(in), routerEngineConfig.Events, l.logger, - "localhost", - "localhost:8080", + l.resolver.InstanceData().hostName, + l.resolver.InstanceData().listenAddress, ) if err != nil { return nil, err diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 7902711674..310b05f7a5 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -84,8 +84,7 @@ type ( runtimeMetrics *rmetric.RuntimeMetrics otlpEngineMetrics *rmetric.EngineMetrics prometheusEngineMetrics *rmetric.EngineMetrics - hostName string - routerListenAddr string + instanceData InstanceData } ) @@ -113,8 +112,10 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC baseRouterConfigVersion: routerConfig.GetVersion(), inFlightRequests: &atomic.Uint64{}, graphMuxList: make([]*graphMux, 0, 1), - routerListenAddr: r.listenAddr, - hostName: r.hostName, + instanceData: InstanceData{ + hostName: r.hostName, + listenAddress: r.listenAddr, + }, } baseOtelAttributes := []attribute.KeyValue{ @@ -925,6 +926,7 @@ func (s *graphServer) buildGraphMux(ctx context.Context, ApolloCompatibilityFlags: s.apolloCompatibilityFlags, ApolloRouterCompatibilityFlags: s.apolloRouterCompatibilityFlags, HeartbeatInterval: s.multipartHeartbeatInterval, + InstanceData: s.instanceData, }, ) if err != nil { From 3b9ac24629ea12614d405ef0b6d45a7770aa607d Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 8 Apr 2025 16:33:09 +0200 Subject: [PATCH 12/49] chore: improved naming, simplify code --- .../subgraphs/availability/availability.go | 2 +- .../availability/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/countries/countries.go | 2 +- .../subgraphs/countries/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/employees/employees.go | 2 +- .../subgraphs/employees/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/family/family.go | 2 +- .../pkg/subgraphs/family/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/hobbies/hobbies.go | 2 +- .../subgraphs/hobbies/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/mood/mood.go | 2 +- demo/pkg/subgraphs/mood/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/products/products.go | 2 +- .../subgraphs/products/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/products_fg/products.go | 2 +- .../products_fg/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/subgraphs.go | 17 +- demo/pkg/subgraphs/test1/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/test1/test1.go | 2 +- router-tests/testenv/testenv.go | 8 +- router/core/factoryresolver.go | 10 +- router/core/plan_generator.go | 4 +- router/pkg/pubsub/datasource.go | 20 +- router/pkg/pubsub/datasource/datasource.go | 28 +- router/pkg/pubsub/datasource/factory.go | 10 +- router/pkg/pubsub/datasource/planner.go | 83 ++-- router/pkg/pubsub/datasource/provider.go | 28 ++ .../pkg/pubsub/kafka/{kafka.go => adapter.go} | 54 +-- router/pkg/pubsub/kafka/datasource.go | 129 ------ router/pkg/pubsub/kafka/datasource_impl.go | 121 ------ ...ion_datasource.go => engine_datasource.go} | 4 +- router/pkg/pubsub/kafka/provider.go | 107 +++++ router/pkg/pubsub/kafka/pubsub_datasource.go | 77 ++++ .../pkg/pubsub/nats/{nats.go => adapter.go} | 64 ++- router/pkg/pubsub/nats/datasource.go | 373 ------------------ ...ion_datasource.go => engine_datasource.go} | 6 +- router/pkg/pubsub/nats/provider.go | 155 ++++++++ router/pkg/pubsub/nats/pubsub_datasource.go | 138 +++++++ router/pkg/pubsub/nats/utils.go | 35 ++ 39 files changed, 685 insertions(+), 822 deletions(-) create mode 100644 router/pkg/pubsub/datasource/provider.go rename router/pkg/pubsub/kafka/{kafka.go => adapter.go} (84%) delete mode 100644 router/pkg/pubsub/kafka/datasource.go delete mode 100644 router/pkg/pubsub/kafka/datasource_impl.go rename router/pkg/pubsub/kafka/{subscription_datasource.go => engine_datasource.go} (98%) create mode 100644 router/pkg/pubsub/kafka/provider.go create mode 100644 router/pkg/pubsub/kafka/pubsub_datasource.go rename router/pkg/pubsub/nats/{nats.go => adapter.go} (86%) delete mode 100644 router/pkg/pubsub/nats/datasource.go rename router/pkg/pubsub/nats/{subscription_datasource.go => engine_datasource.go} (98%) create mode 100644 router/pkg/pubsub/nats/provider.go create mode 100644 router/pkg/pubsub/nats/pubsub_datasource.go create mode 100644 router/pkg/pubsub/nats/utils.go diff --git a/demo/pkg/subgraphs/availability/availability.go b/demo/pkg/subgraphs/availability/availability.go index ec693711cb..3bdd747d48 100644 --- a/demo/pkg/subgraphs/availability/availability.go +++ b/demo/pkg/subgraphs/availability/availability.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/availability/subgraph/generated" ) -func NewSchema(pubSubBySourceName map[string]*nats.NatsPubSub, pubSubName func(string) string) graphql.ExecutableSchema { +func NewSchema(pubSubBySourceName map[string]*nats.Adapter, pubSubName func(string) string) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: pubSubBySourceName, GetPubSubName: pubSubName, diff --git a/demo/pkg/subgraphs/availability/subgraph/resolver.go b/demo/pkg/subgraphs/availability/subgraph/resolver.go index 554b11af76..42feac1afd 100644 --- a/demo/pkg/subgraphs/availability/subgraph/resolver.go +++ b/demo/pkg/subgraphs/availability/subgraph/resolver.go @@ -9,6 +9,6 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter GetPubSubName func(string) string } diff --git a/demo/pkg/subgraphs/countries/countries.go b/demo/pkg/subgraphs/countries/countries.go index 3d5940c1e6..ed876f91a0 100644 --- a/demo/pkg/subgraphs/countries/countries.go +++ b/demo/pkg/subgraphs/countries/countries.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/countries/subgraph/generated" ) -func NewSchema(pubSubBySourceName map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(pubSubBySourceName map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: pubSubBySourceName, }}) diff --git a/demo/pkg/subgraphs/countries/subgraph/resolver.go b/demo/pkg/subgraphs/countries/subgraph/resolver.go index c1f3775ab1..814b836a0c 100644 --- a/demo/pkg/subgraphs/countries/subgraph/resolver.go +++ b/demo/pkg/subgraphs/countries/subgraph/resolver.go @@ -12,5 +12,5 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter } diff --git a/demo/pkg/subgraphs/employees/employees.go b/demo/pkg/subgraphs/employees/employees.go index 5b6b6eae19..d7c67b5daa 100644 --- a/demo/pkg/subgraphs/employees/employees.go +++ b/demo/pkg/subgraphs/employees/employees.go @@ -7,7 +7,7 @@ import ( "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, EmployeesData: subgraph.Employees, diff --git a/demo/pkg/subgraphs/employees/subgraph/resolver.go b/demo/pkg/subgraphs/employees/subgraph/resolver.go index 1f75150132..651f1d8321 100644 --- a/demo/pkg/subgraphs/employees/subgraph/resolver.go +++ b/demo/pkg/subgraphs/employees/subgraph/resolver.go @@ -15,7 +15,7 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter EmployeesData []*model.Employee } diff --git a/demo/pkg/subgraphs/family/family.go b/demo/pkg/subgraphs/family/family.go index 4a682b538d..6978fd4da1 100644 --- a/demo/pkg/subgraphs/family/family.go +++ b/demo/pkg/subgraphs/family/family.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/family/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/family/subgraph/resolver.go b/demo/pkg/subgraphs/family/subgraph/resolver.go index 7c404a867c..d0f9debc31 100644 --- a/demo/pkg/subgraphs/family/subgraph/resolver.go +++ b/demo/pkg/subgraphs/family/subgraph/resolver.go @@ -9,5 +9,5 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter } diff --git a/demo/pkg/subgraphs/hobbies/hobbies.go b/demo/pkg/subgraphs/hobbies/hobbies.go index 926b397c6f..32fa49e785 100644 --- a/demo/pkg/subgraphs/hobbies/hobbies.go +++ b/demo/pkg/subgraphs/hobbies/hobbies.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/hobbies/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/hobbies/subgraph/resolver.go b/demo/pkg/subgraphs/hobbies/subgraph/resolver.go index 696412a06a..0a73c02da2 100644 --- a/demo/pkg/subgraphs/hobbies/subgraph/resolver.go +++ b/demo/pkg/subgraphs/hobbies/subgraph/resolver.go @@ -12,7 +12,7 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter } func (r *Resolver) Employees(hobby model.Hobby) ([]*model.Employee, error) { diff --git a/demo/pkg/subgraphs/mood/mood.go b/demo/pkg/subgraphs/mood/mood.go index 72b6c7118b..f838df05e5 100644 --- a/demo/pkg/subgraphs/mood/mood.go +++ b/demo/pkg/subgraphs/mood/mood.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/mood/subgraph/resolver.go b/demo/pkg/subgraphs/mood/subgraph/resolver.go index 554b11af76..42feac1afd 100644 --- a/demo/pkg/subgraphs/mood/subgraph/resolver.go +++ b/demo/pkg/subgraphs/mood/subgraph/resolver.go @@ -9,6 +9,6 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter GetPubSubName func(string) string } diff --git a/demo/pkg/subgraphs/products/products.go b/demo/pkg/subgraphs/products/products.go index 6a19c4baba..ae165ca566 100644 --- a/demo/pkg/subgraphs/products/products.go +++ b/demo/pkg/subgraphs/products/products.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, TopSecretFederationFactsData: subgraph.TopSecretFederationFacts, diff --git a/demo/pkg/subgraphs/products/subgraph/resolver.go b/demo/pkg/subgraphs/products/subgraph/resolver.go index 64543f0a6d..c8a77e87d5 100644 --- a/demo/pkg/subgraphs/products/subgraph/resolver.go +++ b/demo/pkg/subgraphs/products/subgraph/resolver.go @@ -13,6 +13,6 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter TopSecretFederationFactsData []model.TopSecretFact } diff --git a/demo/pkg/subgraphs/products_fg/products.go b/demo/pkg/subgraphs/products_fg/products.go index 436c34582b..580998919a 100644 --- a/demo/pkg/subgraphs/products_fg/products.go +++ b/demo/pkg/subgraphs/products_fg/products.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products_fg/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, TopSecretFederationFactsData: subgraph.TopSecretFederationFacts, diff --git a/demo/pkg/subgraphs/products_fg/subgraph/resolver.go b/demo/pkg/subgraphs/products_fg/subgraph/resolver.go index c6dca70e73..e20dd809ee 100644 --- a/demo/pkg/subgraphs/products_fg/subgraph/resolver.go +++ b/demo/pkg/subgraphs/products_fg/subgraph/resolver.go @@ -13,6 +13,6 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter TopSecretFederationFactsData []model.TopSecretFact } diff --git a/demo/pkg/subgraphs/subgraphs.go b/demo/pkg/subgraphs/subgraphs.go index 0b43b17aff..42858691bc 100644 --- a/demo/pkg/subgraphs/subgraphs.go +++ b/demo/pkg/subgraphs/subgraphs.go @@ -161,7 +161,7 @@ func subgraphHandler(schema graphql.ExecutableSchema) http.Handler { } type SubgraphOptions struct { - NatsPubSubByProviderID map[string]*natsPubsub.NatsPubSub + NatsPubSubByProviderID map[string]*natsPubsub.Adapter GetPubSubName func(string) string } @@ -207,10 +207,19 @@ func New(ctx context.Context, config *Config) (*Subgraphs, error) { url = defaultSourceNameURL } - natsPubSubByProviderID := map[string]*natsPubsub.NatsPubSub{ - "default": natsPubsub.NewConnector(zap.NewNop(), url, []nats.Option{}, "hostname", "test").New(ctx), - "my-nats": natsPubsub.NewConnector(zap.NewNop(), url, []nats.Option{}, "hostname", "test").New(ctx), + natsPubSubByProviderID := map[string]*natsPubsub.Adapter{} + + defaultAdapter, err := natsPubsub.NewAdapter(ctx, zap.NewNop(), url, []nats.Option{}, "hostname", "test") + if err != nil { + return nil, fmt.Errorf("failed to create default nats adapter: %w", err) + } + natsPubSubByProviderID["default"] = defaultAdapter + + myNatsAdapter, err := natsPubsub.NewAdapter(ctx, zap.NewNop(), url, []nats.Option{}, "hostname", "test") + if err != nil { + return nil, fmt.Errorf("failed to create my-nats adapter: %w", err) } + natsPubSubByProviderID["my-nats"] = myNatsAdapter defaultConnection, err := nats.Connect(url) if err != nil { diff --git a/demo/pkg/subgraphs/test1/subgraph/resolver.go b/demo/pkg/subgraphs/test1/subgraph/resolver.go index 7c404a867c..d0f9debc31 100644 --- a/demo/pkg/subgraphs/test1/subgraph/resolver.go +++ b/demo/pkg/subgraphs/test1/subgraph/resolver.go @@ -9,5 +9,5 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.NatsPubSub + NatsPubSubByProviderID map[string]*nats.Adapter } diff --git a/demo/pkg/subgraphs/test1/test1.go b/demo/pkg/subgraphs/test1/test1.go index 8b74990bc4..5af91d1172 100644 --- a/demo/pkg/subgraphs/test1/test1.go +++ b/demo/pkg/subgraphs/test1/test1.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/test1/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.NatsPubSub) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index e567cf0186..d0d915999d 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -2090,13 +2090,15 @@ func WSWriteJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) func subgraphOptions(ctx context.Context, t testing.TB, logger *zap.Logger, natsData *NatsData, pubSubName func(string) string) *subgraphs.SubgraphOptions { if natsData == nil { return &subgraphs.SubgraphOptions{ - NatsPubSubByProviderID: map[string]*pubsubNats.NatsPubSub{}, + NatsPubSubByProviderID: map[string]*pubsubNats.Adapter{}, GetPubSubName: pubSubName, } } - natsPubSubByProviderID := make(map[string]*pubsubNats.NatsPubSub, len(demoNatsProviders)) + natsPubSubByProviderID := make(map[string]*pubsubNats.Adapter, len(demoNatsProviders)) for _, sourceName := range demoNatsProviders { - natsPubSubByProviderID[sourceName] = pubsubNats.NewConnector(logger, natsData.Params[0].Url, natsData.Params[0].Opts, "hostname", "listenaddr").New(ctx) + adapter, err := pubsubNats.NewAdapter(ctx, logger, natsData.Params[0].Url, natsData.Params[0].Opts, "hostname", "listenaddr") + require.NoError(t, err) + natsPubSubByProviderID[sourceName] = adapter } return &subgraphs.SubgraphOptions{ diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 8d44ecf25d..4499700dcb 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -272,7 +272,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod } for _, in := range engineConfig.DatasourceConfigurations { - var outs []plan.DataSource + var out plan.DataSource switch in.Kind { case nodev1.DataSourceKind_STATIC: @@ -281,7 +281,6 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod return nil, err } - var out plan.DataSource out, err = plan.NewDataSourceConfiguration[staticdatasource.Configuration]( in.Id, factory, @@ -293,7 +292,6 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod if err != nil { return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) } - outs = append(outs, out) case nodev1.DataSourceKind_GRAPHQL: header := http.Header{} @@ -401,7 +399,6 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod return nil, err } - var out plan.DataSource out, err = plan.NewDataSourceConfigurationWithName[graphql_datasource.Configuration]( in.Id, dataSourceName, @@ -412,12 +409,11 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod if err != nil { return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) } - outs = append(outs, out) case nodev1.DataSourceKind_PUBSUB: var err error - outs, err = pubsub.GetDataSourcesFromConfig( + out, err = pubsub.GetDataSourceFromConfig( context.Background(), in, l.dataSourceMetaData(in), @@ -433,7 +429,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod return nil, fmt.Errorf("unknown data source type %q", in.Kind) } - outConfig.DataSources = append(outConfig.DataSources, outs...) + outConfig.DataSources = append(outConfig.DataSources, out) } return &outConfig, nil } diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index e5c8025684..43e8e344b6 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -238,8 +238,8 @@ func (pg *PlanGenerator) buildRouterConfig(configFilePath string) (*nodev1.Route func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, logger *zap.Logger, maxDataSourceCollectorsConcurrency uint) error { routerEngineConfig := RouterEngineConfiguration{} - natSources := map[string]*nats.NatsPubSub{} - kafkaSources := map[string]*kafka.KafkaPubSub{} + natSources := map[string]*nats.Adapter{} + kafkaSources := map[string]*kafka.Adapter{} for _, ds := range routerConfig.GetEngineConfig().GetDatasourceConfigurations() { if ds.GetKind() != nodev1.DataSourceKind_PUBSUB || ds.GetCustomEvents() == nil { continue diff --git a/router/pkg/pubsub/datasource.go b/router/pkg/pubsub/datasource.go index b763cd38e3..d55b8f5532 100644 --- a/router/pkg/pubsub/datasource.go +++ b/router/pkg/pubsub/datasource.go @@ -15,32 +15,32 @@ import ( _ "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) -func GetDataSourcesFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) ([]plan.DataSource, error) { - var pubSubs datasource.PubSubGeneralImplementerList +func GetDataSourceFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (plan.DataSource, error) { + var providers []datasource.PubSubProvider - for _, pubSub := range datasource.GetRegisteredPubSubs() { - pubSub, err := pubSub(ctx, in, dsMeta, config, logger, hostName, routerListenAddr) + for _, providerFactory := range datasource.GetRegisteredProviderFactories() { + provider, err := providerFactory(ctx, in, dsMeta, config, logger, hostName, routerListenAddr) if err != nil { return nil, err } - if pubSub != nil { - pubSubs = append(pubSubs, pubSub) + if provider != nil { + providers = append(providers, provider) } } - if len(pubSubs) == 0 { + if len(providers) == 0 { return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) } ds, err := plan.NewDataSourceConfiguration( in.Id, - datasource.NewFactory(ctx, config, pubSubs), + datasource.NewFactory(ctx, config, providers), dsMeta, - pubSubs, + providers, ) if err != nil { return nil, err } - return []plan.DataSource{ds}, nil + return ds, nil } diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index 04b5ab0552..458a44e709 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -1,28 +1,14 @@ package datasource import ( - "context" - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "go.uber.org/zap" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type Getter func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (PubSubGeneralImplementer, error) - -var pubSubs []Getter - -func RegisterPubSub(pubSub Getter) { - pubSubs = append(pubSubs, pubSub) -} - -func GetRegisteredPubSubs() []Getter { - return pubSubs -} - -type ArgumentTemplateCallback func(tpl string) (string, error) - -type PubSubImplementer interface { - PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error +type PubSubDataSource interface { + GetResolveDataSource() (resolve.DataSource, error) + GetResolveDataSourceInput(event []byte) (string, error) + GetEngineEventConfiguration() *nodev1.EngineEventConfiguration + GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) + GetResolveDataSourceSubscriptionInput() (string, error) } diff --git a/router/pkg/pubsub/datasource/factory.go b/router/pkg/pubsub/datasource/factory.go index 25a1564ae2..ce22f9934b 100644 --- a/router/pkg/pubsub/datasource/factory.go +++ b/router/pkg/pubsub/datasource/factory.go @@ -9,7 +9,7 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" ) -func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers PubSubGeneralImplementerList) *Factory { +func NewFactory(executionContext context.Context, config config.EventsConfiguration, providers []PubSubProvider) *Factory { return &Factory{ providers: providers, executionContext: executionContext, @@ -18,14 +18,14 @@ func NewFactory(executionContext context.Context, config config.EventsConfigurat } type Factory struct { - providers PubSubGeneralImplementerList + providers []PubSubProvider executionContext context.Context config config.EventsConfiguration } -func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[PubSubGeneralImplementerList] { +func (f *Factory) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[[]PubSubProvider] { return &Planner{ - pubSubs: f.providers, + providers: f.providers, } } @@ -33,6 +33,6 @@ func (f *Factory) Context() context.Context { return f.executionContext } -func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[PubSubGeneralImplementerList]) (*ast.Document, bool) { +func (f *Factory) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[[]PubSubProvider]) (*ast.Document, bool) { return nil, false } diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index 21f0244c65..bc035e10ba 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -4,7 +4,6 @@ import ( "fmt" "strings" - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" @@ -12,40 +11,13 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type EventConfigType interface { - GetResolveDataSource() (resolve.DataSource, error) - GetResolveDataSourceInput(event []byte) (string, error) - GetEngineEventConfiguration() *nodev1.EngineEventConfiguration - GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) - GetResolveDataSourceSubscriptionInput() (string, error) -} - -type PubSubGeneralImplementerList []PubSubGeneralImplementer - -func (p *PubSubGeneralImplementerList) FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType, error) { - for _, pubSub := range *p { - eventConfigV, err := pubSub.FindEventConfig(typeName, fieldName, extractFn) - if err != nil { - return nil, err - } - if eventConfigV != nil { - return eventConfigV, nil - } - } - return nil, fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName) -} - -type PubSubGeneralImplementer interface { - FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (EventConfigType, error) -} - type Planner struct { - id int - pubSubs PubSubGeneralImplementerList - eventConfig EventConfigType - rootFieldRef int - variables resolve.Variables - visitor *plan.Visitor + id int + providers []PubSubProvider + pubSubDataSource PubSubDataSource + rootFieldRef int + variables resolve.Variables + visitor *plan.Visitor } func (p *Planner) SetID(id int) { @@ -68,23 +40,23 @@ func (p *Planner) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior { } } -func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[PubSubGeneralImplementerList], _ plan.DataSourcePlannerConfiguration) error { +func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[[]PubSubProvider], _ plan.DataSourcePlannerConfiguration) error { p.visitor = visitor visitor.Walker.RegisterEnterFieldVisitor(p) visitor.Walker.RegisterEnterDocumentVisitor(p) - p.pubSubs = configuration.CustomConfiguration() + p.providers = configuration.CustomConfiguration() return nil } func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { - if p.eventConfig == nil { + if p.pubSubDataSource == nil { // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) return resolve.FetchConfiguration{} } var dataSource resolve.DataSource - dataSource, err := p.eventConfig.GetResolveDataSource() + dataSource, err := p.pubSubDataSource.GetResolveDataSource() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get data source: %w", err)) return resolve.FetchConfiguration{} @@ -96,7 +68,7 @@ func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { return resolve.FetchConfiguration{} } - input, err := p.eventConfig.GetResolveDataSourceInput(event) + input, err := p.pubSubDataSource.GetResolveDataSourceInput(event) if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) return resolve.FetchConfiguration{} @@ -107,24 +79,24 @@ func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { Variables: p.variables, DataSource: dataSource, PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + MergePath: []string{p.pubSubDataSource.GetEngineEventConfiguration().GetFieldName()}, }, } } func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { - if p.eventConfig == nil { + if p.pubSubDataSource == nil { // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure subscription: event manager is nil")) return plan.SubscriptionConfiguration{} } - dataSource, err := p.eventConfig.GetResolveDataSourceSubscription() + dataSource, err := p.pubSubDataSource.GetResolveDataSourceSubscription() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription: %w", err)) return plan.SubscriptionConfiguration{} } - input, err := p.eventConfig.GetResolveDataSourceSubscriptionInput() + input, err := p.pubSubDataSource.GetResolveDataSourceSubscriptionInput() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription input: %w", err)) return plan.SubscriptionConfiguration{} @@ -135,7 +107,7 @@ func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { Variables: p.variables, DataSource: dataSource, PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.eventConfig.GetEngineEventConfiguration().GetFieldName()}, + MergePath: []string{p.pubSubDataSource.GetEngineEventConfiguration().GetFieldName()}, }, } } @@ -221,10 +193,23 @@ func (p *Planner) EnterField(ref int) { return p.extractArgumentTemplate(ref, tpl) } - eventConfigV, err := p.pubSubs.FindEventConfig(typeName, fieldName, extractFn) - if err != nil { - p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\": %w", typeName, fieldName, err)) - return + var pubSubDataSource PubSubDataSource + var err error + + for _, pubSub := range p.providers { + pubSubDataSource, err = pubSub.FindPubSubDataSource(typeName, fieldName, extractFn) + if err != nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\": %w", typeName, fieldName, err)) + return + } + if pubSubDataSource != nil { + break + } } - p.eventConfig = eventConfigV + + if pubSubDataSource == nil { + p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to find event config for type name \"%s\" and field name \"%s\"", typeName, fieldName)) + } + + p.pubSubDataSource = pubSubDataSource } diff --git a/router/pkg/pubsub/datasource/provider.go b/router/pkg/pubsub/datasource/provider.go new file mode 100644 index 0000000000..206d46ba43 --- /dev/null +++ b/router/pkg/pubsub/datasource/provider.go @@ -0,0 +1,28 @@ +package datasource + +import ( + "context" + + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "go.uber.org/zap" +) + +type ProviderFactory func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (PubSubProvider, error) + +var pubSubsFactories []ProviderFactory + +func RegisterProviderFactory(pubSub ProviderFactory) { + pubSubsFactories = append(pubSubsFactories, pubSub) +} + +func GetRegisteredProviderFactories() []ProviderFactory { + return pubSubsFactories +} + +type ArgumentTemplateCallback func(tpl string) (string, error) + +type PubSubProvider interface { + FindPubSubDataSource(typeName string, fieldName string, extractFn ArgumentTemplateCallback) (PubSubDataSource, error) +} diff --git a/router/pkg/pubsub/kafka/kafka.go b/router/pkg/pubsub/kafka/adapter.go similarity index 84% rename from router/pkg/pubsub/kafka/kafka.go rename to router/pkg/pubsub/kafka/adapter.go index ef4920b83e..3f40f1e8a1 100644 --- a/router/pkg/pubsub/kafka/kafka.go +++ b/router/pkg/pubsub/kafka/adapter.go @@ -19,12 +19,6 @@ var ( errClientClosed = errors.New("client closed") ) -type connector struct { - writeClient *LazyClient - opts []kgo.Opt - logger *zap.Logger -} - type LazyClient struct { once sync.Once client *kgo.Client @@ -55,45 +49,33 @@ func NewLazyClient(opts ...kgo.Opt) *LazyClient { } } -func NewConnector(logger *zap.Logger, opts []kgo.Opt) (*connector, error) { - writeClient := NewLazyClient(append(opts, +func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (*Adapter, error) { + ctx, cancel := context.WithCancel(ctx) + if logger == nil { + logger = zap.NewNop() + } + + client := NewLazyClient(append(opts, // For observability, we set the client ID to "router" kgo.ClientID("cosmo.router.producer"))..., ) - return &connector{ - writeClient: writeClient, - opts: opts, - logger: logger, - }, nil -} - -func (c *connector) New(ctx context.Context) *KafkaPubSub { - - ctx, cancel := context.WithCancel(ctx) - - if c.logger == nil { - c.logger = zap.NewNop() - } - - ps := &KafkaPubSub{ + return &Adapter{ ctx: ctx, - logger: c.logger.With(zap.String("pubsub", "kafka")), - opts: c.opts, - writeClient: c.writeClient, + logger: logger.With(zap.String("pubsub", "kafka")), + opts: opts, + writeClient: client, closeWg: sync.WaitGroup{}, cancel: cancel, - } - - return ps + }, nil } -// KafkaPubSub is a Kafka pubsub implementation. +// Adapter is a Kafka pubsub implementation. // It uses the franz-go Kafka client to consume and produce messages. // The pubsub is stateless and does not store any messages. // It uses a single write client to produce messages and a client per topic to consume messages. // Each client polls the Kafka topic for new records and updates the subscriptions with the new data. -type KafkaPubSub struct { +type Adapter struct { ctx context.Context opts []kgo.Opt logger *zap.Logger @@ -103,7 +85,7 @@ type KafkaPubSub struct { } // topicPoller polls the Kafka topic for new records and calls the updateTriggers function. -func (p *KafkaPubSub) topicPoller(ctx context.Context, client *kgo.Client, updater resolve.SubscriptionUpdater) error { +func (p *Adapter) topicPoller(ctx context.Context, client *kgo.Client, updater resolve.SubscriptionUpdater) error { for { select { case <-p.ctx.Done(): // Close the poller if the application context was canceled @@ -157,7 +139,7 @@ func (p *KafkaPubSub) topicPoller(ctx context.Context, client *kgo.Client, updat // Subscribe subscribes to the given topics and updates the subscription updater. // The engine already deduplicates subscriptions with the same topics, stream configuration, extensions, headers, etc. -func (p *KafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { +func (p *Adapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), @@ -207,7 +189,7 @@ func (p *KafkaPubSub) Subscribe(ctx context.Context, event SubscriptionEventConf // Publish publishes the given event to the Kafka topic in a non-blocking way. // Publish errors are logged and returned as a pubsub error. // The event is written with a dedicated write client. -func (p *KafkaPubSub) Publish(ctx context.Context, event PublishEventConfiguration) error { +func (p *Adapter) Publish(ctx context.Context, event PublishEventConfiguration) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "publish"), @@ -241,7 +223,7 @@ func (p *KafkaPubSub) Publish(ctx context.Context, event PublishEventConfigurati return nil } -func (p *KafkaPubSub) Shutdown(ctx context.Context) error { +func (p *Adapter) Shutdown(ctx context.Context) error { err := p.writeClient.GetClient().Flush(ctx) if err != nil { diff --git a/router/pkg/pubsub/kafka/datasource.go b/router/pkg/pubsub/kafka/datasource.go deleted file mode 100644 index e40b5f4f2c..0000000000 --- a/router/pkg/pubsub/kafka/datasource.go +++ /dev/null @@ -1,129 +0,0 @@ -package kafka - -import ( - "context" - "crypto/tls" - "fmt" - "time" - - "github.com/twmb/franz-go/pkg/kgo" - "github.com/twmb/franz-go/pkg/sasl/plain" - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "go.uber.org/zap" - - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" -) - -func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubGeneralImplementer, error) { - if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { - k := NewPubSub(logger) - err := k.PrepareProviders(ctx, in, dsMeta, config) - if err != nil { - return nil, err - } - return k.config, nil - } - - return nil, nil -} - -func init() { - datasource.RegisterPubSub(GetPlanDataSource) -} - -// var _ datasource.PubSubImplementer[*KafkaPubSub] = &Kafka{} - -type Kafka struct { - logger *zap.Logger - providers map[string]*KafkaPubSub - config *Configuration -} - -// buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. -// Only general options like TLS, SASL, etc. are configured here. Specific options like topics, etc. are -// configured in the KafkaPubSub implementation. -func buildKafkaOptions(eventSource config.KafkaEventSource) ([]kgo.Opt, error) { - opts := []kgo.Opt{ - kgo.SeedBrokers(eventSource.Brokers...), - // Ensure proper timeouts are set - kgo.ProduceRequestTimeout(10 * time.Second), - kgo.ConnIdleTimeout(60 * time.Second), - } - - if eventSource.TLS != nil && eventSource.TLS.Enabled { - opts = append(opts, - // Configure TLS. Uses SystemCertPool for RootCAs by default. - kgo.DialTLSConfig(new(tls.Config)), - ) - } - - if eventSource.Authentication != nil && eventSource.Authentication.SASLPlain.Username != nil && eventSource.Authentication.SASLPlain.Password != nil { - opts = append(opts, kgo.SASL(plain.Auth{ - User: *eventSource.Authentication.SASLPlain.Username, - Pass: *eventSource.Authentication.SASLPlain.Password, - }.AsMechanism())) - } - - return opts, nil -} - -func (k *Kafka) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { - definedProviders := make(map[string]bool) - for _, provider := range config.Providers.Kafka { - definedProviders[provider.ID] = true - } - usedProviders := make(map[string]bool) - for _, event := range in.CustomEvents.GetKafka() { - if !definedProviders[event.EngineEventConfiguration.ProviderId] { - return fmt.Errorf("failed to find Kafka provider with ID %s", event.EngineEventConfiguration.ProviderId) - } - usedProviders[event.EngineEventConfiguration.ProviderId] = true - } - for _, provider := range config.Providers.Kafka { - if !usedProviders[provider.ID] { - continue - } - options, err := buildKafkaOptions(provider) - if err != nil { - return fmt.Errorf("failed to build options for Kafka provider with ID \"%s\": %w", provider.ID, err) - } - ps, err := NewConnector(k.logger, options) - if err != nil { - return fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", provider.ID, err) - } - k.providers[provider.ID] = ps.New(ctx) - } - k.config = &Configuration{ - EventConfiguration: in.CustomEvents.GetKafka(), - Logger: k.logger, - Providers: k.providers, - } - return nil -} - -func (k *Kafka) GetPubSubGeneralImplementerList() datasource.PubSubGeneralImplementer { - return k.config -} - -// func (k *Kafka) ConnectProviders(ctx context.Context) error { -// for _, provider := range k.providers { -// err := provider.Connect() -// if err != nil { -// return fmt.Errorf("failed to connect to Kafka provider with ID \"%s\": %w", provider.ID, err) -// } -// } -// return nil -// } - -// func (k *Kafka) GetFactory(executionContext context.Context, config config.EventsConfiguration) *datasource.Factory { -// return datasource.NewFactory(executionContext, config, k.config) -// } - -func NewPubSub(logger *zap.Logger) Kafka { - return Kafka{ - providers: map[string]*KafkaPubSub{}, - logger: logger, - } -} diff --git a/router/pkg/pubsub/kafka/datasource_impl.go b/router/pkg/pubsub/kafka/datasource_impl.go deleted file mode 100644 index a3766b2994..0000000000 --- a/router/pkg/pubsub/kafka/datasource_impl.go +++ /dev/null @@ -1,121 +0,0 @@ -package kafka - -import ( - "encoding/json" - "fmt" - - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" - "go.uber.org/zap" -) - -type Configuration struct { - EventConfiguration []*nodev1.KafkaEventConfiguration - Logger *zap.Logger - Providers map[string]*KafkaPubSub -} - -func (c *Configuration) GetEventsDataConfigurations() []*nodev1.KafkaEventConfiguration { - return c.EventConfiguration -} - -func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.KafkaEventConfiguration, pubsub *KafkaPubSub) (resolve.DataSource, error) { - var dataSource resolve.DataSource - - typeName := eventConfig.GetEngineEventConfiguration().GetType() - switch typeName { - case nodev1.EventType_PUBLISH: - dataSource = &KafkaPublishDataSource{ - pubSub: pubsub, - } - default: - return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String()) - } - - return dataSource, nil -} - -func (c *Configuration) GetResolveDataSourceSubscription(eventConfig *nodev1.KafkaEventConfiguration, pubsub *KafkaPubSub) (resolve.SubscriptionDataSource, error) { - return &SubscriptionSource{ - pubSub: pubsub, - }, nil -} - -func (c *Configuration) GetResolveDataSourceSubscriptionInput(eventConfig *nodev1.KafkaEventConfiguration, pubsub *KafkaPubSub) (string, error) { - providerId := c.GetProviderId(eventConfig) - evtCfg := SubscriptionEventConfiguration{ - ProviderID: providerId, - Topics: eventConfig.GetTopics(), - } - object, err := json.Marshal(evtCfg) - if err != nil { - return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") - } - return string(object), nil -} - -func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.KafkaEventConfiguration, event []byte) (string, error) { - topics := eventConfig.GetTopics() - - if len(topics) != 1 { - return "", fmt.Errorf("publish and request events should define one topic but received %d", len(topics)) - } - - topic := topics[0] - - providerId := c.GetProviderId(eventConfig) - - evtCfg := PublishEventConfiguration{ - ProviderID: providerId, - Topic: topic, - Data: event, - } - - return evtCfg.MarshalJSONTemplate(), nil -} - -func (c *Configuration) GetProviderId(eventConfig *nodev1.KafkaEventConfiguration) string { - return eventConfig.GetEngineEventConfiguration().GetProviderId() -} -func (c *Configuration) FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType, error) { - for _, cfg := range c.EventConfiguration { - if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - return &SelectedConfiguration{ - Provider: c.Providers[c.GetProviderId(cfg)], - EventConfiguration: cfg, - }, nil - } - } - return nil, nil -} - -type SelectedConfiguration struct { - Config *Configuration - EventConfiguration *nodev1.KafkaEventConfiguration - Provider *KafkaPubSub -} - -func (c *SelectedConfiguration) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { - return c.EventConfiguration.GetEngineEventConfiguration() -} - -func (c *SelectedConfiguration) GetResolveDataSource() (resolve.DataSource, error) { - return c.Config.GetResolveDataSource(c.EventConfiguration, c.Provider) -} - -func (c *SelectedConfiguration) GetResolveDataSourceInput(event []byte) (string, error) { - return c.Config.GetResolveDataSourceInput(c.EventConfiguration, event) -} - -func (c *SelectedConfiguration) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { - return c.Config.GetResolveDataSourceSubscription(c.EventConfiguration, c.Provider) -} - -func (c *SelectedConfiguration) GetResolveDataSourceSubscriptionInput() (string, error) { - return c.Config.GetResolveDataSourceSubscriptionInput(c.EventConfiguration, c.Provider) -} - -func (c *SelectedConfiguration) GetProviderId() string { - return c.Config.GetProviderId(c.EventConfiguration) -} diff --git a/router/pkg/pubsub/kafka/subscription_datasource.go b/router/pkg/pubsub/kafka/engine_datasource.go similarity index 98% rename from router/pkg/pubsub/kafka/subscription_datasource.go rename to router/pkg/pubsub/kafka/engine_datasource.go index 0cdc230e87..9b71b24b89 100644 --- a/router/pkg/pubsub/kafka/subscription_datasource.go +++ b/router/pkg/pubsub/kafka/engine_datasource.go @@ -29,7 +29,7 @@ func (s *PublishEventConfiguration) MarshalJSONTemplate() string { } type SubscriptionSource struct { - pubSub *KafkaPubSub + pubSub *Adapter } func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { @@ -63,7 +63,7 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r } type KafkaPublishDataSource struct { - pubSub *KafkaPubSub + pubSub *Adapter } func (s *KafkaPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go new file mode 100644 index 0000000000..f2cc6d5924 --- /dev/null +++ b/router/pkg/pubsub/kafka/provider.go @@ -0,0 +1,107 @@ +package kafka + +import ( + "context" + "crypto/tls" + "fmt" + "time" + + "github.com/twmb/franz-go/pkg/kgo" + "github.com/twmb/franz-go/pkg/sasl/plain" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "go.uber.org/zap" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" +) + +func init() { + datasource.RegisterProviderFactory(GetProvider) +} + +// buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. +// Only general options like TLS, SASL, etc. are configured here. Specific options like topics, etc. are +// configured in the KafkaPubSub implementation. +func buildKafkaOptions(eventSource config.KafkaEventSource) ([]kgo.Opt, error) { + opts := []kgo.Opt{ + kgo.SeedBrokers(eventSource.Brokers...), + // Ensure proper timeouts are set + kgo.ProduceRequestTimeout(10 * time.Second), + kgo.ConnIdleTimeout(60 * time.Second), + } + + if eventSource.TLS != nil && eventSource.TLS.Enabled { + opts = append(opts, + // Configure TLS. Uses SystemCertPool for RootCAs by default. + kgo.DialTLSConfig(new(tls.Config)), + ) + } + + if eventSource.Authentication != nil && eventSource.Authentication.SASLPlain.Username != nil && eventSource.Authentication.SASLPlain.Password != nil { + opts = append(opts, kgo.SASL(plain.Auth{ + User: *eventSource.Authentication.SASLPlain.Username, + Pass: *eventSource.Authentication.SASLPlain.Password, + }.AsMechanism())) + } + + return opts, nil +} + +func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubProvider, error) { + providers := make(map[string]*Adapter) + definedProviders := make(map[string]bool) + for _, provider := range config.Providers.Kafka { + definedProviders[provider.ID] = true + } + usedProviders := make(map[string]bool) + if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { + for _, event := range in.CustomEvents.GetKafka() { + if !definedProviders[event.EngineEventConfiguration.ProviderId] { + return nil, fmt.Errorf("failed to find Kafka provider with ID %s", event.EngineEventConfiguration.ProviderId) + } + usedProviders[event.EngineEventConfiguration.ProviderId] = true + } + + for _, provider := range config.Providers.Kafka { + if !usedProviders[provider.ID] { + continue + } + options, err := buildKafkaOptions(provider) + if err != nil { + return nil, fmt.Errorf("failed to build options for Kafka provider with ID \"%s\": %w", provider.ID, err) + } + adapter, err := NewAdapter(ctx, logger, options) + if err != nil { + return nil, fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", provider.ID, err) + } + providers[provider.ID] = adapter + } + + return &PubSubProvider{ + EventConfiguration: in.CustomEvents.GetKafka(), + Logger: logger, + Providers: providers, + }, nil + } + + return nil, nil +} + +type PubSubProvider struct { + EventConfiguration []*nodev1.KafkaEventConfiguration + Logger *zap.Logger + Providers map[string]*Adapter +} + +func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, extractFn datasource.ArgumentTemplateCallback) (datasource.PubSubDataSource, error) { + for _, cfg := range c.EventConfiguration { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + return &PubSubDataSource{ + KafkaAdapter: c.Providers[cfg.GetEngineEventConfiguration().GetProviderId()], + EventConfiguration: cfg, + }, nil + } + } + return nil, nil +} diff --git a/router/pkg/pubsub/kafka/pubsub_datasource.go b/router/pkg/pubsub/kafka/pubsub_datasource.go new file mode 100644 index 0000000000..0e1e631af5 --- /dev/null +++ b/router/pkg/pubsub/kafka/pubsub_datasource.go @@ -0,0 +1,77 @@ +package kafka + +import ( + "encoding/json" + "fmt" + + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type PubSubDataSource struct { + EventConfiguration *nodev1.KafkaEventConfiguration + KafkaAdapter *Adapter +} + +func (c *PubSubDataSource) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { + return c.EventConfiguration.GetEngineEventConfiguration() +} + +func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { + var dataSource resolve.DataSource + + typeName := c.EventConfiguration.GetEngineEventConfiguration().GetType() + switch typeName { + case nodev1.EventType_PUBLISH: + dataSource = &KafkaPublishDataSource{ + pubSub: c.KafkaAdapter, + } + default: + return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Kafka", typeName.String()) + } + + return dataSource, nil +} + +func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, error) { + topics := c.EventConfiguration.GetTopics() + + if len(topics) != 1 { + return "", fmt.Errorf("publish and request events should define one topic but received %d", len(topics)) + } + + topic := topics[0] + + providerId := c.GetProviderId() + + evtCfg := PublishEventConfiguration{ + ProviderID: providerId, + Topic: topic, + Data: event, + } + + return evtCfg.MarshalJSONTemplate(), nil +} + +func (c *PubSubDataSource) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { + return &SubscriptionSource{ + pubSub: c.KafkaAdapter, + }, nil +} + +func (c *PubSubDataSource) GetResolveDataSourceSubscriptionInput() (string, error) { + providerId := c.GetProviderId() + evtCfg := SubscriptionEventConfiguration{ + ProviderID: providerId, + Topics: c.EventConfiguration.GetTopics(), + } + object, err := json.Marshal(evtCfg) + if err != nil { + return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") + } + return string(object), nil +} + +func (c *PubSubDataSource) GetProviderId() string { + return c.EventConfiguration.GetEngineEventConfiguration().GetProviderId() +} diff --git a/router/pkg/pubsub/nats/nats.go b/router/pkg/pubsub/nats/adapter.go similarity index 86% rename from router/pkg/pubsub/nats/nats.go rename to router/pkg/pubsub/nats/adapter.go index 496e23bc1c..679feeac82 100644 --- a/router/pkg/pubsub/nats/nats.go +++ b/router/pkg/pubsub/nats/adapter.go @@ -16,38 +16,7 @@ import ( "go.uber.org/zap" ) -type connector struct { - client *LazyClient - logger *zap.Logger - hostName string - routerListenAddr string -} - -func NewConnector(logger *zap.Logger, url string, opts []nats.Option, hostName string, routerListenAddr string) *connector { - client := NewLazyClient(url, opts...) - return &connector{ - client: client, - logger: logger, - hostName: hostName, - routerListenAddr: routerListenAddr, - } -} - -func (c *connector) New(ctx context.Context) *NatsPubSub { - if c.logger == nil { - c.logger = zap.NewNop() - } - return &NatsPubSub{ - ctx: ctx, - client: c.client, - logger: c.logger.With(zap.String("pubsub", "nats")), - closeWg: sync.WaitGroup{}, - hostName: c.hostName, - routerListenAddr: c.routerListenAddr, - } -} - -type NatsPubSub struct { +type Adapter struct { ctx context.Context client *LazyClient logger *zap.Logger @@ -60,7 +29,7 @@ type NatsPubSub struct { // We use the hostname and the address the router is listening on, which should provide a good representation // of what a unique instance is from the perspective of the client that has started a subscription to this instance // and want to restart the subscription after a failure on the client or router side. -func (p *NatsPubSub) getInstanceIdentifier() string { +func (p *Adapter) getInstanceIdentifier() string { return fmt.Sprintf("%s-%s", p.hostName, p.routerListenAddr) } @@ -68,7 +37,7 @@ func (p *NatsPubSub) getInstanceIdentifier() string { // we need to make sure that the durable consumer name is unique for each instance and subjects to prevent // multiple routers from changing the same consumer, which would lead to message loss and wrong messages delivered // to the subscribers -func (p *NatsPubSub) getDurableConsumerName(durableName string, subjects []string) (string, error) { +func (p *Adapter) getDurableConsumerName(durableName string, subjects []string) (string, error) { subjHash := xxhash.New() _, err := subjHash.WriteString(p.getInstanceIdentifier()) if err != nil { @@ -84,7 +53,7 @@ func (p *NatsPubSub) getDurableConsumerName(durableName string, subjects []strin return fmt.Sprintf("%s-%x", durableName, subjHash.Sum64()), nil } -func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { +func (p *Adapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "subscribe"), @@ -212,7 +181,7 @@ func (p *NatsPubSub) Subscribe(ctx context.Context, event SubscriptionEventConfi return nil } -func (p *NatsPubSub) Publish(_ context.Context, event PublishAndRequestEventConfiguration) error { +func (p *Adapter) Publish(_ context.Context, event PublishAndRequestEventConfiguration) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "publish"), @@ -236,7 +205,7 @@ func (p *NatsPubSub) Publish(_ context.Context, event PublishAndRequestEventConf return nil } -func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { +func (p *Adapter) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { log := p.logger.With( zap.String("provider_id", event.ProviderID), zap.String("method", "request"), @@ -266,7 +235,7 @@ func (p *NatsPubSub) Request(ctx context.Context, event PublishAndRequestEventCo return err } -func (p *NatsPubSub) flush(ctx context.Context) error { +func (p *Adapter) flush(ctx context.Context) error { nc, err := p.client.GetClient() if err != nil { return err @@ -274,7 +243,7 @@ func (p *NatsPubSub) flush(ctx context.Context) error { return nc.FlushWithContext(ctx) } -func (p *NatsPubSub) Shutdown(ctx context.Context) error { +func (p *Adapter) Shutdown(ctx context.Context) error { nc, err := p.client.GetClient() if err != nil { return nil // Already disconnected or failed to connect @@ -305,3 +274,20 @@ func (p *NatsPubSub) Shutdown(ctx context.Context) error { return nil } + +func NewAdapter(ctx context.Context, logger *zap.Logger, url string, opts []nats.Option, hostName string, routerListenAddr string) (*Adapter, error) { + if logger == nil { + logger = zap.NewNop() + } + + client := NewLazyClient(url, opts...) + + return &Adapter{ + ctx: ctx, + client: client, + logger: logger.With(zap.String("pubsub", "nats")), + closeWg: sync.WaitGroup{}, + hostName: hostName, + routerListenAddr: routerListenAddr, + }, nil +} diff --git a/router/pkg/pubsub/nats/datasource.go b/router/pkg/pubsub/nats/datasource.go deleted file mode 100644 index 0b4fa8d4b3..0000000000 --- a/router/pkg/pubsub/nats/datasource.go +++ /dev/null @@ -1,373 +0,0 @@ -package nats - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "slices" - "strings" - "sync" - "time" - - "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "go.uber.org/zap" - - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" -) - -const ( - fwc = '>' - tsep = "." -) - -func GetPlanDataSource(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubGeneralImplementer, error) { - if natsData := in.GetCustomEvents().GetNats(); natsData != nil { - k := NewPubSub(logger, hostName, routerListenAddr) - err := k.PrepareProviders(ctx, in, dsMeta, config) - if err != nil { - return nil, err - } - return k.config, nil - } - - return nil, nil -} - -func init() { - datasource.RegisterPubSub(GetPlanDataSource) -} - -func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { - opts := []nats.Option{ - nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), - nats.ReconnectJitter(500*time.Millisecond, 2*time.Second), - nats.ClosedHandler(func(conn *nats.Conn) { - logger.Info("NATS connection closed", zap.String("provider_id", eventSource.ID), zap.Error(conn.LastError())) - }), - nats.ConnectHandler(func(nc *nats.Conn) { - logger.Info("NATS connection established", zap.String("provider_id", eventSource.ID), zap.String("url", nc.ConnectedUrlRedacted())) - }), - nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { - if err != nil { - logger.Error("NATS disconnected; will attempt to reconnect", zap.Error(err), zap.String("provider_id", eventSource.ID)) - } else { - logger.Info("NATS disconnected", zap.String("provider_id", eventSource.ID)) - } - }), - nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { - if errors.Is(err, nats.ErrSlowConsumer) { - logger.Warn( - "NATS slow consumer detected. Events are being dropped. Please consider increasing the buffer size or reducing the number of messages being sent.", - zap.Error(err), - zap.String("provider_id", eventSource.ID), - ) - } else { - logger.Error("NATS error", zap.Error(err)) - } - }), - nats.ReconnectHandler(func(conn *nats.Conn) { - logger.Info("NATS reconnected", zap.String("provider_id", eventSource.ID), zap.String("url", conn.ConnectedUrlRedacted())) - }), - } - - if eventSource.Authentication != nil { - if eventSource.Authentication.Token != nil { - opts = append(opts, nats.Token(*eventSource.Authentication.Token)) - } else if eventSource.Authentication.UserInfo.Username != nil && eventSource.Authentication.UserInfo.Password != nil { - opts = append(opts, nats.UserInfo(*eventSource.Authentication.UserInfo.Username, *eventSource.Authentication.UserInfo.Password)) - } - } - - return opts, nil -} - -type Nats struct { - providers map[string]*NatsPubSub - logger *zap.Logger - hostName string - routerListenAddr string - config *Configuration -} - -type LazyClient struct { - once sync.Once - url string - opts []nats.Option - client *nats.Conn - js jetstream.JetStream - err error -} - -func (c *LazyClient) Connect(opts ...nats.Option) error { - c.once.Do(func() { - c.client, c.err = nats.Connect(c.url, opts...) - if c.err != nil { - return - } - c.js, c.err = jetstream.New(c.client) - }) - return c.err -} - -func (c *LazyClient) GetClient() (*nats.Conn, error) { - if c.client == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.client, c.err -} - -func (c *LazyClient) GetJetStream() (jetstream.JetStream, error) { - if c.js == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.js, c.err -} - -func NewLazyClient(url string, opts ...nats.Option) *LazyClient { - return &LazyClient{ - url: url, - opts: opts, - } -} - -func (n *Nats) PrepareProviders(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration) error { - definedProviders := make(map[string]bool) - for _, provider := range config.Providers.Nats { - definedProviders[provider.ID] = true - } - usedProviders := make(map[string]bool) - for _, event := range in.CustomEvents.GetNats() { - if _, found := definedProviders[event.EngineEventConfiguration.ProviderId]; !found { - return fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) - } - usedProviders[event.EngineEventConfiguration.ProviderId] = true - } - n.providers = map[string]*NatsPubSub{} - for _, provider := range config.Providers.Nats { - if !usedProviders[provider.ID] { - continue - } - options, err := buildNatsOptions(provider, n.logger) - if err != nil { - return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", provider.ID, err) - } - - n.providers[provider.ID] = NewConnector(n.logger, provider.URL, options, n.hostName, n.routerListenAddr).New(ctx) - } - n.config = &Configuration{ - EventConfiguration: in.CustomEvents.GetNats(), - Logger: n.logger, - Providers: n.providers, - } - return nil -} - -func (n *Nats) GetPubSubGeneralImplementerList() datasource.PubSubGeneralImplementer { - return n.config -} - -// func (n *Nats) GetFactory(executionContext context.Context, config config.EventsConfiguration, providers map[string]*NatsPubSub) *datasource.Factory[*nodev1.NatsEventConfiguration, *NatsPubSub] { -// return datasource.NewFactory[*nodev1.NatsEventConfiguration](executionContext, config, n.providers) -// } - -func NewPubSub(logger *zap.Logger, hostName string, routerListenAddr string) Nats { - return Nats{ - logger: logger, - hostName: hostName, - routerListenAddr: routerListenAddr, - } -} - -type Configuration struct { - EventConfiguration []*nodev1.NatsEventConfiguration - Logger *zap.Logger - Providers map[string]*NatsPubSub -} - -func (c *Configuration) GetEventsDataConfigurations() []*nodev1.NatsEventConfiguration { - return c.EventConfiguration -} - -func (c *Configuration) GetResolveDataSource(eventConfig *nodev1.NatsEventConfiguration, pubsub *NatsPubSub) (resolve.DataSource, error) { - var dataSource resolve.DataSource - - typeName := eventConfig.GetEngineEventConfiguration().GetType() - switch typeName { - case nodev1.EventType_PUBLISH: - dataSource = &NatsPublishDataSource{ - pubSub: pubsub, - } - case nodev1.EventType_REQUEST: - dataSource = &NatsRequestDataSource{ - pubSub: pubsub, - } - default: - return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Nats", typeName.String()) - } - - return dataSource, nil -} - -func (c *Configuration) GetResolveDataSourceSubscription(eventConfig *nodev1.NatsEventConfiguration, pubsub *NatsPubSub) (resolve.SubscriptionDataSource, error) { - return &SubscriptionSource{ - pubSub: pubsub, - }, nil -} - -func (c *Configuration) GetResolveDataSourceSubscriptionInput(eventConfig *nodev1.NatsEventConfiguration, pubsub *NatsPubSub) (string, error) { - providerId := c.GetProviderId(eventConfig) - - evtCfg := SubscriptionEventConfiguration{ - ProviderID: providerId, - Subjects: eventConfig.GetSubjects(), - } - if eventConfig.StreamConfiguration != nil { - evtCfg.StreamConfiguration = &StreamConfiguration{ - Consumer: eventConfig.StreamConfiguration.ConsumerName, - StreamName: eventConfig.StreamConfiguration.StreamName, - ConsumerInactiveThreshold: eventConfig.StreamConfiguration.ConsumerInactiveThreshold, - } - } - object, err := json.Marshal(evtCfg) - if err != nil { - return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") - } - return string(object), nil -} - -func (c *Configuration) GetResolveDataSourceInput(eventConfig *nodev1.NatsEventConfiguration, event []byte) (string, error) { - subjects := eventConfig.GetSubjects() - - if len(subjects) != 1 { - return "", fmt.Errorf("publish and request events should define one subject but received %d", len(subjects)) - } - - subject := subjects[0] - - providerId := c.GetProviderId(eventConfig) - - evtCfg := PublishEventConfiguration{ - ProviderID: providerId, - Subject: subject, - Data: event, - } - - return evtCfg.MarshalJSONTemplate(), nil -} - -func (c *Configuration) GetProviderId(eventConfig *nodev1.NatsEventConfiguration) string { - return eventConfig.GetEngineEventConfiguration().GetProviderId() -} - -func (c *Configuration) transformEventConfig(cfg *nodev1.NatsEventConfiguration, fn datasource.ArgumentTemplateCallback) (*nodev1.NatsEventConfiguration, error) { - switch v := cfg.GetEngineEventConfiguration().GetType(); v { - case nodev1.EventType_PUBLISH, nodev1.EventType_REQUEST: - extractedSubject, err := fn(cfg.GetSubjects()[0]) - if err != nil { - return cfg, fmt.Errorf("unable to parse subject with id %s", cfg.GetSubjects()[0]) - } - if !isValidNatsSubject(extractedSubject) { - return cfg, fmt.Errorf("invalid subject: %s", extractedSubject) - } - cfg.Subjects = []string{extractedSubject} - case nodev1.EventType_SUBSCRIBE: - extractedSubjects := make([]string, 0, len(cfg.Subjects)) - for _, rawSubject := range cfg.Subjects { - extractedSubject, err := fn(rawSubject) - if err != nil { - return cfg, nil - } - if !isValidNatsSubject(extractedSubject) { - return cfg, fmt.Errorf("invalid subject: %s", extractedSubject) - } - extractedSubjects = append(extractedSubjects, extractedSubject) - } - slices.Sort(extractedSubjects) - cfg.Subjects = extractedSubjects - } - return cfg, nil -} - -func (c *Configuration) FindEventConfig(typeName string, fieldName string, extractFn func(string) (string, error)) (datasource.EventConfigType, error) { - for _, cfg := range c.EventConfiguration { - if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { - transformedCfg, err := c.transformEventConfig(cfg, extractFn) - if err != nil { - return nil, err - } - return &SelectedConfiguration{ - Config: c, - EventConfiguration: transformedCfg, - Provider: c.Providers[c.GetProviderId(transformedCfg)], - }, nil - } - } - return nil, nil -} - -type SelectedConfiguration struct { - Config *Configuration - EventConfiguration *nodev1.NatsEventConfiguration - Provider *NatsPubSub -} - -func (c *SelectedConfiguration) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { - return c.EventConfiguration.GetEngineEventConfiguration() -} - -func (c *SelectedConfiguration) GetResolveDataSource() (resolve.DataSource, error) { - return c.Config.GetResolveDataSource(c.EventConfiguration, c.Provider) -} - -func (c *SelectedConfiguration) GetResolveDataSourceInput(event []byte) (string, error) { - return c.Config.GetResolveDataSourceInput(c.EventConfiguration, event) -} - -func (c *SelectedConfiguration) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { - return c.Config.GetResolveDataSourceSubscription(c.EventConfiguration, c.Provider) -} - -func (c *SelectedConfiguration) GetResolveDataSourceSubscriptionInput() (string, error) { - return c.Config.GetResolveDataSourceSubscriptionInput(c.EventConfiguration, c.Provider) -} - -func (c *SelectedConfiguration) GetProviderId() string { - return c.Config.GetProviderId(c.EventConfiguration) -} - -func isValidNatsSubject(subject string) bool { - if subject == "" { - return false - } - sfwc := false - tokens := strings.Split(subject, tsep) - for _, t := range tokens { - length := len(t) - if length == 0 || sfwc { - return false - } - if length > 1 { - if strings.ContainsAny(t, "\t\n\f\r ") { - return false - } - continue - } - switch t[0] { - case fwc: - sfwc = true - case ' ', '\t', '\n', '\r', '\f': - return false - } - } - return true -} diff --git a/router/pkg/pubsub/nats/subscription_datasource.go b/router/pkg/pubsub/nats/engine_datasource.go similarity index 98% rename from router/pkg/pubsub/nats/subscription_datasource.go rename to router/pkg/pubsub/nats/engine_datasource.go index ef5e2e5fe2..ba7e0ae1e6 100644 --- a/router/pkg/pubsub/nats/subscription_datasource.go +++ b/router/pkg/pubsub/nats/engine_datasource.go @@ -46,7 +46,7 @@ func (s *PublishEventConfiguration) MarshalJSONTemplate() string { } type SubscriptionSource struct { - pubSub *NatsPubSub + pubSub *Adapter } func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { @@ -81,7 +81,7 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r } type NatsPublishDataSource struct { - pubSub *NatsPubSub + pubSub *Adapter } func (s *NatsPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { @@ -104,7 +104,7 @@ func (s *NatsPublishDataSource) LoadWithFiles(ctx context.Context, input []byte, } type NatsRequestDataSource struct { - pubSub *NatsPubSub + pubSub *Adapter } func (s *NatsRequestDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { diff --git a/router/pkg/pubsub/nats/provider.go b/router/pkg/pubsub/nats/provider.go new file mode 100644 index 0000000000..edcba478be --- /dev/null +++ b/router/pkg/pubsub/nats/provider.go @@ -0,0 +1,155 @@ +package nats + +import ( + "context" + "errors" + "fmt" + "slices" + "time" + + "github.com/nats-io/nats.go" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "go.uber.org/zap" +) + +func init() { + datasource.RegisterProviderFactory(GetProvider) +} + +func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { + opts := []nats.Option{ + nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), + nats.ReconnectJitter(500*time.Millisecond, 2*time.Second), + nats.ClosedHandler(func(conn *nats.Conn) { + logger.Info("NATS connection closed", zap.String("provider_id", eventSource.ID), zap.Error(conn.LastError())) + }), + nats.ConnectHandler(func(nc *nats.Conn) { + logger.Info("NATS connection established", zap.String("provider_id", eventSource.ID), zap.String("url", nc.ConnectedUrlRedacted())) + }), + nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { + if err != nil { + logger.Error("NATS disconnected; will attempt to reconnect", zap.Error(err), zap.String("provider_id", eventSource.ID)) + } else { + logger.Info("NATS disconnected", zap.String("provider_id", eventSource.ID)) + } + }), + nats.ErrorHandler(func(conn *nats.Conn, subscription *nats.Subscription, err error) { + if errors.Is(err, nats.ErrSlowConsumer) { + logger.Warn( + "NATS slow consumer detected. Events are being dropped. Please consider increasing the buffer size or reducing the number of messages being sent.", + zap.Error(err), + zap.String("provider_id", eventSource.ID), + ) + } else { + logger.Error("NATS error", zap.Error(err)) + } + }), + nats.ReconnectHandler(func(conn *nats.Conn) { + logger.Info("NATS reconnected", zap.String("provider_id", eventSource.ID), zap.String("url", conn.ConnectedUrlRedacted())) + }), + } + + if eventSource.Authentication != nil { + if eventSource.Authentication.Token != nil { + opts = append(opts, nats.Token(*eventSource.Authentication.Token)) + } else if eventSource.Authentication.UserInfo.Username != nil && eventSource.Authentication.UserInfo.Password != nil { + opts = append(opts, nats.UserInfo(*eventSource.Authentication.UserInfo.Username, *eventSource.Authentication.UserInfo.Password)) + } + } + + return opts, nil +} + +func transformEventConfig(cfg *nodev1.NatsEventConfiguration, fn datasource.ArgumentTemplateCallback) (*nodev1.NatsEventConfiguration, error) { + switch v := cfg.GetEngineEventConfiguration().GetType(); v { + case nodev1.EventType_PUBLISH, nodev1.EventType_REQUEST: + extractedSubject, err := fn(cfg.GetSubjects()[0]) + if err != nil { + return cfg, fmt.Errorf("unable to parse subject with id %s", cfg.GetSubjects()[0]) + } + if !isValidNatsSubject(extractedSubject) { + return cfg, fmt.Errorf("invalid subject: %s", extractedSubject) + } + cfg.Subjects = []string{extractedSubject} + case nodev1.EventType_SUBSCRIBE: + extractedSubjects := make([]string, 0, len(cfg.Subjects)) + for _, rawSubject := range cfg.Subjects { + extractedSubject, err := fn(rawSubject) + if err != nil { + return cfg, nil + } + if !isValidNatsSubject(extractedSubject) { + return cfg, fmt.Errorf("invalid subject: %s", extractedSubject) + } + extractedSubjects = append(extractedSubjects, extractedSubject) + } + slices.Sort(extractedSubjects) + cfg.Subjects = extractedSubjects + } + return cfg, nil +} + +type PubSubProvider struct { + EventConfiguration []*nodev1.NatsEventConfiguration + Logger *zap.Logger + Providers map[string]*Adapter +} + +func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, extractFn datasource.ArgumentTemplateCallback) (datasource.PubSubDataSource, error) { + for _, cfg := range c.EventConfiguration { + if cfg.GetEngineEventConfiguration().GetTypeName() == typeName && cfg.GetEngineEventConfiguration().GetFieldName() == fieldName { + transformedCfg, err := transformEventConfig(cfg, extractFn) + if err != nil { + return nil, err + } + return &PubSubDataSource{ + EventConfiguration: transformedCfg, + NatsAdapter: c.Providers[cfg.GetEngineEventConfiguration().GetProviderId()], + }, nil + } + } + return nil, nil +} + +func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubProvider, error) { + var providers map[string]*Adapter + if natsData := in.GetCustomEvents().GetNats(); natsData != nil { + definedProviders := make(map[string]bool) + for _, provider := range config.Providers.Nats { + definedProviders[provider.ID] = true + } + usedProviders := make(map[string]bool) + for _, event := range in.CustomEvents.GetNats() { + if _, found := definedProviders[event.EngineEventConfiguration.ProviderId]; !found { + return nil, fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) + } + usedProviders[event.EngineEventConfiguration.ProviderId] = true + } + providers = map[string]*Adapter{} + for _, provider := range config.Providers.Nats { + if !usedProviders[provider.ID] { + continue + } + options, err := buildNatsOptions(provider, logger) + if err != nil { + return nil, fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", provider.ID, err) + } + + adapter, err := NewAdapter(ctx, logger, provider.URL, options, hostName, routerListenAddr) + if err != nil { + return nil, fmt.Errorf("failed to create adapter for Nats provider with ID \"%s\": %w", provider.ID, err) + } + providers[provider.ID] = adapter + } + return &PubSubProvider{ + EventConfiguration: in.CustomEvents.GetNats(), + Logger: logger, + Providers: providers, + }, nil + } + + return nil, nil +} diff --git a/router/pkg/pubsub/nats/pubsub_datasource.go b/router/pkg/pubsub/nats/pubsub_datasource.go new file mode 100644 index 0000000000..b084cb9edc --- /dev/null +++ b/router/pkg/pubsub/nats/pubsub_datasource.go @@ -0,0 +1,138 @@ +package nats + +import ( + "encoding/json" + "fmt" + "sync" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type LazyClient struct { + once sync.Once + url string + opts []nats.Option + client *nats.Conn + js jetstream.JetStream + err error +} + +func (c *LazyClient) Connect(opts ...nats.Option) error { + c.once.Do(func() { + c.client, c.err = nats.Connect(c.url, opts...) + if c.err != nil { + return + } + c.js, c.err = jetstream.New(c.client) + }) + return c.err +} + +func (c *LazyClient) GetClient() (*nats.Conn, error) { + if c.client == nil { + if err := c.Connect(c.opts...); err != nil { + return nil, err + } + } + return c.client, c.err +} + +func (c *LazyClient) GetJetStream() (jetstream.JetStream, error) { + if c.js == nil { + if err := c.Connect(c.opts...); err != nil { + return nil, err + } + } + return c.js, c.err +} + +func NewLazyClient(url string, opts ...nats.Option) *LazyClient { + return &LazyClient{ + url: url, + opts: opts, + } +} + +type PubSubDataSource struct { + EventConfiguration *nodev1.NatsEventConfiguration + NatsAdapter *Adapter +} + +func (c *PubSubDataSource) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { + return c.EventConfiguration.GetEngineEventConfiguration() +} + +func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { + var dataSource resolve.DataSource + + typeName := c.EventConfiguration.GetEngineEventConfiguration().GetType() + switch typeName { + case nodev1.EventType_PUBLISH: + dataSource = &NatsPublishDataSource{ + pubSub: c.NatsAdapter, + } + case nodev1.EventType_REQUEST: + dataSource = &NatsRequestDataSource{ + pubSub: c.NatsAdapter, + } + default: + return nil, fmt.Errorf("failed to configure fetch: invalid event type \"%s\" for Nats", typeName.String()) + } + + return dataSource, nil +} + +func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, error) { + subjects := c.EventConfiguration.GetSubjects() + + if len(subjects) != 1 { + return "", fmt.Errorf("publish and request events should define one subject but received %d", len(subjects)) + } + + subject := subjects[0] + + providerId := c.GetProviderId() + + evtCfg := PublishEventConfiguration{ + ProviderID: providerId, + Subject: subject, + Data: event, + } + + return evtCfg.MarshalJSONTemplate(), nil +} + +func (c *PubSubDataSource) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { + return &SubscriptionSource{ + pubSub: c.NatsAdapter, + }, nil +} + +func (c *PubSubDataSource) GetResolveDataSourceSubscriptionInput() (string, error) { + providerId := c.GetProviderId() + + evtCfg := SubscriptionEventConfiguration{ + ProviderID: providerId, + Subjects: c.EventConfiguration.GetSubjects(), + } + if c.EventConfiguration.StreamConfiguration != nil { + evtCfg.StreamConfiguration = &StreamConfiguration{ + Consumer: c.EventConfiguration.StreamConfiguration.ConsumerName, + StreamName: c.EventConfiguration.StreamConfiguration.StreamName, + ConsumerInactiveThreshold: c.EventConfiguration.StreamConfiguration.ConsumerInactiveThreshold, + } + } + object, err := json.Marshal(evtCfg) + if err != nil { + return "", fmt.Errorf("failed to marshal event subscription streamConfiguration") + } + return string(object), nil +} + +func (c *PubSubDataSource) GetProviderId() string { + return c.EventConfiguration.GetEngineEventConfiguration().GetProviderId() +} diff --git a/router/pkg/pubsub/nats/utils.go b/router/pkg/pubsub/nats/utils.go new file mode 100644 index 0000000000..010b666c61 --- /dev/null +++ b/router/pkg/pubsub/nats/utils.go @@ -0,0 +1,35 @@ +package nats + +import "strings" + +const ( + fwc = '>' + tsep = "." +) + +func isValidNatsSubject(subject string) bool { + if subject == "" { + return false + } + sfwc := false + tokens := strings.Split(subject, tsep) + for _, t := range tokens { + length := len(t) + if length == 0 || sfwc { + return false + } + if length > 1 { + if strings.ContainsAny(t, "\t\n\f\r ") { + return false + } + continue + } + switch t[0] { + case fwc: + sfwc = true + case ' ', '\t', '\n', '\r', '\f': + return false + } + } + return true +} From 51f134374925bad3b2e81b5c861abd424edfb969 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 8 Apr 2025 16:42:06 +0200 Subject: [PATCH 13/49] chore: remove unused pub/sub configuration code --- router/core/graph_server.go | 85 ------------------------------------- 1 file changed, 85 deletions(-) diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 310b05f7a5..bf795e9f8e 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -880,11 +880,6 @@ func (s *graphServer) buildGraphMux(ctx context.Context, SubgraphErrorPropagation: s.subgraphErrorPropagation, } - //err = s.buildPubSubConfiguration(ctx, engineConfig, routerEngineConfig) - //if err != nil { - // return nil, fmt.Errorf("failed to build pubsub configuration: %w", err) - //} - ecb := &ExecutorConfigurationBuilder{ introspection: s.introspection, baseURL: s.baseURL, @@ -1178,86 +1173,6 @@ func (s *graphServer) buildGraphMux(ctx context.Context, return gm, nil } -//func (s *graphServer) buildPubSubConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, routerEngineCfg *RouterEngineConfiguration) error { -// datasourceConfigurations := engineConfig.GetDatasourceConfigurations() -// for _, datasourceConfiguration := range datasourceConfigurations { -// if datasourceConfiguration.CustomEvents == nil { -// continue -// } -// -// for _, eventConfiguration := range datasourceConfiguration.GetCustomEvents().GetNats() { -// -// providerID := eventConfiguration.EngineEventConfiguration.GetProviderId() -// // if this source name's provider has already been initiated, do not try to initiate again -// _, ok := s.pubSubProviders.nats[providerID] -// if ok { -// continue -// } -// -// for _, eventSource := range routerEngineCfg.Events.Providers.Nats { -// if eventSource.ID == eventConfiguration.EngineEventConfiguration.GetProviderId() { -// options, err := buildNatsOptions(eventSource, s.logger) -// if err != nil { -// return fmt.Errorf("failed to build options for Nats provider with ID \"%s\": %w", providerID, err) -// } -// natsConnection, err := nats.Connect(eventSource.URL, options...) -// if err != nil { -// return fmt.Errorf("failed to create connection for Nats provider with ID \"%s\": %w", providerID, err) -// } -// js, err := jetstream.New(natsConnection) -// if err != nil { -// return err -// } -// -// s.pubSubProviders.nats[providerID] = pubsubNats.NewConnector(s.logger, natsConnection, js, s.hostName, s.routerListenAddr).New(ctx) -// -// break -// } -// } -// -// _, ok = s.pubSubProviders.nats[providerID] -// if !ok { -// return fmt.Errorf("failed to find Nats provider with ID \"%s\". Ensure the provider definition is part of the config", providerID) -// } -// } -// -// for _, eventConfiguration := range datasourceConfiguration.GetCustomEvents().GetKafka() { -// -// providerID := eventConfiguration.EngineEventConfiguration.GetProviderId() -// // if this source name's provider has already been initiated, do not try to initiate again -// _, ok := s.pubSubProviders.kafka[providerID] -// if ok { -// continue -// } -// -// for _, eventSource := range routerEngineCfg.Events.Providers.Kafka { -// if eventSource.ID == providerID { -// options, err := buildKafkaOptions(eventSource) -// if err != nil { -// return fmt.Errorf("failed to build options for Kafka provider with ID \"%s\": %w", providerID, err) -// } -// ps, err := kafka.NewConnector(s.logger, options) -// if err != nil { -// return fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", providerID, err) -// } -// -// s.pubSubProviders.kafka[providerID] = ps.New(ctx) -// -// break -// } -// } -// -// _, ok = s.pubSubProviders.kafka[providerID] -// if !ok { -// return fmt.Errorf("failed to find Kafka provider with ID \"%s\". Ensure the provider definition is part of the config", providerID) -// } -// } -// -// } -// -// return nil -//} - // wait waits for all in-flight requests to finish. Similar to http.Server.Shutdown we wait in intervals + jitter // to make the shutdown process more efficient. func (s *graphServer) wait(ctx context.Context) error { From 60c9b4b8ff639b6cdd8e52b5bbc48577f1d0a289 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 8 Apr 2025 16:50:23 +0200 Subject: [PATCH 14/49] chore: remove unused field --- router/core/factoryresolver.go | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 4499700dcb..30633f004a 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -29,9 +29,8 @@ import ( type Loader struct { resolver FactoryResolver // includeInfo controls whether additional information like type usage and field usage is included in the plan de - includeInfo bool - logger *zap.Logger - instanceData InstanceData + includeInfo bool + logger *zap.Logger } type InstanceData struct { From 4017c9f532af7d19ae9826ecd801c2db507c0ede Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 8 Apr 2025 17:24:17 +0200 Subject: [PATCH 15/49] chore: remove registry --- router/pkg/pubsub/datasource.go | 46 ---------------- router/pkg/pubsub/datasource/provider.go | 10 ---- router/pkg/pubsub/kafka/provider.go | 4 -- router/pkg/pubsub/nats/provider.go | 4 -- router/pkg/pubsub/pubsub.go | 70 ++++++++++++++++++++++++ 5 files changed, 70 insertions(+), 64 deletions(-) delete mode 100644 router/pkg/pubsub/datasource.go create mode 100644 router/pkg/pubsub/pubsub.go diff --git a/router/pkg/pubsub/datasource.go b/router/pkg/pubsub/datasource.go deleted file mode 100644 index d55b8f5532..0000000000 --- a/router/pkg/pubsub/datasource.go +++ /dev/null @@ -1,46 +0,0 @@ -package pubsub - -import ( - "context" - "fmt" - - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "go.uber.org/zap" - - // Register all PubSub implementations - _ "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" - _ "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" -) - -func GetDataSourceFromConfig(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (plan.DataSource, error) { - var providers []datasource.PubSubProvider - - for _, providerFactory := range datasource.GetRegisteredProviderFactories() { - provider, err := providerFactory(ctx, in, dsMeta, config, logger, hostName, routerListenAddr) - if err != nil { - return nil, err - } - if provider != nil { - providers = append(providers, provider) - } - } - - if len(providers) == 0 { - return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) - } - - ds, err := plan.NewDataSourceConfiguration( - in.Id, - datasource.NewFactory(ctx, config, providers), - dsMeta, - providers, - ) - if err != nil { - return nil, err - } - - return ds, nil -} diff --git a/router/pkg/pubsub/datasource/provider.go b/router/pkg/pubsub/datasource/provider.go index 206d46ba43..d496a4e3a0 100644 --- a/router/pkg/pubsub/datasource/provider.go +++ b/router/pkg/pubsub/datasource/provider.go @@ -11,16 +11,6 @@ import ( type ProviderFactory func(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (PubSubProvider, error) -var pubSubsFactories []ProviderFactory - -func RegisterProviderFactory(pubSub ProviderFactory) { - pubSubsFactories = append(pubSubsFactories, pubSub) -} - -func GetRegisteredProviderFactories() []ProviderFactory { - return pubSubsFactories -} - type ArgumentTemplateCallback func(tpl string) (string, error) type PubSubProvider interface { diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go index f2cc6d5924..664745697e 100644 --- a/router/pkg/pubsub/kafka/provider.go +++ b/router/pkg/pubsub/kafka/provider.go @@ -16,10 +16,6 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" ) -func init() { - datasource.RegisterProviderFactory(GetProvider) -} - // buildKafkaOptions creates a list of kgo.Opt options for the given Kafka event source configuration. // Only general options like TLS, SASL, etc. are configured here. Specific options like topics, etc. are // configured in the KafkaPubSub implementation. diff --git a/router/pkg/pubsub/nats/provider.go b/router/pkg/pubsub/nats/provider.go index edcba478be..8e7815722f 100644 --- a/router/pkg/pubsub/nats/provider.go +++ b/router/pkg/pubsub/nats/provider.go @@ -15,10 +15,6 @@ import ( "go.uber.org/zap" ) -func init() { - datasource.RegisterProviderFactory(GetProvider) -} - func buildNatsOptions(eventSource config.NatsEventSource, logger *zap.Logger) ([]nats.Option, error) { opts := []nats.Option{ nats.Name(fmt.Sprintf("cosmo.router.edfs.nats.%s", eventSource.ID)), diff --git a/router/pkg/pubsub/pubsub.go b/router/pkg/pubsub/pubsub.go new file mode 100644 index 0000000000..8ae599be28 --- /dev/null +++ b/router/pkg/pubsub/pubsub.go @@ -0,0 +1,70 @@ +package pubsub + +import ( + "context" + "fmt" + + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "go.uber.org/zap" + + "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" + "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" +) + +var additionalProviders []datasource.ProviderFactory + +// RegisterAdditionalProvider registers an additional PubSub provider +func RegisterAdditionalProvider(provider datasource.ProviderFactory) { + additionalProviders = append(additionalProviders, provider) +} + +// GetProviderFactories returns a list of all PubSub implementations +func GetProviderFactories() []datasource.ProviderFactory { + return append([]datasource.ProviderFactory{ + kafka.GetProvider, + nats.GetProvider, + }, additionalProviders...) +} + +// GetDataSourceFromConfig returns a new plan.DataSource from the given configuration, +// using the registered PubSub providers. +func GetDataSourceFromConfig( + ctx context.Context, + in *nodev1.DataSourceConfiguration, + dsMeta *plan.DataSourceMetadata, + config config.EventsConfiguration, + logger *zap.Logger, + hostName string, + routerListenAddr string, +) (plan.DataSource, error) { + var providers []datasource.PubSubProvider + + for _, providerFactory := range GetProviderFactories() { + provider, err := providerFactory(ctx, in, dsMeta, config, logger, hostName, routerListenAddr) + if err != nil { + return nil, err + } + if provider != nil { + providers = append(providers, provider) + } + } + + if len(providers) == 0 { + return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) + } + + ds, err := plan.NewDataSourceConfiguration( + in.Id, + datasource.NewFactory(ctx, config, providers), + dsMeta, + providers, + ) + if err != nil { + return nil, err + } + + return ds, nil +} From e76ceb303bc22c95b72931232498b9cdb3730acc Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 9 Apr 2025 12:46:07 +0200 Subject: [PATCH 16/49] fix: improve error message in Kafka provider --- router/pkg/pubsub/kafka/provider.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go index 664745697e..9ecb5b95d0 100644 --- a/router/pkg/pubsub/kafka/provider.go +++ b/router/pkg/pubsub/kafka/provider.go @@ -69,7 +69,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta } adapter, err := NewAdapter(ctx, logger, options) if err != nil { - return nil, fmt.Errorf("failed to create connection for Kafka provider with ID \"%s\": %w", provider.ID, err) + return nil, fmt.Errorf("failed to create adapter for Kafka provider with ID \"%s\": %w", provider.ID, err) } providers[provider.ID] = adapter } From 47e83bd64e426aaf000fce4d9facf98e11df7f48 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 9 Apr 2025 13:11:23 +0200 Subject: [PATCH 17/49] chore: remove commented error handling code in Planner.ConfigureFetch --- router/pkg/pubsub/datasource/planner.go | 1 - 1 file changed, 1 deletion(-) diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index bc035e10ba..9f9ae3326a 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -50,7 +50,6 @@ func (p *Planner) Register(visitor *plan.Visitor, configuration plan.DataSourceC func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { if p.pubSubDataSource == nil { - // p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to configure fetch: event config is nil")) return resolve.FetchConfiguration{} } From 546be70aa144a2395c81346fc63c9e523ea94a9a Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 9 Apr 2025 18:51:29 +0200 Subject: [PATCH 18/49] chore: add adapter interfaces and tests --- .../subgraphs/availability/availability.go | 2 +- .../availability/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/countries/countries.go | 2 +- .../subgraphs/countries/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/employees/employees.go | 2 +- .../subgraphs/employees/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/family/family.go | 2 +- .../pkg/subgraphs/family/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/hobbies/hobbies.go | 2 +- .../subgraphs/hobbies/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/mood/mood.go | 2 +- demo/pkg/subgraphs/mood/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/products/products.go | 2 +- .../subgraphs/products/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/products_fg/products.go | 2 +- .../products_fg/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/subgraphs.go | 4 +- demo/pkg/subgraphs/test1/subgraph/resolver.go | 2 +- demo/pkg/subgraphs/test1/test1.go | 2 +- router-tests/testenv/testenv.go | 4 +- router/pkg/pubsub/datasource/test_utils.go | 47 +++ router/pkg/pubsub/kafka/adapter.go | 11 +- router/pkg/pubsub/kafka/engine_datasource.go | 4 +- router/pkg/pubsub/kafka/provider.go | 4 +- router/pkg/pubsub/kafka/pubsub_datasource.go | 2 +- .../pubsub/kafka/pubsub_datasource_test.go | 170 ++++++++++ router/pkg/pubsub/nats/adapter.go | 15 +- router/pkg/pubsub/nats/engine_datasource.go | 6 +- router/pkg/pubsub/nats/provider.go | 6 +- router/pkg/pubsub/nats/pubsub_datasource.go | 2 +- .../pkg/pubsub/nats/pubsub_datasource_test.go | 300 ++++++++++++++++++ 31 files changed, 575 insertions(+), 36 deletions(-) create mode 100644 router/pkg/pubsub/datasource/test_utils.go create mode 100644 router/pkg/pubsub/kafka/pubsub_datasource_test.go create mode 100644 router/pkg/pubsub/nats/pubsub_datasource_test.go diff --git a/demo/pkg/subgraphs/availability/availability.go b/demo/pkg/subgraphs/availability/availability.go index 3bdd747d48..412ffd1ac3 100644 --- a/demo/pkg/subgraphs/availability/availability.go +++ b/demo/pkg/subgraphs/availability/availability.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/availability/subgraph/generated" ) -func NewSchema(pubSubBySourceName map[string]*nats.Adapter, pubSubName func(string) string) graphql.ExecutableSchema { +func NewSchema(pubSubBySourceName map[string]nats.AdapterInterface, pubSubName func(string) string) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: pubSubBySourceName, GetPubSubName: pubSubName, diff --git a/demo/pkg/subgraphs/availability/subgraph/resolver.go b/demo/pkg/subgraphs/availability/subgraph/resolver.go index 42feac1afd..a5461fac82 100644 --- a/demo/pkg/subgraphs/availability/subgraph/resolver.go +++ b/demo/pkg/subgraphs/availability/subgraph/resolver.go @@ -9,6 +9,6 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface GetPubSubName func(string) string } diff --git a/demo/pkg/subgraphs/countries/countries.go b/demo/pkg/subgraphs/countries/countries.go index ed876f91a0..726ef1834a 100644 --- a/demo/pkg/subgraphs/countries/countries.go +++ b/demo/pkg/subgraphs/countries/countries.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/countries/subgraph/generated" ) -func NewSchema(pubSubBySourceName map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(pubSubBySourceName map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: pubSubBySourceName, }}) diff --git a/demo/pkg/subgraphs/countries/subgraph/resolver.go b/demo/pkg/subgraphs/countries/subgraph/resolver.go index 814b836a0c..e2f350a70a 100644 --- a/demo/pkg/subgraphs/countries/subgraph/resolver.go +++ b/demo/pkg/subgraphs/countries/subgraph/resolver.go @@ -12,5 +12,5 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface } diff --git a/demo/pkg/subgraphs/employees/employees.go b/demo/pkg/subgraphs/employees/employees.go index d7c67b5daa..aa8c38b134 100644 --- a/demo/pkg/subgraphs/employees/employees.go +++ b/demo/pkg/subgraphs/employees/employees.go @@ -7,7 +7,7 @@ import ( "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, EmployeesData: subgraph.Employees, diff --git a/demo/pkg/subgraphs/employees/subgraph/resolver.go b/demo/pkg/subgraphs/employees/subgraph/resolver.go index 651f1d8321..8f416f6390 100644 --- a/demo/pkg/subgraphs/employees/subgraph/resolver.go +++ b/demo/pkg/subgraphs/employees/subgraph/resolver.go @@ -15,7 +15,7 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface EmployeesData []*model.Employee } diff --git a/demo/pkg/subgraphs/family/family.go b/demo/pkg/subgraphs/family/family.go index 6978fd4da1..aa3234ab45 100644 --- a/demo/pkg/subgraphs/family/family.go +++ b/demo/pkg/subgraphs/family/family.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/family/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/family/subgraph/resolver.go b/demo/pkg/subgraphs/family/subgraph/resolver.go index d0f9debc31..6cea8bd318 100644 --- a/demo/pkg/subgraphs/family/subgraph/resolver.go +++ b/demo/pkg/subgraphs/family/subgraph/resolver.go @@ -9,5 +9,5 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface } diff --git a/demo/pkg/subgraphs/hobbies/hobbies.go b/demo/pkg/subgraphs/hobbies/hobbies.go index 32fa49e785..cf43f3ddfb 100644 --- a/demo/pkg/subgraphs/hobbies/hobbies.go +++ b/demo/pkg/subgraphs/hobbies/hobbies.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/hobbies/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/hobbies/subgraph/resolver.go b/demo/pkg/subgraphs/hobbies/subgraph/resolver.go index 0a73c02da2..9206076fc1 100644 --- a/demo/pkg/subgraphs/hobbies/subgraph/resolver.go +++ b/demo/pkg/subgraphs/hobbies/subgraph/resolver.go @@ -12,7 +12,7 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface } func (r *Resolver) Employees(hobby model.Hobby) ([]*model.Employee, error) { diff --git a/demo/pkg/subgraphs/mood/mood.go b/demo/pkg/subgraphs/mood/mood.go index f838df05e5..8bce0d004b 100644 --- a/demo/pkg/subgraphs/mood/mood.go +++ b/demo/pkg/subgraphs/mood/mood.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/demo/pkg/subgraphs/mood/subgraph/resolver.go b/demo/pkg/subgraphs/mood/subgraph/resolver.go index 42feac1afd..a5461fac82 100644 --- a/demo/pkg/subgraphs/mood/subgraph/resolver.go +++ b/demo/pkg/subgraphs/mood/subgraph/resolver.go @@ -9,6 +9,6 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface GetPubSubName func(string) string } diff --git a/demo/pkg/subgraphs/products/products.go b/demo/pkg/subgraphs/products/products.go index ae165ca566..6a2c8c5984 100644 --- a/demo/pkg/subgraphs/products/products.go +++ b/demo/pkg/subgraphs/products/products.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, TopSecretFederationFactsData: subgraph.TopSecretFederationFacts, diff --git a/demo/pkg/subgraphs/products/subgraph/resolver.go b/demo/pkg/subgraphs/products/subgraph/resolver.go index c8a77e87d5..db8de70257 100644 --- a/demo/pkg/subgraphs/products/subgraph/resolver.go +++ b/demo/pkg/subgraphs/products/subgraph/resolver.go @@ -13,6 +13,6 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface TopSecretFederationFactsData []model.TopSecretFact } diff --git a/demo/pkg/subgraphs/products_fg/products.go b/demo/pkg/subgraphs/products_fg/products.go index 580998919a..73e2082af4 100644 --- a/demo/pkg/subgraphs/products_fg/products.go +++ b/demo/pkg/subgraphs/products_fg/products.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/products_fg/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, TopSecretFederationFactsData: subgraph.TopSecretFederationFacts, diff --git a/demo/pkg/subgraphs/products_fg/subgraph/resolver.go b/demo/pkg/subgraphs/products_fg/subgraph/resolver.go index e20dd809ee..78a07d51ce 100644 --- a/demo/pkg/subgraphs/products_fg/subgraph/resolver.go +++ b/demo/pkg/subgraphs/products_fg/subgraph/resolver.go @@ -13,6 +13,6 @@ import ( type Resolver struct { mux sync.Mutex - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface TopSecretFederationFactsData []model.TopSecretFact } diff --git a/demo/pkg/subgraphs/subgraphs.go b/demo/pkg/subgraphs/subgraphs.go index 42858691bc..c5e91ab034 100644 --- a/demo/pkg/subgraphs/subgraphs.go +++ b/demo/pkg/subgraphs/subgraphs.go @@ -161,7 +161,7 @@ func subgraphHandler(schema graphql.ExecutableSchema) http.Handler { } type SubgraphOptions struct { - NatsPubSubByProviderID map[string]*natsPubsub.Adapter + NatsPubSubByProviderID map[string]natsPubsub.AdapterInterface GetPubSubName func(string) string } @@ -207,7 +207,7 @@ func New(ctx context.Context, config *Config) (*Subgraphs, error) { url = defaultSourceNameURL } - natsPubSubByProviderID := map[string]*natsPubsub.Adapter{} + natsPubSubByProviderID := map[string]natsPubsub.AdapterInterface{} defaultAdapter, err := natsPubsub.NewAdapter(ctx, zap.NewNop(), url, []nats.Option{}, "hostname", "test") if err != nil { diff --git a/demo/pkg/subgraphs/test1/subgraph/resolver.go b/demo/pkg/subgraphs/test1/subgraph/resolver.go index d0f9debc31..6cea8bd318 100644 --- a/demo/pkg/subgraphs/test1/subgraph/resolver.go +++ b/demo/pkg/subgraphs/test1/subgraph/resolver.go @@ -9,5 +9,5 @@ import ( // It serves as dependency injection for your app, add any dependencies you require here. type Resolver struct { - NatsPubSubByProviderID map[string]*nats.Adapter + NatsPubSubByProviderID map[string]nats.AdapterInterface } diff --git a/demo/pkg/subgraphs/test1/test1.go b/demo/pkg/subgraphs/test1/test1.go index 5af91d1172..3234af4209 100644 --- a/demo/pkg/subgraphs/test1/test1.go +++ b/demo/pkg/subgraphs/test1/test1.go @@ -8,7 +8,7 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/test1/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]*nats.Adapter) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, }}) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index d0d915999d..f0a2003ba5 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -2090,11 +2090,11 @@ func WSWriteJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) func subgraphOptions(ctx context.Context, t testing.TB, logger *zap.Logger, natsData *NatsData, pubSubName func(string) string) *subgraphs.SubgraphOptions { if natsData == nil { return &subgraphs.SubgraphOptions{ - NatsPubSubByProviderID: map[string]*pubsubNats.Adapter{}, + NatsPubSubByProviderID: map[string]pubsubNats.AdapterInterface{}, GetPubSubName: pubSubName, } } - natsPubSubByProviderID := make(map[string]*pubsubNats.Adapter, len(demoNatsProviders)) + natsPubSubByProviderID := make(map[string]pubsubNats.AdapterInterface, len(demoNatsProviders)) for _, sourceName := range demoNatsProviders { adapter, err := pubsubNats.NewAdapter(ctx, logger, natsData.Params[0].Url, natsData.Params[0].Opts, "hostname", "listenaddr") require.NoError(t, err) diff --git a/router/pkg/pubsub/datasource/test_utils.go b/router/pkg/pubsub/datasource/test_utils.go new file mode 100644 index 0000000000..33500536ce --- /dev/null +++ b/router/pkg/pubsub/datasource/test_utils.go @@ -0,0 +1,47 @@ +package datasource + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// VerifyPubSubDataSourceImplementation is a common test function to verify any PubSubDataSource implementation +// This function can be used by other packages to test their PubSubDataSource implementations +func VerifyPubSubDataSourceImplementation(t *testing.T, pubSub PubSubDataSource) { + // Test GetEngineEventConfiguration + engineEventConfig := pubSub.GetEngineEventConfiguration() + require.NotNil(t, engineEventConfig, "Expected non-nil EngineEventConfiguration") + + // Test GetResolveDataSource + dataSource, err := pubSub.GetResolveDataSource() + require.NoError(t, err, "Expected no error from GetResolveDataSource") + require.NotNil(t, dataSource, "Expected non-nil DataSource") + + // Test GetResolveDataSourceInput with sample event data + testEvent := []byte(`{"test":"data"}`) + input, err := pubSub.GetResolveDataSourceInput(testEvent) + require.NoError(t, err, "Expected no error from GetResolveDataSourceInput") + assert.NotEmpty(t, input, "Expected non-empty input") + + // Make sure the input is valid JSON + var result interface{} + err = json.Unmarshal([]byte(input), &result) + assert.NoError(t, err, "Expected valid JSON from GetResolveDataSourceInput") + + // Test GetResolveDataSourceSubscription + subscription, err := pubSub.GetResolveDataSourceSubscription() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscription") + require.NotNil(t, subscription, "Expected non-nil SubscriptionDataSource") + + // Test GetResolveDataSourceSubscriptionInput + subscriptionInput, err := pubSub.GetResolveDataSourceSubscriptionInput() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") + assert.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") + + // Make sure the subscription input is valid JSON + err = json.Unmarshal([]byte(subscriptionInput), &result) + assert.NoError(t, err, "Expected valid JSON from GetResolveDataSourceSubscriptionInput") +} diff --git a/router/pkg/pubsub/kafka/adapter.go b/router/pkg/pubsub/kafka/adapter.go index 3f40f1e8a1..675c225c9d 100644 --- a/router/pkg/pubsub/kafka/adapter.go +++ b/router/pkg/pubsub/kafka/adapter.go @@ -49,7 +49,7 @@ func NewLazyClient(opts ...kgo.Opt) *LazyClient { } } -func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (*Adapter, error) { +func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (AdapterInterface, error) { ctx, cancel := context.WithCancel(ctx) if logger == nil { logger = zap.NewNop() @@ -70,6 +70,12 @@ func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (*Adapt }, nil } +// AdapterInterface defines the interface for Kafka adapter operations +type AdapterInterface interface { + Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error + Publish(ctx context.Context, event PublishEventConfiguration) error +} + // Adapter is a Kafka pubsub implementation. // It uses the franz-go Kafka client to consume and produce messages. // The pubsub is stateless and does not store any messages. @@ -84,6 +90,9 @@ type Adapter struct { cancel context.CancelFunc } +// Ensure Adapter implements AdapterInterface +var _ AdapterInterface = (*Adapter)(nil) + // topicPoller polls the Kafka topic for new records and calls the updateTriggers function. func (p *Adapter) topicPoller(ctx context.Context, client *kgo.Client, updater resolve.SubscriptionUpdater) error { for { diff --git a/router/pkg/pubsub/kafka/engine_datasource.go b/router/pkg/pubsub/kafka/engine_datasource.go index 9b71b24b89..f342469f52 100644 --- a/router/pkg/pubsub/kafka/engine_datasource.go +++ b/router/pkg/pubsub/kafka/engine_datasource.go @@ -29,7 +29,7 @@ func (s *PublishEventConfiguration) MarshalJSONTemplate() string { } type SubscriptionSource struct { - pubSub *Adapter + pubSub AdapterInterface } func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { @@ -63,7 +63,7 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r } type KafkaPublishDataSource struct { - pubSub *Adapter + pubSub AdapterInterface } func (s *KafkaPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go index 9ecb5b95d0..ff767a7de0 100644 --- a/router/pkg/pubsub/kafka/provider.go +++ b/router/pkg/pubsub/kafka/provider.go @@ -45,7 +45,7 @@ func buildKafkaOptions(eventSource config.KafkaEventSource) ([]kgo.Opt, error) { } func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubProvider, error) { - providers := make(map[string]*Adapter) + providers := make(map[string]AdapterInterface) definedProviders := make(map[string]bool) for _, provider := range config.Providers.Kafka { definedProviders[provider.ID] = true @@ -87,7 +87,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta type PubSubProvider struct { EventConfiguration []*nodev1.KafkaEventConfiguration Logger *zap.Logger - Providers map[string]*Adapter + Providers map[string]AdapterInterface } func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, extractFn datasource.ArgumentTemplateCallback) (datasource.PubSubDataSource, error) { diff --git a/router/pkg/pubsub/kafka/pubsub_datasource.go b/router/pkg/pubsub/kafka/pubsub_datasource.go index 0e1e631af5..fac984411e 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource.go @@ -10,7 +10,7 @@ import ( type PubSubDataSource struct { EventConfiguration *nodev1.KafkaEventConfiguration - KafkaAdapter *Adapter + KafkaAdapter AdapterInterface } func (c *PubSubDataSource) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go new file mode 100644 index 0000000000..a090e3ca87 --- /dev/null +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -0,0 +1,170 @@ +package kafka + +import ( + "bytes" + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +// MockAdapter mocks the required functionality from the Adapter for testing +type MockAdapter struct { + mock.Mock +} + +// Ensure MockAdapter implements KafkaAdapterInterface +var _ AdapterInterface = (*MockAdapter)(nil) + +func (m *MockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { + args := m.Called(ctx, event, updater) + return args.Error(0) +} + +func (m *MockAdapter) Publish(ctx context.Context, event PublishEventConfiguration) error { + args := m.Called(ctx, event) + return args.Error(0) +} + +func TestKafkaPubSubDataSource(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + kafkaCfg := &nodev1.KafkaEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Topics: []string{"test-topic"}, + } + + // Create the data source to test with a real adapter + adapter := &Adapter{} + pubsub := &PubSubDataSource{ + EventConfiguration: kafkaCfg, + KafkaAdapter: adapter, + } + + // Run the standard test suite + datasource.VerifyPubSubDataSourceImplementation(t, pubsub) +} + +// TestPubSubDataSourceWithMockAdapter tests the PubSubDataSource with a mocked adapter +func TestPubSubDataSourceWithMockAdapter(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + kafkaCfg := &nodev1.KafkaEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Topics: []string{"test-topic"}, + } + + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Configure mock expectations for Publish + mockAdapter.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishEventConfiguration) bool { + return event.ProviderID == "test-provider" && event.Topic == "test-topic" + })).Return(nil) + + // Create the data source with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: kafkaCfg, + KafkaAdapter: mockAdapter, + } + + // Get the data source + ds, err := pubsub.GetResolveDataSource() + require.NoError(t, err) + + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.NoError(t, err) + + // Call Load on the data source + out := &bytes.Buffer{} + err = ds.Load(context.Background(), []byte(input), out) + require.NoError(t, err) + require.Equal(t, `{"success": true}`, out.String()) + + // Verify mock expectations + mockAdapter.AssertExpectations(t) +} + +// TestKafkaPubSubDataSourceMultiTopicSubscription tests only the subscription functionality +// for multiple topics. The publish and resolve datasource tests are skipped since they +// do not support multiple topics. +func TestKafkaPubSubDataSourceMultiTopicSubscription(t *testing.T) { + // Create event configuration with multiple topics + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, // Must be PUBLISH as it's the only supported type + TypeName: "TestType", + FieldName: "testField", + } + + kafkaCfg := &nodev1.KafkaEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Topics: []string{"test-topic-1", "test-topic-2"}, + } + + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Set up expectations for subscribe with both topics + mockAdapter.On("Subscribe", mock.Anything, mock.MatchedBy(func(event SubscriptionEventConfiguration) bool { + return event.ProviderID == "test-provider" && + len(event.Topics) == 2 && + event.Topics[0] == "test-topic-1" && + event.Topics[1] == "test-topic-2" + }), mock.Anything).Return(nil) + + // Create the data source to test with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: kafkaCfg, + KafkaAdapter: mockAdapter, + } + + // Test GetEngineEventConfiguration + testConfig := pubsub.GetEngineEventConfiguration() + require.NotNil(t, testConfig, "Expected non-nil EngineEventConfiguration") + + // Test GetResolveDataSourceSubscription + subscription, err := pubsub.GetResolveDataSourceSubscription() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscription") + require.NotNil(t, subscription, "Expected non-nil SubscriptionDataSource") + + // Test GetResolveDataSourceSubscriptionInput + subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") + require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") + + // Verify the subscription input contains both topics + var subscriptionConfig SubscriptionEventConfiguration + err = json.Unmarshal([]byte(subscriptionInput), &subscriptionConfig) + require.NoError(t, err, "Expected valid JSON from GetResolveDataSourceSubscriptionInput") + require.Equal(t, 2, len(subscriptionConfig.Topics), "Expected 2 topics in subscription configuration") + require.Equal(t, "test-topic-1", subscriptionConfig.Topics[0], "Expected first topic to be 'test-topic-1'") + require.Equal(t, "test-topic-2", subscriptionConfig.Topics[1], "Expected second topic to be 'test-topic-2'") + +} + +// bytesBuffer is a helper that implements io.Writer to capture output +type bytesBuffer []byte + +func (b *bytesBuffer) Write(p []byte) (n int, err error) { + *b = append(*b, p...) + return len(p), nil +} diff --git a/router/pkg/pubsub/nats/adapter.go b/router/pkg/pubsub/nats/adapter.go index 679feeac82..bcdad2f3d1 100644 --- a/router/pkg/pubsub/nats/adapter.go +++ b/router/pkg/pubsub/nats/adapter.go @@ -16,6 +16,19 @@ import ( "go.uber.org/zap" ) +// AdapterInterface defines the methods that a NATS adapter should implement +type AdapterInterface interface { + // Subscribe subscribes to the given events and sends updates to the updater + Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error + // Publish publishes the given event to the specified subject + Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error + // Request sends a request to the specified subject and writes the response to the given writer + Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error + // Shutdown gracefully shuts down the adapter + Shutdown(ctx context.Context) error +} + +// Adapter implements the AdapterInterface for NATS pub/sub type Adapter struct { ctx context.Context client *LazyClient @@ -275,7 +288,7 @@ func (p *Adapter) Shutdown(ctx context.Context) error { return nil } -func NewAdapter(ctx context.Context, logger *zap.Logger, url string, opts []nats.Option, hostName string, routerListenAddr string) (*Adapter, error) { +func NewAdapter(ctx context.Context, logger *zap.Logger, url string, opts []nats.Option, hostName string, routerListenAddr string) (AdapterInterface, error) { if logger == nil { logger = zap.NewNop() } diff --git a/router/pkg/pubsub/nats/engine_datasource.go b/router/pkg/pubsub/nats/engine_datasource.go index ba7e0ae1e6..c4ef8e1e2f 100644 --- a/router/pkg/pubsub/nats/engine_datasource.go +++ b/router/pkg/pubsub/nats/engine_datasource.go @@ -46,7 +46,7 @@ func (s *PublishEventConfiguration) MarshalJSONTemplate() string { } type SubscriptionSource struct { - pubSub *Adapter + pubSub AdapterInterface } func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { @@ -81,7 +81,7 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r } type NatsPublishDataSource struct { - pubSub *Adapter + pubSub AdapterInterface } func (s *NatsPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { @@ -104,7 +104,7 @@ func (s *NatsPublishDataSource) LoadWithFiles(ctx context.Context, input []byte, } type NatsRequestDataSource struct { - pubSub *Adapter + pubSub AdapterInterface } func (s *NatsRequestDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { diff --git a/router/pkg/pubsub/nats/provider.go b/router/pkg/pubsub/nats/provider.go index 8e7815722f..22212c735b 100644 --- a/router/pkg/pubsub/nats/provider.go +++ b/router/pkg/pubsub/nats/provider.go @@ -91,7 +91,7 @@ func transformEventConfig(cfg *nodev1.NatsEventConfiguration, fn datasource.Argu type PubSubProvider struct { EventConfiguration []*nodev1.NatsEventConfiguration Logger *zap.Logger - Providers map[string]*Adapter + Providers map[string]AdapterInterface } func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, extractFn datasource.ArgumentTemplateCallback) (datasource.PubSubDataSource, error) { @@ -111,7 +111,7 @@ func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, } func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta *plan.DataSourceMetadata, config config.EventsConfiguration, logger *zap.Logger, hostName string, routerListenAddr string) (datasource.PubSubProvider, error) { - var providers map[string]*Adapter + var providers map[string]AdapterInterface if natsData := in.GetCustomEvents().GetNats(); natsData != nil { definedProviders := make(map[string]bool) for _, provider := range config.Providers.Nats { @@ -124,7 +124,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta } usedProviders[event.EngineEventConfiguration.ProviderId] = true } - providers = map[string]*Adapter{} + providers = map[string]AdapterInterface{} for _, provider := range config.Providers.Nats { if !usedProviders[provider.ID] { continue diff --git a/router/pkg/pubsub/nats/pubsub_datasource.go b/router/pkg/pubsub/nats/pubsub_datasource.go index b084cb9edc..5f3b37d702 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource.go +++ b/router/pkg/pubsub/nats/pubsub_datasource.go @@ -59,7 +59,7 @@ func NewLazyClient(url string, opts ...nats.Option) *LazyClient { type PubSubDataSource struct { EventConfiguration *nodev1.NatsEventConfiguration - NatsAdapter *Adapter + NatsAdapter AdapterInterface } func (c *PubSubDataSource) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { diff --git a/router/pkg/pubsub/nats/pubsub_datasource_test.go b/router/pkg/pubsub/nats/pubsub_datasource_test.go new file mode 100644 index 0000000000..3f1622cec2 --- /dev/null +++ b/router/pkg/pubsub/nats/pubsub_datasource_test.go @@ -0,0 +1,300 @@ +package nats + +import ( + "bytes" + "context" + "encoding/json" + "io" + "testing" + + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +// MockAdapter mocks the required functionality from the Adapter for testing +type MockAdapter struct { + mock.Mock +} + +// Ensure MockAdapter implements AdapterInterface +var _ AdapterInterface = (*MockAdapter)(nil) + +func (m *MockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { + args := m.Called(ctx, event, updater) + return args.Error(0) +} + +func (m *MockAdapter) Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error { + args := m.Called(ctx, event) + return args.Error(0) +} + +func (m *MockAdapter) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { + args := m.Called(ctx, event, w) + return args.Error(0) +} + +func (m *MockAdapter) Shutdown(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + +func TestNatsPubSubDataSource(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject"}, + } + + // Create the data source to test + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + } + + // Run the standard test suite + datasource.VerifyPubSubDataSourceImplementation(t, pubsub) +} + +func TestNatsPubSubDataSourceWithStreamConfiguration(t *testing.T) { + // Create event configuration with required fields and stream config + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + streamCfg := &nodev1.NatsStreamConfiguration{ + ConsumerName: "test-consumer", + StreamName: "test-stream", + ConsumerInactiveThreshold: 60, + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject"}, + StreamConfiguration: streamCfg, + } + + // Create the data source to test + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + } + + // Run the standard test suite + datasource.VerifyPubSubDataSourceImplementation(t, pubsub) +} + +func TestNatsPubSubDataSourceRequestType(t *testing.T) { + // Create event configuration with REQUEST type + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_REQUEST, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject"}, + } + + // Create the data source to test + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + } + + // Run the standard test suite + datasource.VerifyPubSubDataSourceImplementation(t, pubsub) +} + +func TestNatsPubSubDataSourceSubscribeType(t *testing.T) { + // Create event configuration with SUBSCRIBE type + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject"}, + } + + // Create the data source to test + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + } + + // Run the standard test suite + datasource.VerifyPubSubDataSourceImplementation(t, pubsub) +} + +// TestPubSubDataSourceWithMockAdapter tests the PubSubDataSource with a mocked adapter +func TestPubSubDataSourceWithMockAdapter(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject"}, + } + + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Configure mock expectations for Publish + mockAdapter.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { + return event.ProviderID == "test-provider" && event.Subject == "test.subject" + })).Return(nil) + + // Create the data source with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: mockAdapter, + } + + // Get the data source + ds, err := pubsub.GetResolveDataSource() + require.NoError(t, err) + + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.NoError(t, err) + + // Call Load on the data source + out := &bytes.Buffer{} + err = ds.Load(context.Background(), []byte(input), out) + require.NoError(t, err) + require.Equal(t, `{"success": true}`, out.String()) + + // Verify mock expectations + mockAdapter.AssertExpectations(t) +} + +// TestNatsPubSubDataSourceRequestWithMockAdapter tests the REQUEST functionality with a mocked adapter +func TestNatsPubSubDataSourceRequestWithMockAdapter(t *testing.T) { + // Create event configuration with REQUEST type + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_REQUEST, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject"}, + } + + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Configure mock expectations for Request + mockAdapter.On("Request", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { + return event.ProviderID == "test-provider" && event.Subject == "test.subject" + }), mock.Anything).Run(func(args mock.Arguments) { + // Simulate writing a response + w := args.Get(2).(io.Writer) + w.Write([]byte(`{"response":"data"}`)) + }).Return(nil) + + // Create the data source with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: mockAdapter, + } + + // Get the data source + ds, err := pubsub.GetResolveDataSource() + require.NoError(t, err) + + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.NoError(t, err) + + // Call Load on the data source + out := &bytes.Buffer{} + err = ds.Load(context.Background(), []byte(input), out) + require.NoError(t, err) + require.Equal(t, `{"response":"data"}`, out.String()) + + // Verify mock expectations + mockAdapter.AssertExpectations(t) +} + +// TestNatsPubSubDataSourceMultiSubjectSubscription tests the subscription functionality +// for multiple subjects with a mocked adapter +func TestNatsPubSubDataSourceMultiSubjectSubscription(t *testing.T) { + // Create event configuration with multiple subjects + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test.subject.1", "test.subject.2"}, + } + + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Set up expectations for subscribe with both subjects + mockAdapter.On("Subscribe", mock.Anything, mock.MatchedBy(func(event SubscriptionEventConfiguration) bool { + return event.ProviderID == "test-provider" && + len(event.Subjects) == 2 && + event.Subjects[0] == "test.subject.1" && + event.Subjects[1] == "test.subject.2" + }), mock.Anything).Return(nil) + + // Create the data source to test with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + NatsAdapter: mockAdapter, + } + + // Test GetEngineEventConfiguration + testConfig := pubsub.GetEngineEventConfiguration() + require.NotNil(t, testConfig, "Expected non-nil EngineEventConfiguration") + + // Test GetResolveDataSourceSubscription + subscription, err := pubsub.GetResolveDataSourceSubscription() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscription") + require.NotNil(t, subscription, "Expected non-nil SubscriptionDataSource") + + // Test GetResolveDataSourceSubscriptionInput + subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") + require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") + + // Verify the subscription input contains both subjects + var subscriptionConfig SubscriptionEventConfiguration + err = json.Unmarshal([]byte(subscriptionInput), &subscriptionConfig) + require.NoError(t, err, "Expected valid JSON from GetResolveDataSourceSubscriptionInput") + require.Equal(t, 2, len(subscriptionConfig.Subjects), "Expected 2 subjects in subscription configuration") + require.Equal(t, "test.subject.1", subscriptionConfig.Subjects[0], "Expected first subject to be 'test.subject.1'") + require.Equal(t, "test.subject.2", subscriptionConfig.Subjects[1], "Expected second subject to be 'test.subject.2'") +} From 3a023e266b282e894c67da8361cc88864c74b3b0 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 9 Apr 2025 19:07:17 +0200 Subject: [PATCH 19/49] chore: add github.com/stretchr/objx v0.5.2 dependency --- router/go.mod | 1 + router/go.sum | 1 + 2 files changed, 2 insertions(+) diff --git a/router/go.mod b/router/go.mod index c7e45704b8..2e06e70b34 100644 --- a/router/go.mod +++ b/router/go.mod @@ -125,6 +125,7 @@ require ( github.com/sergi/go-diff v1.3.1 // indirect github.com/shoenig/go-m1cpu v0.1.6 // indirect github.com/sirupsen/logrus v1.9.3 // indirect + github.com/stretchr/objx v0.5.2 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tklauser/go-sysconf v0.3.12 // indirect diff --git a/router/go.sum b/router/go.sum index fe8e6ee3f0..58c3f6a810 100644 --- a/router/go.sum +++ b/router/go.sum @@ -228,6 +228,7 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= From 2a96355caefc41c479e4213b6df60465cb70db72 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 9 Apr 2025 19:12:19 +0200 Subject: [PATCH 20/49] chore: remove unused bytesBuffer helper from Kafka pubsub tests --- router/pkg/pubsub/kafka/pubsub_datasource_test.go | 8 -------- 1 file changed, 8 deletions(-) diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go index a090e3ca87..abfa4ce03b 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource_test.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -160,11 +160,3 @@ func TestKafkaPubSubDataSourceMultiTopicSubscription(t *testing.T) { require.Equal(t, "test-topic-2", subscriptionConfig.Topics[1], "Expected second topic to be 'test-topic-2'") } - -// bytesBuffer is a helper that implements io.Writer to capture output -type bytesBuffer []byte - -func (b *bytesBuffer) Write(p []byte) (n int, err error) { - *b = append(*b, p...) - return len(p), nil -} From fe2f388dd44f7ee3c9d7a8f3288aacedfcef8f77 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Thu, 10 Apr 2025 16:44:07 +0200 Subject: [PATCH 21/49] chore: add tests --- .../pubsub/kafka/engine_datasource_test.go | 281 ++++++++++++ router/pkg/pubsub/kafka/provider_test.go | 206 +++++++++ .../pubsub/kafka/pubsub_datasource_test.go | 103 ++++- .../pkg/pubsub/nats/engine_datasource_test.go | 401 +++++++++++++++++ router/pkg/pubsub/nats/lazy_client.go | 53 +++ router/pkg/pubsub/nats/lazy_client_test.go | 413 ++++++++++++++++++ router/pkg/pubsub/nats/provider_test.go | 282 ++++++++++++ router/pkg/pubsub/nats/pubsub_datasource.go | 48 -- .../pkg/pubsub/nats/pubsub_datasource_test.go | 257 ++++++----- router/pkg/pubsub/nats/utils.go | 4 +- router/pkg/pubsub/nats/utils_test.go | 83 ++++ 11 files changed, 1940 insertions(+), 191 deletions(-) create mode 100644 router/pkg/pubsub/kafka/engine_datasource_test.go create mode 100644 router/pkg/pubsub/kafka/provider_test.go create mode 100644 router/pkg/pubsub/nats/engine_datasource_test.go create mode 100644 router/pkg/pubsub/nats/lazy_client.go create mode 100644 router/pkg/pubsub/nats/lazy_client_test.go create mode 100644 router/pkg/pubsub/nats/provider_test.go create mode 100644 router/pkg/pubsub/nats/utils_test.go diff --git a/router/pkg/pubsub/kafka/engine_datasource_test.go b/router/pkg/pubsub/kafka/engine_datasource_test.go new file mode 100644 index 0000000000..aa58bb939e --- /dev/null +++ b/router/pkg/pubsub/kafka/engine_datasource_test.go @@ -0,0 +1,281 @@ +package kafka + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "testing" + + "github.com/cespare/xxhash/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +// EngineDataSourceMockAdapter is a mock implementation of AdapterInterface for testing +type EngineDataSourceMockAdapter struct { + mock.Mock +} + +func (m *EngineDataSourceMockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { + args := m.Called(ctx, event, updater) + return args.Error(0) +} + +func (m *EngineDataSourceMockAdapter) Publish(ctx context.Context, event PublishEventConfiguration) error { + args := m.Called(ctx, event) + return args.Error(0) +} + +// MockSubscriptionUpdater implements resolve.SubscriptionUpdater +type MockSubscriptionUpdater struct { + mock.Mock +} + +func (m *MockSubscriptionUpdater) Update(data []byte) { + m.Called(data) +} + +func (m *MockSubscriptionUpdater) Done() { + m.Called() +} + +func TestPublishEventConfiguration_MarshalJSONTemplate(t *testing.T) { + tests := []struct { + name string + config PublishEventConfiguration + wantPattern string + }{ + { + name: "simple configuration", + config: PublishEventConfiguration{ + ProviderID: "test-provider", + Topic: "test-topic", + Data: json.RawMessage(`{"message":"hello"}`), + }, + wantPattern: `{"topic":"test-topic", "data": {"message":"hello"}, "providerId":"test-provider"}`, + }, + { + name: "with special characters", + config: PublishEventConfiguration{ + ProviderID: "test-provider-id", + Topic: "topic-with-hyphens", + Data: json.RawMessage(`{"message":"special \"quotes\" here"}`), + }, + wantPattern: `{"topic":"topic-with-hyphens", "data": {"message":"special \"quotes\" here"}, "providerId":"test-provider-id"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.config.MarshalJSONTemplate() + assert.Equal(t, tt.wantPattern, result) + }) + } +} + +func TestSubscriptionSource_UniqueRequestID(t *testing.T) { + tests := []struct { + name string + input string + expectError bool + expectedError error + }{ + { + name: "valid input", + input: `{"topics":["topic1", "topic2"], "providerId":"test-provider"}`, + expectError: false, + }, + { + name: "missing topics", + input: `{"providerId":"test-provider"}`, + expectError: true, + expectedError: errors.New("Key path not found"), + }, + { + name: "missing providerId", + input: `{"topics":["topic1", "topic2"]}`, + expectError: true, + expectedError: errors.New("Key path not found"), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + source := &SubscriptionSource{ + pubSub: &EngineDataSourceMockAdapter{}, + } + ctx := &resolve.Context{} + input := []byte(tt.input) + xxh := xxhash.New() + + err := source.UniqueRequestID(ctx, input, xxh) + + if tt.expectError { + require.Error(t, err) + if tt.expectedError != nil { + // For jsonparser errors, just check if the error message contains the expected text + assert.Contains(t, err.Error(), tt.expectedError.Error()) + } + } else { + require.NoError(t, err) + // Check that the hash has been updated + assert.NotEqual(t, 0, xxh.Sum64()) + } + }) + } +} + +func TestSubscriptionSource_Start(t *testing.T) { + tests := []struct { + name string + input string + mockSetup func(*EngineDataSourceMockAdapter) + expectError bool + }{ + { + name: "successful subscription", + input: `{"topics":["topic1", "topic2"], "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ + ProviderID: "test-provider", + Topics: []string{"topic1", "topic2"}, + }, mock.Anything).Return(nil) + }, + expectError: false, + }, + { + name: "adapter returns error", + input: `{"topics":["topic1"], "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ + ProviderID: "test-provider", + Topics: []string{"topic1"}, + }, mock.Anything).Return(errors.New("subscription error")) + }, + expectError: true, + }, + { + name: "invalid input json", + input: `{"invalid json":`, + mockSetup: func(m *EngineDataSourceMockAdapter) {}, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockAdapter := new(EngineDataSourceMockAdapter) + tt.mockSetup(mockAdapter) + + source := &SubscriptionSource{ + pubSub: mockAdapter, + } + + // Set up go context + goCtx := context.Background() + + // Create a resolve.Context with the standard context + resolveCtx := &resolve.Context{} + resolveCtx = resolveCtx.WithContext(goCtx) + + // Create a proper mock updater + updater := new(MockSubscriptionUpdater) + updater.On("Done").Return() + + input := []byte(tt.input) + err := source.Start(resolveCtx, input, updater) + + if tt.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + } + mockAdapter.AssertExpectations(t) + }) + } +} + +func TestKafkaPublishDataSource_Load(t *testing.T) { + tests := []struct { + name string + input string + mockSetup func(*EngineDataSourceMockAdapter) + expectError bool + expectedOutput string + expectPublished bool + }{ + { + name: "successful publish", + input: `{"topic":"test-topic", "data":{"message":"hello"}, "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishEventConfiguration) bool { + return event.ProviderID == "test-provider" && + event.Topic == "test-topic" && + string(event.Data) == `{"message":"hello"}` + })).Return(nil) + }, + expectError: false, + expectedOutput: `{"success": true}`, + expectPublished: true, + }, + { + name: "publish error", + input: `{"topic":"test-topic", "data":{"message":"hello"}, "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Publish", mock.Anything, mock.Anything).Return(errors.New("publish error")) + }, + expectError: false, // The Load method doesn't return the publish error directly + expectedOutput: `{"success": false}`, + expectPublished: true, + }, + { + name: "invalid input json", + input: `{"invalid json":`, + mockSetup: func(m *EngineDataSourceMockAdapter) {}, + expectError: true, + expectPublished: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockAdapter := new(EngineDataSourceMockAdapter) + tt.mockSetup(mockAdapter) + + dataSource := &KafkaPublishDataSource{ + pubSub: mockAdapter, + } + ctx := context.Background() + input := []byte(tt.input) + out := &bytes.Buffer{} + + err := dataSource.Load(ctx, input, out) + + if tt.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + assert.Equal(t, tt.expectedOutput, out.String()) + } + + if tt.expectPublished { + mockAdapter.AssertExpectations(t) + } + }) + } +} + +func TestKafkaPublishDataSource_LoadWithFiles(t *testing.T) { + t.Run("panic on not implemented", func(t *testing.T) { + dataSource := &KafkaPublishDataSource{ + pubSub: &EngineDataSourceMockAdapter{}, + } + + assert.Panics(t, func() { + dataSource.LoadWithFiles(context.Background(), nil, nil, &bytes.Buffer{}) + }) + }) +} diff --git a/router/pkg/pubsub/kafka/provider_test.go b/router/pkg/pubsub/kafka/provider_test.go new file mode 100644 index 0000000000..de54da2dfa --- /dev/null +++ b/router/pkg/pubsub/kafka/provider_test.go @@ -0,0 +1,206 @@ +package kafka + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" + "go.uber.org/zap" + "go.uber.org/zap/zaptest" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" +) + +// mockAdapter is a mock of AdapterInterface +type mockAdapter struct { + mock.Mock +} + +func (m *mockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { + args := m.Called(ctx, event, updater) + return args.Error(0) +} + +func (m *mockAdapter) Publish(ctx context.Context, event PublishEventConfiguration) error { + args := m.Called(ctx, event) + return args.Error(0) +} + +func TestBuildKafkaOptions(t *testing.T) { + t.Run("basic configuration", func(t *testing.T) { + cfg := config.KafkaEventSource{ + Brokers: []string{"localhost:9092"}, + } + + opts, err := buildKafkaOptions(cfg) + require.NoError(t, err) + require.NotEmpty(t, opts) + }) + + t.Run("with TLS", func(t *testing.T) { + enabled := true + cfg := config.KafkaEventSource{ + Brokers: []string{"localhost:9092"}, + TLS: &config.KafkaTLSConfiguration{ + Enabled: enabled, + }, + } + + opts, err := buildKafkaOptions(cfg) + require.NoError(t, err) + require.NotEmpty(t, opts) + // Can't directly check for TLS options, but we can verify more options are present + require.Equal(t, len(opts), 4) + }) + + t.Run("with auth", func(t *testing.T) { + username := "user" + password := "pass" + cfg := config.KafkaEventSource{ + Brokers: []string{"localhost:9092"}, + Authentication: &config.KafkaAuthentication{ + SASLPlain: config.KafkaSASLPlainAuthentication{ + Username: &username, + Password: &password, + }, + }, + } + + opts, err := buildKafkaOptions(cfg) + require.NoError(t, err) + require.NotEmpty(t, opts) + // Can't directly check for SASL options, but we can verify more options are present + require.Greater(t, len(opts), 1) + }) +} + +func TestGetProvider(t *testing.T) { + t.Run("returns nil if no Kafka configuration", func(t *testing.T) { + ctx := context.Background() + in := &nodev1.DataSourceConfiguration{ + CustomEvents: &nodev1.DataSourceCustomEvents{}, + } + + dsMeta := &plan.DataSourceMetadata{} + cfg := config.EventsConfiguration{} + logger := zaptest.NewLogger(t) + + provider, err := GetProvider(ctx, in, dsMeta, cfg, logger, "host", "addr") + require.NoError(t, err) + require.Nil(t, provider) + }) + + t.Run("errors if provider not found", func(t *testing.T) { + ctx := context.Background() + in := &nodev1.DataSourceConfiguration{ + CustomEvents: &nodev1.DataSourceCustomEvents{ + Kafka: []*nodev1.KafkaEventConfiguration{ + { + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + ProviderId: "unknown", + }, + }, + }, + }, + } + + dsMeta := &plan.DataSourceMetadata{} + cfg := config.EventsConfiguration{ + Providers: config.EventProviders{ + Kafka: []config.KafkaEventSource{ + {ID: "provider1", Brokers: []string{"localhost:9092"}}, + }, + }, + } + logger := zaptest.NewLogger(t) + + provider, err := GetProvider(ctx, in, dsMeta, cfg, logger, "host", "addr") + require.Error(t, err) + require.Nil(t, provider) + assert.Contains(t, err.Error(), "failed to find Kafka provider with ID") + }) + + t.Run("creates provider with configured adapters", func(t *testing.T) { + providerId := "test-provider" + + in := &nodev1.DataSourceConfiguration{ + CustomEvents: &nodev1.DataSourceCustomEvents{ + Kafka: []*nodev1.KafkaEventConfiguration{ + { + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + ProviderId: providerId, + }, + }, + }, + }, + } + + cfg := config.EventsConfiguration{ + Providers: config.EventProviders{ + Kafka: []config.KafkaEventSource{ + {ID: providerId, Brokers: []string{"localhost:9092"}}, + }, + }, + } + + logger := zaptest.NewLogger(t) + + // Create mock adapter for testing + provider, err := GetProvider(context.Background(), in, &plan.DataSourceMetadata{}, cfg, logger, "host", "addr") + require.NoError(t, err) + require.NotNil(t, provider) + + // Check the returned provider + kafkaProvider, ok := provider.(*PubSubProvider) + require.True(t, ok) + assert.NotNil(t, kafkaProvider.Logger) + assert.NotNil(t, kafkaProvider.Providers) + assert.Contains(t, kafkaProvider.Providers, providerId) + }) +} + +func TestPubSubProvider_FindPubSubDataSource(t *testing.T) { + mock := &mockAdapter{} + providerId := "test-provider" + typeName := "TestType" + fieldName := "testField" + + provider := &PubSubProvider{ + EventConfiguration: []*nodev1.KafkaEventConfiguration{ + { + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + TypeName: typeName, + FieldName: fieldName, + ProviderId: providerId, + }, + }, + }, + Logger: zap.NewNop(), + Providers: map[string]AdapterInterface{ + providerId: mock, + }, + } + + t.Run("find matching datasource", func(t *testing.T) { + ds, err := provider.FindPubSubDataSource(typeName, fieldName, nil) + require.NoError(t, err) + require.NotNil(t, ds) + + // Check the returned datasource + kafkaDs, ok := ds.(*PubSubDataSource) + require.True(t, ok) + assert.Equal(t, mock, kafkaDs.KafkaAdapter) + assert.Equal(t, provider.EventConfiguration[0], kafkaDs.EventConfiguration) + }) + + t.Run("return nil if no match", func(t *testing.T) { + ds, err := provider.FindPubSubDataSource("OtherType", fieldName, nil) + require.NoError(t, err) + require.Nil(t, ds) + }) +} diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go index abfa4ce03b..b97bbd034b 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource_test.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -103,6 +103,88 @@ func TestPubSubDataSourceWithMockAdapter(t *testing.T) { mockAdapter.AssertExpectations(t) } +// TestPubSubDataSource_GetResolveDataSource_WrongType tests the PubSubDataSource with a mocked adapter +func TestPubSubDataSource_GetResolveDataSource_WrongType(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_SUBSCRIBE, + TypeName: "TestType", + FieldName: "testField", + } + + kafkaCfg := &nodev1.KafkaEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Topics: []string{"test-topic"}, + } + + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Create the data source with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: kafkaCfg, + KafkaAdapter: mockAdapter, + } + + // Get the data source + ds, err := pubsub.GetResolveDataSource() + require.Error(t, err) + require.Nil(t, ds) +} + +// TestPubSubDataSource_GetResolveDataSourceInput_MultipleTopics tests the PubSubDataSource with a mocked adapter +func TestPubSubDataSource_GetResolveDataSourceInput_MultipleTopics(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + kafkaCfg := &nodev1.KafkaEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Topics: []string{"test-topic-1", "test-topic-2"}, + } + + // Create the data source with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: kafkaCfg, + } + + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.Error(t, err) + require.Empty(t, input) +} + +// TestPubSubDataSource_GetResolveDataSourceInput_NoTopics tests the PubSubDataSource with a mocked adapter +func TestPubSubDataSource_GetResolveDataSourceInput_NoTopics(t *testing.T) { + // Create event configuration with required fields + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + kafkaCfg := &nodev1.KafkaEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Topics: []string{}, + } + + // Create the data source with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: kafkaCfg, + } + + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.Error(t, err) + require.Empty(t, input) +} + // TestKafkaPubSubDataSourceMultiTopicSubscription tests only the subscription functionality // for multiple topics. The publish and resolve datasource tests are skipped since they // do not support multiple topics. @@ -120,32 +202,11 @@ func TestKafkaPubSubDataSourceMultiTopicSubscription(t *testing.T) { Topics: []string{"test-topic-1", "test-topic-2"}, } - // Create mock adapter - mockAdapter := new(MockAdapter) - - // Set up expectations for subscribe with both topics - mockAdapter.On("Subscribe", mock.Anything, mock.MatchedBy(func(event SubscriptionEventConfiguration) bool { - return event.ProviderID == "test-provider" && - len(event.Topics) == 2 && - event.Topics[0] == "test-topic-1" && - event.Topics[1] == "test-topic-2" - }), mock.Anything).Return(nil) - // Create the data source to test with mock adapter pubsub := &PubSubDataSource{ EventConfiguration: kafkaCfg, - KafkaAdapter: mockAdapter, } - // Test GetEngineEventConfiguration - testConfig := pubsub.GetEngineEventConfiguration() - require.NotNil(t, testConfig, "Expected non-nil EngineEventConfiguration") - - // Test GetResolveDataSourceSubscription - subscription, err := pubsub.GetResolveDataSourceSubscription() - require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscription") - require.NotNil(t, subscription, "Expected non-nil SubscriptionDataSource") - // Test GetResolveDataSourceSubscriptionInput subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") diff --git a/router/pkg/pubsub/nats/engine_datasource_test.go b/router/pkg/pubsub/nats/engine_datasource_test.go new file mode 100644 index 0000000000..4fb1ae767d --- /dev/null +++ b/router/pkg/pubsub/nats/engine_datasource_test.go @@ -0,0 +1,401 @@ +package nats + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "testing" + + "github.com/cespare/xxhash/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +// EngineDataSourceMockAdapter is a mock implementation of AdapterInterface for testing +type EngineDataSourceMockAdapter struct { + mock.Mock +} + +func (m *EngineDataSourceMockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { + args := m.Called(ctx, event, updater) + return args.Error(0) +} + +func (m *EngineDataSourceMockAdapter) Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error { + args := m.Called(ctx, event) + return args.Error(0) +} + +func (m *EngineDataSourceMockAdapter) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { + args := m.Called(ctx, event, w) + return args.Error(0) +} + +func (m *EngineDataSourceMockAdapter) Shutdown(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + +// MockSubscriptionUpdater implements resolve.SubscriptionUpdater +type MockSubscriptionUpdater struct { + mock.Mock +} + +func (m *MockSubscriptionUpdater) Update(data []byte) { + m.Called(data) +} + +func (m *MockSubscriptionUpdater) Done() { + m.Called() +} + +func TestPublishEventConfiguration_MarshalJSONTemplate(t *testing.T) { + tests := []struct { + name string + config PublishEventConfiguration + wantPattern string + }{ + { + name: "simple configuration", + config: PublishEventConfiguration{ + ProviderID: "test-provider", + Subject: "test-subject", + Data: json.RawMessage(`{"message":"hello"}`), + }, + wantPattern: `{"subject":"test-subject", "data": {"message":"hello"}, "providerId":"test-provider"}`, + }, + { + name: "with special characters", + config: PublishEventConfiguration{ + ProviderID: "test-provider-id", + Subject: "subject-with-hyphens", + Data: json.RawMessage(`{"message":"special \"quotes\" here"}`), + }, + wantPattern: `{"subject":"subject-with-hyphens", "data": {"message":"special \"quotes\" here"}, "providerId":"test-provider-id"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.config.MarshalJSONTemplate() + assert.Equal(t, tt.wantPattern, result) + }) + } +} + +func TestPublishAndRequestEventConfiguration_MarshalJSONTemplate(t *testing.T) { + tests := []struct { + name string + config PublishAndRequestEventConfiguration + wantPattern string + }{ + { + name: "simple configuration", + config: PublishAndRequestEventConfiguration{ + ProviderID: "test-provider", + Subject: "test-subject", + Data: json.RawMessage(`{"message":"hello"}`), + }, + wantPattern: `{"subject":"test-subject", "data": {"message":"hello"}, "providerId":"test-provider"}`, + }, + { + name: "with special characters", + config: PublishAndRequestEventConfiguration{ + ProviderID: "test-provider-id", + Subject: "subject-with-hyphens", + Data: json.RawMessage(`{"message":"special \"quotes\" here"}`), + }, + wantPattern: `{"subject":"subject-with-hyphens", "data": {"message":"special \"quotes\" here"}, "providerId":"test-provider-id"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.config.MarshalJSONTemplate() + assert.Equal(t, tt.wantPattern, result) + }) + } +} + +func TestSubscriptionSource_UniqueRequestID(t *testing.T) { + tests := []struct { + name string + input string + expectError bool + expectedError error + }{ + { + name: "valid input", + input: `{"subjects":["subject1", "subject2"], "providerId":"test-provider"}`, + expectError: false, + }, + { + name: "missing subjects", + input: `{"providerId":"test-provider"}`, + expectError: true, + expectedError: errors.New("Key path not found"), + }, + { + name: "missing providerId", + input: `{"subjects":["subject1", "subject2"]}`, + expectError: true, + expectedError: errors.New("Key path not found"), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + source := &SubscriptionSource{ + pubSub: &EngineDataSourceMockAdapter{}, + } + ctx := &resolve.Context{} + input := []byte(tt.input) + xxh := xxhash.New() + + err := source.UniqueRequestID(ctx, input, xxh) + + if tt.expectError { + require.Error(t, err) + if tt.expectedError != nil { + // For jsonparser errors, just check if the error message contains the expected text + assert.Contains(t, err.Error(), tt.expectedError.Error()) + } + } else { + require.NoError(t, err) + // Check that the hash has been updated + assert.NotEqual(t, 0, xxh.Sum64()) + } + }) + } +} + +func TestSubscriptionSource_Start(t *testing.T) { + tests := []struct { + name string + input string + mockSetup func(*EngineDataSourceMockAdapter) + expectError bool + }{ + { + name: "successful subscription", + input: `{"subjects":["subject1", "subject2"], "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ + ProviderID: "test-provider", + Subjects: []string{"subject1", "subject2"}, + }, mock.Anything).Return(nil) + }, + expectError: false, + }, + { + name: "adapter returns error", + input: `{"subjects":["subject1"], "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ + ProviderID: "test-provider", + Subjects: []string{"subject1"}, + }, mock.Anything).Return(errors.New("subscription error")) + }, + expectError: true, + }, + { + name: "invalid input json", + input: `{"invalid json":`, + mockSetup: func(m *EngineDataSourceMockAdapter) {}, + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockAdapter := new(EngineDataSourceMockAdapter) + tt.mockSetup(mockAdapter) + + source := &SubscriptionSource{ + pubSub: mockAdapter, + } + + // Set up go context + goCtx := context.Background() + + // Create a resolve.Context with the standard context + resolveCtx := &resolve.Context{} + resolveCtx = resolveCtx.WithContext(goCtx) + + // Create a proper mock updater + updater := new(MockSubscriptionUpdater) + updater.On("Done").Return() + + input := []byte(tt.input) + err := source.Start(resolveCtx, input, updater) + + if tt.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + } + mockAdapter.AssertExpectations(t) + }) + } +} + +func TestNatsPublishDataSource_Load(t *testing.T) { + tests := []struct { + name string + input string + mockSetup func(*EngineDataSourceMockAdapter) + expectError bool + expectedOutput string + expectPublished bool + }{ + { + name: "successful publish", + input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { + return event.ProviderID == "test-provider" && + event.Subject == "test-subject" && + string(event.Data) == `{"message":"hello"}` + })).Return(nil) + }, + expectError: false, + expectedOutput: `{"success": true}`, + expectPublished: true, + }, + { + name: "publish error", + input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Publish", mock.Anything, mock.Anything).Return(errors.New("publish error")) + }, + expectError: false, // The Load method doesn't return the publish error directly + expectedOutput: `{"success": false}`, + expectPublished: true, + }, + { + name: "invalid input json", + input: `{"invalid json":`, + mockSetup: func(m *EngineDataSourceMockAdapter) {}, + expectError: true, + expectPublished: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockAdapter := new(EngineDataSourceMockAdapter) + tt.mockSetup(mockAdapter) + + dataSource := &NatsPublishDataSource{ + pubSub: mockAdapter, + } + + ctx := context.Background() + input := []byte(tt.input) + var out bytes.Buffer + + err := dataSource.Load(ctx, input, &out) + + if tt.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + if tt.expectPublished { + mockAdapter.AssertExpectations(t) + } + if tt.expectedOutput != "" { + assert.Equal(t, tt.expectedOutput, out.String()) + } + } + }) + } +} + +func TestNatsPublishDataSource_LoadWithFiles(t *testing.T) { + dataSource := &NatsPublishDataSource{} + assert.Panics(t, func() { + dataSource.LoadWithFiles(context.Background(), []byte{}, nil, &bytes.Buffer{}) + }, "Expected LoadWithFiles to panic with 'not implemented'") +} + +func TestNatsRequestDataSource_Load(t *testing.T) { + tests := []struct { + name string + input string + mockSetup func(*EngineDataSourceMockAdapter) + expectError bool + expectedOutput string + }{ + { + name: "successful request", + input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Request", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { + return event.ProviderID == "test-provider" && + event.Subject == "test-subject" && + string(event.Data) == `{"message":"hello"}` + }), mock.Anything).Run(func(args mock.Arguments) { + // Write response to the output buffer + w := args.Get(2).(io.Writer) + _, _ = w.Write([]byte(`{"response":"success"}`)) + }).Return(nil) + }, + expectError: false, + expectedOutput: `{"response":"success"}`, + }, + { + name: "request error", + input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, + mockSetup: func(m *EngineDataSourceMockAdapter) { + m.On("Request", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("request error")) + }, + expectError: true, + expectedOutput: "", + }, + { + name: "invalid input json", + input: `{"invalid json":`, + mockSetup: func(m *EngineDataSourceMockAdapter) {}, + expectError: true, + expectedOutput: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mockAdapter := new(EngineDataSourceMockAdapter) + tt.mockSetup(mockAdapter) + + dataSource := &NatsRequestDataSource{ + pubSub: mockAdapter, + } + + ctx := context.Background() + input := []byte(tt.input) + var out bytes.Buffer + + err := dataSource.Load(ctx, input, &out) + + if tt.expectError { + require.Error(t, err) + } else { + require.NoError(t, err) + mockAdapter.AssertExpectations(t) + if tt.expectedOutput != "" { + assert.Equal(t, tt.expectedOutput, out.String()) + } + } + }) + } +} + +func TestNatsRequestDataSource_LoadWithFiles(t *testing.T) { + dataSource := &NatsRequestDataSource{} + assert.Panics(t, func() { + dataSource.LoadWithFiles(context.Background(), []byte{}, nil, &bytes.Buffer{}) + }, "Expected LoadWithFiles to panic with 'not implemented'") +} diff --git a/router/pkg/pubsub/nats/lazy_client.go b/router/pkg/pubsub/nats/lazy_client.go new file mode 100644 index 0000000000..a5abc82bda --- /dev/null +++ b/router/pkg/pubsub/nats/lazy_client.go @@ -0,0 +1,53 @@ +package nats + +import ( + "sync" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" +) + +type LazyClient struct { + once sync.Once + url string + opts []nats.Option + client *nats.Conn + js jetstream.JetStream + err error +} + +func (c *LazyClient) Connect(opts ...nats.Option) error { + c.once.Do(func() { + c.client, c.err = nats.Connect(c.url, opts...) + if c.err != nil { + return + } + c.js, c.err = jetstream.New(c.client) + }) + return c.err +} + +func (c *LazyClient) GetClient() (*nats.Conn, error) { + if c.client == nil { + if err := c.Connect(c.opts...); err != nil { + return nil, err + } + } + return c.client, c.err +} + +func (c *LazyClient) GetJetStream() (jetstream.JetStream, error) { + if c.js == nil { + if err := c.Connect(c.opts...); err != nil { + return nil, err + } + } + return c.js, c.err +} + +func NewLazyClient(url string, opts ...nats.Option) *LazyClient { + return &LazyClient{ + url: url, + opts: opts, + } +} diff --git a/router/pkg/pubsub/nats/lazy_client_test.go b/router/pkg/pubsub/nats/lazy_client_test.go new file mode 100644 index 0000000000..03b63b27a8 --- /dev/null +++ b/router/pkg/pubsub/nats/lazy_client_test.go @@ -0,0 +1,413 @@ +package nats + +import ( + "errors" + "sync" + "testing" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Define interfaces for our testing +type natsConnector interface { + Connect(url string, opts ...nats.Option) (*nats.Conn, error) +} + +type jetStreamCreator interface { + Create(nc *nats.Conn) (jetstream.JetStream, error) +} + +// mockNatsConnector implements natsConnector for testing +type mockNatsConnector struct { + mockConn *nats.Conn + mockErr error + callCount int + lastOpts []nats.Option +} + +func (m *mockNatsConnector) Connect(url string, opts ...nats.Option) (*nats.Conn, error) { + m.callCount++ + m.lastOpts = opts + return m.mockConn, m.mockErr +} + +// threadSafeConnector is a specialized version for thread safety testing +type threadSafeConnector struct { + mockConn *nats.Conn + mockErr error + callCount int + countMutex sync.Mutex + sleepTime time.Duration +} + +func (t *threadSafeConnector) Connect(url string, opts ...nats.Option) (*nats.Conn, error) { + // Simulate a slow connection + time.Sleep(t.sleepTime) + t.countMutex.Lock() + t.callCount++ + t.countMutex.Unlock() + return t.mockConn, t.mockErr +} + +// mockJetStreamCreator implements jetStreamCreator for testing +type mockJetStreamCreator struct { + mockJS jetstream.JetStream + mockErr error + callCount int +} + +func (m *mockJetStreamCreator) Create(nc *nats.Conn) (jetstream.JetStream, error) { + m.callCount++ + return m.mockJS, m.mockErr +} + +// testLazyClient is a wrapper for LazyClient that allows us to inject mocks +type testLazyClient struct { + *LazyClient + connector natsConnector + jetStreamCreator jetStreamCreator +} + +// newTestLazyClient creates a LazyClient with mocked dependencies +func newTestLazyClient(url string, connector natsConnector, creator jetStreamCreator, opts ...nats.Option) *testLazyClient { + c := &testLazyClient{ + LazyClient: NewLazyClient(url, opts...), + connector: connector, + jetStreamCreator: creator, + } + + return c +} + +// Connect overrides LazyClient.Connect to use our mocks +func (c *testLazyClient) Connect(opts ...nats.Option) error { + c.once.Do(func() { + // If no options are provided, use the ones stored during initialization + optionsToUse := opts + if len(optionsToUse) == 0 { + optionsToUse = c.opts + } + c.client, c.err = c.connector.Connect(c.url, optionsToUse...) + if c.err != nil { + return + } + c.js, c.err = c.jetStreamCreator.Create(c.client) + }) + return c.err +} + +// GetClient overrides LazyClient.GetClient to ensure we use our Connect method +func (c *testLazyClient) GetClient() (*nats.Conn, error) { + if c.client == nil { + if err := c.Connect(c.opts...); err != nil { + return nil, err + } + } + return c.client, c.err +} + +// GetJetStream overrides LazyClient.GetJetStream to ensure we use our Connect method +func (c *testLazyClient) GetJetStream() (jetstream.JetStream, error) { + if c.js == nil { + if err := c.Connect(c.opts...); err != nil { + return nil, err + } + } + return c.js, c.err +} + +func TestNewLazyClient(t *testing.T) { + url := "nats://localhost:4222" + opts := []nats.Option{nats.Name("test-client")} + + client := NewLazyClient(url, opts...) + + assert.Equal(t, url, client.url) + assert.Equal(t, opts, client.opts) + assert.Nil(t, client.client) + assert.Nil(t, client.js) + assert.Nil(t, client.err) +} + +func TestLazyClient_Connect_Success(t *testing.T) { + // Create mocks + connector := &mockNatsConnector{ + mockConn: &nats.Conn{}, + } + + jsCreator := &mockJetStreamCreator{ + mockJS: &struct{ jetstream.JetStream }{}, + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // First call should connect + err := client.Connect() + require.NoError(t, err) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") + assert.Equal(t, connector.mockConn, client.client) + assert.Equal(t, jsCreator.mockJS, client.js) + + // Second call should not connect again due to sync.Once + err = client.Connect() + require.NoError(t, err) + assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should still be called only once") +} + +func TestLazyClient_Connect_Error(t *testing.T) { + // Create mocks with connection error + connector := &mockNatsConnector{ + mockErr: errors.New("mock connect error"), + } + + jsCreator := &mockJetStreamCreator{ + mockJS: &struct{ jetstream.JetStream }{}, + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // Connect should return an error + err := client.Connect() + require.Error(t, err) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should not be called") + assert.Nil(t, client.client) + assert.Nil(t, client.js) + assert.NotNil(t, client.err) + assert.Equal(t, "mock connect error", client.err.Error()) + + // Second call should not connect again and return the same error + connector.mockErr = errors.New("different error") // Should be ignored due to sync.Once + err2 := client.Connect() + require.Error(t, err2) + assert.Equal(t, err, err2, "Should return the same error") + assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") + assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should still not be called") +} + +func TestLazyClient_Connect_JetStreamError(t *testing.T) { + // Create mocks with jetstream error + connector := &mockNatsConnector{ + mockConn: &nats.Conn{}, + } + + jsCreator := &mockJetStreamCreator{ + mockErr: errors.New("mock jetstream error"), + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // Connect should return the jetstream error + err := client.Connect() + require.Error(t, err) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") + assert.NotNil(t, client.client) // NATS connection was successful + assert.Nil(t, client.js) // JetStream failed + assert.NotNil(t, client.err) + assert.Equal(t, "mock jetstream error", client.err.Error()) + + // Second call should not connect again and return the same error + jsCreator.mockErr = errors.New("different error") // Should be ignored due to sync.Once + err2 := client.Connect() + require.Error(t, err2) + assert.Equal(t, err, err2, "Should return the same error") + assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should still be called only once") +} + +func TestLazyClient_GetClient(t *testing.T) { + t.Run("with successful connection", func(t *testing.T) { + // Create mocks + connector := &mockNatsConnector{ + mockConn: &nats.Conn{}, + } + + jsCreator := &mockJetStreamCreator{ + mockJS: &struct{ jetstream.JetStream }{}, + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // First call to GetClient should connect + conn, err := client.GetClient() + require.NoError(t, err) + assert.Equal(t, connector.mockConn, conn) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + + // Second call should not connect again + conn2, err := client.GetClient() + require.NoError(t, err) + assert.Equal(t, conn, conn2) + assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") + }) + + t.Run("with failed connection", func(t *testing.T) { + // Create mocks with connection error + connector := &mockNatsConnector{ + mockErr: errors.New("mock connect error"), + } + + jsCreator := &mockJetStreamCreator{} + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // GetClient should attempt to connect and return an error + conn, err := client.GetClient() + require.Error(t, err) + assert.Nil(t, conn) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should not be called") + assert.Equal(t, "mock connect error", err.Error()) + }) +} + +func TestLazyClient_GetJetStream(t *testing.T) { + t.Run("with successful connection", func(t *testing.T) { + // Create mocks + connector := &mockNatsConnector{ + mockConn: &nats.Conn{}, + } + + jsCreator := &mockJetStreamCreator{ + mockJS: &struct{ jetstream.JetStream }{}, + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // First call to GetJetStream should connect + js, err := client.GetJetStream() + require.NoError(t, err) + assert.Equal(t, jsCreator.mockJS, js) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") + + // Second call should not connect again + js2, err := client.GetJetStream() + require.NoError(t, err) + assert.Equal(t, js, js2) + assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should still be called only once") + }) + + t.Run("with failed connection", func(t *testing.T) { + // Create mocks with connection error + connector := &mockNatsConnector{ + mockErr: errors.New("mock connect error"), + } + + jsCreator := &mockJetStreamCreator{} + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // GetJetStream should attempt to connect and return an error + js, err := client.GetJetStream() + require.Error(t, err) + assert.Nil(t, js) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should not be called") + assert.Equal(t, "mock connect error", err.Error()) + }) + + t.Run("with jetstream creation failure", func(t *testing.T) { + // Create mocks with jetstream error + connector := &mockNatsConnector{ + mockConn: &nats.Conn{}, + } + + jsCreator := &mockJetStreamCreator{ + mockErr: errors.New("mock jetstream error"), + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // GetJetStream should create the connection but fail on jetstream + js, err := client.GetJetStream() + require.Error(t, err) + assert.Nil(t, js) + assert.Equal(t, 1, connector.callCount, "Connect should be called once") + assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") + assert.Equal(t, "mock jetstream error", err.Error()) + }) +} + +func TestLazyClient_WithOptions(t *testing.T) { + // Create mocks that will track options + connector := &mockNatsConnector{ + mockConn: &nats.Conn{}, + } + + jsCreator := &mockJetStreamCreator{ + mockJS: &struct{ jetstream.JetStream }{}, + } + + // Create options that we can verify are passed to Connect + option1 := nats.Name("test-client") + option2 := nats.NoEcho() + + // Create test client with options + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator, option1, option2) + + // Call Connect to trigger the connection + err := client.Connect() + require.NoError(t, err) + + // Verify that the options were passed - they're stored in connector.lastOpts + require.Len(t, connector.lastOpts, 2, "Options should be passed to Connect") + + // Now we can call GetClient() and verify it reuses the connection + prevCallCount := connector.callCount + _, err = client.GetClient() + require.NoError(t, err) + + // Verify Connect wasn't called again + assert.Equal(t, prevCallCount, connector.callCount, "GetClient should reuse the existing connection") +} + +func TestLazyClient_ThreadSafety(t *testing.T) { + // This test checks that LazyClient's sync.Once protection works as expected + // by calling Connect from multiple goroutines simultaneously + + // Create a thread-safe connector that simulates slow connections + connector := &threadSafeConnector{ + mockConn: &nats.Conn{}, + sleepTime: 10 * time.Millisecond, + } + + jsCreator := &mockJetStreamCreator{ + mockJS: &struct{ jetstream.JetStream }{}, + } + + // Create test client + client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) + + // Launch multiple goroutines that call Connect simultaneously + var wg sync.WaitGroup + for i := 0; i < 10; i++ { + wg.Add(1) + go func() { + defer wg.Done() + err := client.Connect() + assert.NoError(t, err) + }() + } + + wg.Wait() + + // Connect should have been called exactly once + assert.Equal(t, 1, connector.callCount, "Connect should have been called exactly once") +} diff --git a/router/pkg/pubsub/nats/provider_test.go b/router/pkg/pubsub/nats/provider_test.go new file mode 100644 index 0000000000..fd2255bf60 --- /dev/null +++ b/router/pkg/pubsub/nats/provider_test.go @@ -0,0 +1,282 @@ +package nats + +import ( + "context" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" + "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" + "go.uber.org/zap" + "go.uber.org/zap/zaptest" +) + +// mockAdapter is a mock of AdapterInterface +type mockAdapter struct { + mock.Mock +} + +func (m *mockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { + args := m.Called(ctx, event, updater) + return args.Error(0) +} + +func (m *mockAdapter) Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error { + args := m.Called(ctx, event) + return args.Error(0) +} + +func (m *mockAdapter) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { + args := m.Called(ctx, event, w) + return args.Error(0) +} + +func (m *mockAdapter) Shutdown(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + +func TestBuildNatsOptions(t *testing.T) { + t.Run("basic configuration", func(t *testing.T) { + cfg := config.NatsEventSource{ + ID: "test-nats", + URL: "nats://localhost:4222", + } + logger := zaptest.NewLogger(t) + + opts, err := buildNatsOptions(cfg, logger) + require.NoError(t, err) + require.NotEmpty(t, opts) + }) + + t.Run("with token authentication", func(t *testing.T) { + token := "test-token" + cfg := config.NatsEventSource{ + ID: "test-nats", + URL: "nats://localhost:4222", + Authentication: &config.NatsAuthentication{ + NatsTokenBasedAuthentication: config.NatsTokenBasedAuthentication{ + Token: &token, + }, + }, + } + logger := zaptest.NewLogger(t) + + opts, err := buildNatsOptions(cfg, logger) + require.NoError(t, err) + require.NotEmpty(t, opts) + // Can't directly check for token options, but we can verify options are present + require.Greater(t, len(opts), 7) // Basic options (7) + token option + }) + + t.Run("with user/password authentication", func(t *testing.T) { + username := "user" + password := "pass" + cfg := config.NatsEventSource{ + ID: "test-nats", + URL: "nats://localhost:4222", + Authentication: &config.NatsAuthentication{ + UserInfo: config.NatsCredentialsAuthentication{ + Username: &username, + Password: &password, + }, + }, + } + logger := zaptest.NewLogger(t) + + opts, err := buildNatsOptions(cfg, logger) + require.NoError(t, err) + require.NotEmpty(t, opts) + // Can't directly check for auth options, but we can verify options are present + require.Greater(t, len(opts), 7) // Basic options (7) + user info option + }) +} + +func TestTransformEventConfig(t *testing.T) { + t.Run("publish event", func(t *testing.T) { + cfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + Type: nodev1.EventType_PUBLISH, + }, + Subjects: []string{"original.subject"}, + } + + // Simple transform function that adds "transformed." prefix + transformFn := func(s string) (string, error) { + return "transformed." + s, nil + } + + transformedCfg, err := transformEventConfig(cfg, transformFn) + require.NoError(t, err) + require.Equal(t, []string{"transformed.original.subject"}, transformedCfg.Subjects) + }) + + t.Run("subscribe event", func(t *testing.T) { + cfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + Type: nodev1.EventType_SUBSCRIBE, + }, + Subjects: []string{"original.subject1", "original.subject2"}, + } + + // Simple transform function that adds "transformed." prefix + transformFn := func(s string) (string, error) { + return "transformed." + s, nil + } + + transformedCfg, err := transformEventConfig(cfg, transformFn) + require.NoError(t, err) + // Since the function sorts the subjects + require.Equal(t, []string{"transformed.original.subject1", "transformed.original.subject2"}, transformedCfg.Subjects) + }) + + t.Run("invalid subject", func(t *testing.T) { + cfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + Type: nodev1.EventType_PUBLISH, + }, + Subjects: []string{"invalid subject with spaces"}, + } + + transformFn := func(s string) (string, error) { + return s, nil + } + + _, err := transformEventConfig(cfg, transformFn) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid subject") + }) +} + +func TestGetProvider(t *testing.T) { + t.Run("returns nil if no NATS configuration", func(t *testing.T) { + ctx := context.Background() + in := &nodev1.DataSourceConfiguration{ + CustomEvents: &nodev1.DataSourceCustomEvents{}, + } + + dsMeta := &plan.DataSourceMetadata{} + cfg := config.EventsConfiguration{} + logger := zaptest.NewLogger(t) + + provider, err := GetProvider(ctx, in, dsMeta, cfg, logger, "host", "addr") + require.NoError(t, err) + require.Nil(t, provider) + }) + + t.Run("errors if provider not found", func(t *testing.T) { + ctx := context.Background() + in := &nodev1.DataSourceConfiguration{ + CustomEvents: &nodev1.DataSourceCustomEvents{ + Nats: []*nodev1.NatsEventConfiguration{ + { + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + ProviderId: "unknown", + }, + }, + }, + }, + } + + dsMeta := &plan.DataSourceMetadata{} + cfg := config.EventsConfiguration{ + Providers: config.EventProviders{ + Nats: []config.NatsEventSource{ + {ID: "provider1", URL: "nats://localhost:4222"}, + }, + }, + } + logger := zaptest.NewLogger(t) + + provider, err := GetProvider(ctx, in, dsMeta, cfg, logger, "host", "addr") + require.Error(t, err) + require.Nil(t, provider) + assert.Contains(t, err.Error(), "failed to find Nats provider with ID") + }) +} + +func TestPubSubProvider_FindPubSubDataSource(t *testing.T) { + mockNats := &mockAdapter{} + providerId := "test-provider" + typeName := "TestType" + fieldName := "testField" + + provider := &PubSubProvider{ + EventConfiguration: []*nodev1.NatsEventConfiguration{ + { + EngineEventConfiguration: &nodev1.EngineEventConfiguration{ + TypeName: typeName, + FieldName: fieldName, + ProviderId: providerId, + Type: nodev1.EventType_PUBLISH, + }, + Subjects: []string{"test.subject"}, + }, + }, + Logger: zap.NewNop(), + Providers: map[string]AdapterInterface{ + providerId: mockNats, + }, + } + + t.Run("find matching datasource", func(t *testing.T) { + // Identity transform function + transformFn := func(s string) (string, error) { + return s, nil + } + + ds, err := provider.FindPubSubDataSource(typeName, fieldName, transformFn) + require.NoError(t, err) + require.NotNil(t, ds) + + // Check the returned datasource + natsDs, ok := ds.(*PubSubDataSource) + require.True(t, ok) + assert.Equal(t, mockNats, natsDs.NatsAdapter) + assert.Equal(t, provider.EventConfiguration[0], natsDs.EventConfiguration) + }) + + t.Run("return nil if no match", func(t *testing.T) { + ds, err := provider.FindPubSubDataSource("OtherType", fieldName, nil) + require.NoError(t, err) + require.Nil(t, ds) + }) + + t.Run("handle error in transform function", func(t *testing.T) { + // Function that returns error + errorFn := func(s string) (string, error) { + return "", assert.AnError + } + + ds, err := provider.FindPubSubDataSource(typeName, fieldName, errorFn) + require.Error(t, err) + require.Nil(t, ds) + }) + + t.Run("handle error in transform function", func(t *testing.T) { + // Function that returns error + errorFn := func(s string) (string, error) { + return "", assert.AnError + } + + ds, err := provider.FindPubSubDataSource(typeName, fieldName, errorFn) + require.Error(t, err) + require.Nil(t, ds) + }) + + t.Run("handle error in transform function with invalid subject", func(t *testing.T) { + // Function that returns error + errorFn := func(s string) (string, error) { + return " ", nil + } + + ds, err := provider.FindPubSubDataSource(typeName, fieldName, errorFn) + require.Error(t, err) + require.Nil(t, ds) + }) +} diff --git a/router/pkg/pubsub/nats/pubsub_datasource.go b/router/pkg/pubsub/nats/pubsub_datasource.go index 5f3b37d702..07a2d37542 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource.go +++ b/router/pkg/pubsub/nats/pubsub_datasource.go @@ -3,60 +3,12 @@ package nats import ( "encoding/json" "fmt" - "sync" - "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type LazyClient struct { - once sync.Once - url string - opts []nats.Option - client *nats.Conn - js jetstream.JetStream - err error -} - -func (c *LazyClient) Connect(opts ...nats.Option) error { - c.once.Do(func() { - c.client, c.err = nats.Connect(c.url, opts...) - if c.err != nil { - return - } - c.js, c.err = jetstream.New(c.client) - }) - return c.err -} - -func (c *LazyClient) GetClient() (*nats.Conn, error) { - if c.client == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.client, c.err -} - -func (c *LazyClient) GetJetStream() (jetstream.JetStream, error) { - if c.js == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.js, c.err -} - -func NewLazyClient(url string, opts ...nats.Option) *LazyClient { - return &LazyClient{ - url: url, - opts: opts, - } -} - type PubSubDataSource struct { EventConfiguration *nodev1.NatsEventConfiguration NatsAdapter AdapterInterface diff --git a/router/pkg/pubsub/nats/pubsub_datasource_test.go b/router/pkg/pubsub/nats/pubsub_datasource_test.go index 3f1622cec2..43fb75fef8 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource_test.go +++ b/router/pkg/pubsub/nats/pubsub_datasource_test.go @@ -53,21 +53,22 @@ func TestNatsPubSubDataSource(t *testing.T) { natsCfg := &nodev1.NatsEventConfiguration{ EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject"}, + Subjects: []string{"test-subject"}, } - // Create the data source to test + // Create the data source to test with a real adapter + adapter := &Adapter{} pubsub := &PubSubDataSource{ EventConfiguration: natsCfg, - NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + NatsAdapter: adapter, } // Run the standard test suite datasource.VerifyPubSubDataSourceImplementation(t, pubsub) } -func TestNatsPubSubDataSourceWithStreamConfiguration(t *testing.T) { - // Create event configuration with required fields and stream config +func TestPubSubDataSourceWithMockAdapter(t *testing.T) { + // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ ProviderId: "test-provider", Type: nodev1.EventType_PUBLISH, @@ -75,54 +76,74 @@ func TestNatsPubSubDataSourceWithStreamConfiguration(t *testing.T) { FieldName: "testField", } - streamCfg := &nodev1.NatsStreamConfiguration{ - ConsumerName: "test-consumer", - StreamName: "test-stream", - ConsumerInactiveThreshold: 60, - } - natsCfg := &nodev1.NatsEventConfiguration{ EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject"}, - StreamConfiguration: streamCfg, + Subjects: []string{"test-subject"}, } - // Create the data source to test + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Configure mock expectations for Publish + mockAdapter.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { + return event.ProviderID == "test-provider" && event.Subject == "test-subject" + })).Return(nil) + + // Create the data source with mock adapter pubsub := &PubSubDataSource{ EventConfiguration: natsCfg, - NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + NatsAdapter: mockAdapter, } - // Run the standard test suite - datasource.VerifyPubSubDataSourceImplementation(t, pubsub) + // Get the data source + ds, err := pubsub.GetResolveDataSource() + require.NoError(t, err) + + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.NoError(t, err) + + // Call Load on the data source + out := &bytes.Buffer{} + err = ds.Load(context.Background(), []byte(input), out) + require.NoError(t, err) + require.Equal(t, `{"success": true}`, out.String()) + + // Verify mock expectations + mockAdapter.AssertExpectations(t) } -func TestNatsPubSubDataSourceRequestType(t *testing.T) { - // Create event configuration with REQUEST type +func TestPubSubDataSource_GetResolveDataSource_WrongType(t *testing.T) { + // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ ProviderId: "test-provider", - Type: nodev1.EventType_REQUEST, + Type: nodev1.EventType_SUBSCRIBE, // This is not supported TypeName: "TestType", FieldName: "testField", } natsCfg := &nodev1.NatsEventConfiguration{ EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject"}, + Subjects: []string{"test-subject"}, } - // Create the data source to test + // Create mock adapter + mockAdapter := new(MockAdapter) + + // Create the data source with mock adapter pubsub := &PubSubDataSource{ EventConfiguration: natsCfg, - NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test + NatsAdapter: mockAdapter, } - // Run the standard test suite - datasource.VerifyPubSubDataSourceImplementation(t, pubsub) + // Get the data source + ds, err := pubsub.GetResolveDataSource() + require.Error(t, err) + require.Nil(t, ds) } -func TestNatsPubSubDataSourceSubscribeType(t *testing.T) { - // Create event configuration with SUBSCRIBE type +func TestPubSubDataSource_GetResolveDataSourceInput_MultipleSubjects(t *testing.T) { + // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ ProviderId: "test-provider", Type: nodev1.EventType_PUBLISH, @@ -132,21 +153,21 @@ func TestNatsPubSubDataSourceSubscribeType(t *testing.T) { natsCfg := &nodev1.NatsEventConfiguration{ EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject"}, + Subjects: []string{"test-subject-1", "test-subject-2"}, } - // Create the data source to test + // Create the data source with mock adapter pubsub := &PubSubDataSource{ EventConfiguration: natsCfg, - NatsAdapter: &Adapter{}, // Using a real Adapter type but with nil values for the test } - // Run the standard test suite - datasource.VerifyPubSubDataSourceImplementation(t, pubsub) + // Get the input + input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + require.Error(t, err) + require.Empty(t, input) } -// TestPubSubDataSourceWithMockAdapter tests the PubSubDataSource with a mocked adapter -func TestPubSubDataSourceWithMockAdapter(t *testing.T) { +func TestPubSubDataSource_GetResolveDataSourceInput_NoSubjects(t *testing.T) { // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ ProviderId: "test-provider", @@ -157,43 +178,93 @@ func TestPubSubDataSourceWithMockAdapter(t *testing.T) { natsCfg := &nodev1.NatsEventConfiguration{ EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject"}, + Subjects: []string{}, } - // Create mock adapter - mockAdapter := new(MockAdapter) - - // Configure mock expectations for Publish - mockAdapter.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { - return event.ProviderID == "test-provider" && event.Subject == "test.subject" - })).Return(nil) - // Create the data source with mock adapter pubsub := &PubSubDataSource{ EventConfiguration: natsCfg, - NatsAdapter: mockAdapter, } - // Get the data source - ds, err := pubsub.GetResolveDataSource() - require.NoError(t, err) - // Get the input input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) - require.NoError(t, err) + require.Error(t, err) + require.Empty(t, input) +} - // Call Load on the data source - out := &bytes.Buffer{} - err = ds.Load(context.Background(), []byte(input), out) - require.NoError(t, err) - require.Equal(t, `{"success": true}`, out.String()) +func TestNatsPubSubDataSourceMultiSubjectSubscription(t *testing.T) { + // Create event configuration with multiple subjects + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, // Must be PUBLISH as it's the only supported type + TypeName: "TestType", + FieldName: "testField", + } - // Verify mock expectations - mockAdapter.AssertExpectations(t) + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test-subject-1", "test-subject-2"}, + } + + // Create the data source to test with mock adapter + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + } + + // Test GetResolveDataSourceSubscriptionInput + subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") + require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") + + // Verify the subscription input contains both subjects + var subscriptionConfig SubscriptionEventConfiguration + err = json.Unmarshal([]byte(subscriptionInput), &subscriptionConfig) + require.NoError(t, err, "Expected valid JSON from GetResolveDataSourceSubscriptionInput") + require.Equal(t, 2, len(subscriptionConfig.Subjects), "Expected 2 subjects in subscription configuration") + require.Equal(t, "test-subject-1", subscriptionConfig.Subjects[0], "Expected first subject to be 'test-subject-1'") + require.Equal(t, "test-subject-2", subscriptionConfig.Subjects[1], "Expected second subject to be 'test-subject-2'") +} + +func TestNatsPubSubDataSourceWithStreamConfiguration(t *testing.T) { + // Create event configuration with stream configuration + engineEventConfig := &nodev1.EngineEventConfiguration{ + ProviderId: "test-provider", + Type: nodev1.EventType_PUBLISH, + TypeName: "TestType", + FieldName: "testField", + } + + natsCfg := &nodev1.NatsEventConfiguration{ + EngineEventConfiguration: engineEventConfig, + Subjects: []string{"test-subject"}, + StreamConfiguration: &nodev1.NatsStreamConfiguration{ + StreamName: "test-stream", + ConsumerName: "test-consumer", + ConsumerInactiveThreshold: 30, + }, + } + + // Create the data source to test + pubsub := &PubSubDataSource{ + EventConfiguration: natsCfg, + } + + // Test GetResolveDataSourceSubscriptionInput with stream configuration + subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") + require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") + + // Verify the subscription input contains stream configuration + var subscriptionConfig SubscriptionEventConfiguration + err = json.Unmarshal([]byte(subscriptionInput), &subscriptionConfig) + require.NoError(t, err, "Expected valid JSON from GetResolveDataSourceSubscriptionInput") + require.NotNil(t, subscriptionConfig.StreamConfiguration, "Expected non-nil stream configuration") + require.Equal(t, "test-consumer", subscriptionConfig.StreamConfiguration.Consumer, "Expected consumer to be 'test-consumer'") + require.Equal(t, "test-stream", subscriptionConfig.StreamConfiguration.StreamName, "Expected stream name to be 'test-stream'") + require.Equal(t, int32(30), subscriptionConfig.StreamConfiguration.ConsumerInactiveThreshold, "Expected consumer inactive threshold to be 30") } -// TestNatsPubSubDataSourceRequestWithMockAdapter tests the REQUEST functionality with a mocked adapter -func TestNatsPubSubDataSourceRequestWithMockAdapter(t *testing.T) { +func TestPubSubDataSource_RequestDataSource(t *testing.T) { // Create event configuration with REQUEST type engineEventConfig := &nodev1.EngineEventConfiguration{ ProviderId: "test-provider", @@ -204,7 +275,7 @@ func TestNatsPubSubDataSourceRequestWithMockAdapter(t *testing.T) { natsCfg := &nodev1.NatsEventConfiguration{ EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject"}, + Subjects: []string{"test-subject"}, } // Create mock adapter @@ -212,12 +283,11 @@ func TestNatsPubSubDataSourceRequestWithMockAdapter(t *testing.T) { // Configure mock expectations for Request mockAdapter.On("Request", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { - return event.ProviderID == "test-provider" && event.Subject == "test.subject" - }), mock.Anything).Run(func(args mock.Arguments) { - // Simulate writing a response + return event.ProviderID == "test-provider" && event.Subject == "test-subject" + }), mock.Anything).Return(nil).Run(func(args mock.Arguments) { w := args.Get(2).(io.Writer) - w.Write([]byte(`{"response":"data"}`)) - }).Return(nil) + w.Write([]byte(`{"response": "test"}`)) + }) // Create the data source with mock adapter pubsub := &PubSubDataSource{ @@ -228,6 +298,7 @@ func TestNatsPubSubDataSourceRequestWithMockAdapter(t *testing.T) { // Get the data source ds, err := pubsub.GetResolveDataSource() require.NoError(t, err) + require.NotNil(t, ds) // Get the input input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) @@ -237,64 +308,8 @@ func TestNatsPubSubDataSourceRequestWithMockAdapter(t *testing.T) { out := &bytes.Buffer{} err = ds.Load(context.Background(), []byte(input), out) require.NoError(t, err) - require.Equal(t, `{"response":"data"}`, out.String()) + require.Equal(t, `{"response": "test"}`, out.String()) // Verify mock expectations mockAdapter.AssertExpectations(t) } - -// TestNatsPubSubDataSourceMultiSubjectSubscription tests the subscription functionality -// for multiple subjects with a mocked adapter -func TestNatsPubSubDataSourceMultiSubjectSubscription(t *testing.T) { - // Create event configuration with multiple subjects - engineEventConfig := &nodev1.EngineEventConfiguration{ - ProviderId: "test-provider", - Type: nodev1.EventType_PUBLISH, - TypeName: "TestType", - FieldName: "testField", - } - - natsCfg := &nodev1.NatsEventConfiguration{ - EngineEventConfiguration: engineEventConfig, - Subjects: []string{"test.subject.1", "test.subject.2"}, - } - - // Create mock adapter - mockAdapter := new(MockAdapter) - - // Set up expectations for subscribe with both subjects - mockAdapter.On("Subscribe", mock.Anything, mock.MatchedBy(func(event SubscriptionEventConfiguration) bool { - return event.ProviderID == "test-provider" && - len(event.Subjects) == 2 && - event.Subjects[0] == "test.subject.1" && - event.Subjects[1] == "test.subject.2" - }), mock.Anything).Return(nil) - - // Create the data source to test with mock adapter - pubsub := &PubSubDataSource{ - EventConfiguration: natsCfg, - NatsAdapter: mockAdapter, - } - - // Test GetEngineEventConfiguration - testConfig := pubsub.GetEngineEventConfiguration() - require.NotNil(t, testConfig, "Expected non-nil EngineEventConfiguration") - - // Test GetResolveDataSourceSubscription - subscription, err := pubsub.GetResolveDataSourceSubscription() - require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscription") - require.NotNil(t, subscription, "Expected non-nil SubscriptionDataSource") - - // Test GetResolveDataSourceSubscriptionInput - subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() - require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") - require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") - - // Verify the subscription input contains both subjects - var subscriptionConfig SubscriptionEventConfiguration - err = json.Unmarshal([]byte(subscriptionInput), &subscriptionConfig) - require.NoError(t, err, "Expected valid JSON from GetResolveDataSourceSubscriptionInput") - require.Equal(t, 2, len(subscriptionConfig.Subjects), "Expected 2 subjects in subscription configuration") - require.Equal(t, "test.subject.1", subscriptionConfig.Subjects[0], "Expected first subject to be 'test.subject.1'") - require.Equal(t, "test.subject.2", subscriptionConfig.Subjects[1], "Expected second subject to be 'test.subject.2'") -} diff --git a/router/pkg/pubsub/nats/utils.go b/router/pkg/pubsub/nats/utils.go index 010b666c61..0229f7a41c 100644 --- a/router/pkg/pubsub/nats/utils.go +++ b/router/pkg/pubsub/nats/utils.go @@ -1,6 +1,8 @@ package nats -import "strings" +import ( + "strings" +) const ( fwc = '>' diff --git a/router/pkg/pubsub/nats/utils_test.go b/router/pkg/pubsub/nats/utils_test.go new file mode 100644 index 0000000000..9ec92c78cc --- /dev/null +++ b/router/pkg/pubsub/nats/utils_test.go @@ -0,0 +1,83 @@ +package nats + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestIsValidNatsSubject(t *testing.T) { + tests := []struct { + name string + subject string + want bool + }{ + { + name: "empty string", + subject: "", + want: false, + }, + { + name: "simple valid subject", + subject: "test.subject", + want: true, + }, + { + name: "valid subject with wildcard", + subject: "test.>", + want: true, + }, + { + name: "invalid with space", + subject: "test subject", + want: false, + }, + { + name: "invalid with tab", + subject: "test\tsubject", + want: false, + }, + { + name: "invalid with newline", + subject: "test\nsubject", + want: false, + }, + { + name: "invalid with empty token", + subject: "test..subject", + want: false, + }, + { + name: "wildcard not at end", + subject: "test.>.subject", + want: false, + }, + { + name: "contains a space", + subject: " ", + want: false, + }, + { + name: "contains a tab", + subject: "\t", + want: false, + }, + { + name: "contains a newline", + subject: "\n", + want: false, + }, + { + name: "contains a form feed", + subject: "\f", + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := isValidNatsSubject(tt.subject) + assert.Equal(t, tt.want, got) + }) + } +} From d4ed48ab47f573e2a8ac18deb177c5cd391f7b00 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 11 Apr 2025 14:59:17 +0200 Subject: [PATCH 22/49] chore: cleaned a bit the implementation, add a README to make it easier to add a PubSubDataSource --- router/core/executor.go | 2 +- router/core/factoryresolver.go | 6 +- router/core/plan_generator.go | 5 +- router/pkg/pubsub/README.md | 72 +++++++++++++++++++ router/pkg/pubsub/kafka/engine_datasource.go | 28 ++------ .../pubsub/kafka/engine_datasource_test.go | 8 +-- router/pkg/pubsub/kafka/provider.go | 4 +- router/pkg/pubsub/kafka/pubsub_datasource.go | 19 ++++- router/pkg/pubsub/nats/engine_datasource.go | 33 --------- router/pkg/pubsub/nats/provider.go | 4 +- router/pkg/pubsub/nats/pubsub_datasource.go | 32 +++++++++ 11 files changed, 143 insertions(+), 70 deletions(-) create mode 100644 router/pkg/pubsub/README.md diff --git a/router/core/executor.go b/router/core/executor.go index 13c8d350ee..63073e06aa 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -198,7 +198,7 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con // the plan config is what the engine uses to turn a GraphQL Request into an execution plan // the plan config is stateful as it carries connection pools and other things - loader := NewLoader(b.trackUsageInfo, NewDefaultFactoryResolver( + loader := NewLoader(ctx, b.trackUsageInfo, NewDefaultFactoryResolver( ctx, b.transportOptions, b.subscriptionClientOptions, diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 30633f004a..e12d64b5df 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -27,6 +27,7 @@ import ( ) type Loader struct { + ctx context.Context resolver FactoryResolver // includeInfo controls whether additional information like type usage and field usage is included in the plan de includeInfo bool @@ -158,8 +159,9 @@ func (d *DefaultFactoryResolver) InstanceData() InstanceData { return d.instanceData } -func NewLoader(includeInfo bool, resolver FactoryResolver, logger *zap.Logger) *Loader { +func NewLoader(ctx context.Context, includeInfo bool, resolver FactoryResolver, logger *zap.Logger) *Loader { return &Loader{ + ctx: ctx, resolver: resolver, includeInfo: includeInfo, logger: logger, @@ -413,7 +415,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod var err error out, err = pubsub.GetDataSourceFromConfig( - context.Background(), + l.ctx, in, l.dataSourceMetaData(in), routerEngineConfig.Events, diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index 43e8e344b6..c30f108a23 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -275,8 +275,9 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo graphql_datasource.WithNetPollConfiguration(netPollConfig), ) - loader := NewLoader(false, &DefaultFactoryResolver{ - engineCtx: context.Background(), + ctx := context.Background() + loader := NewLoader(ctx, false, &DefaultFactoryResolver{ + engineCtx: ctx, httpClient: http.DefaultClient, streamingClient: http.DefaultClient, subscriptionClient: subscriptionClient, diff --git a/router/pkg/pubsub/README.md b/router/pkg/pubsub/README.md new file mode 100644 index 0000000000..76e9efcbc2 --- /dev/null +++ b/router/pkg/pubsub/README.md @@ -0,0 +1,72 @@ + +# How to add a PubSub Provider + +## Add the data to the router proto + +You need to change the [router proto](../../../proto/wg/cosmo/node/v1/node.proto) as follows. + +Add the provider configuration like the `KafkaEventConfiguration` and then add it as repeated inside the `DataSourceCustomEvents`. + +The fields of `KafkaEventConfiguration` will depends on the provider. If the providers uses as grouping mechanisms of the messages "channel" it will be called "channels", if it is "Topic" it will be "topics", and so on. + +After this you will have to compile the proto launching from the main folder the command `make generate-go`. + + +## Build the PubSub Provider + +To build a PubSub provider you need to implement 4 things: +- `Adapter` +- `ProviderFactory` +- `PubSubProvider` +- `PubSubDataSource` + +And then add it inside the `GetProviderFactories` function. + +### Adapter + +The Adapter contains the logic that is actually calling the provider, usually it implement an interface as follows: + +```go +type AdapterInterface interface { + Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error + Publish(ctx context.Context, event PublishEventConfiguration) error +} +``` + +The content of `SubscriptionEventConfiguration` and `PublishEventConfiguration` depends on the provider, you can see an example of them in the [kafka implementation](./kafka/pubsub_datasource.go). + + +### ProviderFactory + +The `ProviderFactory` is the initial contact point where you receive: +- `ctx context.Context`, usually passed down to the adapter +- [`*nodev1.DataSourceConfiguration`](../../gen/proto/wg/cosmo/node/v1/node.pb.go#DataSourceConfiguration), that contains everything you need about the provider data parsed from the schema. +- `*plan.DataSourceMetadata`, usually not needed +- [`config.EventsConfiguration`](../config/config.go#EventsConfiguration) that contains the config needed to setup the provider connection +- `*zap.Logger` +- `hostName string`, useful if you need to identify the connection based on the local host name +- `routerListenAddr string`, useful if you need to identify the connection based on different router instances in the same host + +The responsability of the factory is to initialize the PubSubProvider, like in this implementation for an `ExampleProvider`: + +You can see as an example of the `GetProvider` function in the [kafka implementation](./kafka/provider.go). + +### PubSubProvider + +So, the `PubSubProvider` has already the Adapter of the provider initialized, and it will be called on a `Visitor.EnterField` call from the engine to check if the `PubSubProvider` is matching any `EngineEventConfiguration`. + +The responsability of the `PubSubProvider` is to match the `EngineEventConfiguration` and initialize a `PubSubDataSource` with the matching event and the provider `Adapter`. + +You can see as an example of the `PubSubProvider` in the [kafka implementation](./kafka/provider.go). + +### PubSubDataSource + +The `[PubSubDataSource](./datasource/datasource.go)` is the junction between the engine `resolve.DataSource` and the Provider that we are implementing. + +You can see an example in [kafka `PubSubDataSource`](./kafka/pubsub_datasource.go). + +To complete the `PubSubDataSource` implementation you should also add the engine data source. + +So you have to implement the SubscriptionDataSource, a structure that implements all the methods needed by the interface `resolve.SubscriptionDataSource`, like the [kafka implementation](./kafka/engine_datasource.go). + +And also, you have to implement the DataSource, a structure that implements all the methods needed by the interface `resolve.DataSource`, like `PublishDataSource` in the [kafka implementation](./kafka/pubsub_datasource.go). \ No newline at end of file diff --git a/router/pkg/pubsub/kafka/engine_datasource.go b/router/pkg/pubsub/kafka/engine_datasource.go index f342469f52..f31758ffa7 100644 --- a/router/pkg/pubsub/kafka/engine_datasource.go +++ b/router/pkg/pubsub/kafka/engine_datasource.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "encoding/json" - "fmt" "io" "github.com/buger/jsonparser" @@ -13,26 +12,11 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type SubscriptionEventConfiguration struct { - ProviderID string `json:"providerId"` - Topics []string `json:"topics"` -} - -type PublishEventConfiguration struct { - ProviderID string `json:"providerId"` - Topic string `json:"topic"` - Data json.RawMessage `json:"data"` -} - -func (s *PublishEventConfiguration) MarshalJSONTemplate() string { - return fmt.Sprintf(`{"topic":"%s", "data": %s, "providerId":"%s"}`, s.Topic, s.Data, s.ProviderID) -} - -type SubscriptionSource struct { +type SubscriptionDataSource struct { pubSub AdapterInterface } -func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { +func (s *SubscriptionDataSource) UniqueRequestID(ctx *resolve.Context, input []byte, xxh *xxhash.Digest) error { val, _, _, err := jsonparser.Get(input, "topics") if err != nil { return err @@ -52,7 +36,7 @@ func (s *SubscriptionSource) UniqueRequestID(ctx *resolve.Context, input []byte, return err } -func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater resolve.SubscriptionUpdater) error { +func (s *SubscriptionDataSource) Start(ctx *resolve.Context, input []byte, updater resolve.SubscriptionUpdater) error { var subscriptionConfiguration SubscriptionEventConfiguration err := json.Unmarshal(input, &subscriptionConfiguration) if err != nil { @@ -62,11 +46,11 @@ func (s *SubscriptionSource) Start(ctx *resolve.Context, input []byte, updater r return s.pubSub.Subscribe(ctx.Context(), subscriptionConfiguration, updater) } -type KafkaPublishDataSource struct { +type PublishDataSource struct { pubSub AdapterInterface } -func (s *KafkaPublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { +func (s *PublishDataSource) Load(ctx context.Context, input []byte, out *bytes.Buffer) error { var publishConfiguration PublishEventConfiguration err := json.Unmarshal(input, &publishConfiguration) if err != nil { @@ -81,6 +65,6 @@ func (s *KafkaPublishDataSource) Load(ctx context.Context, input []byte, out *by return err } -func (s *KafkaPublishDataSource) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { +func (s *PublishDataSource) LoadWithFiles(ctx context.Context, input []byte, files []*httpclient.FileUpload, out *bytes.Buffer) (err error) { panic("not implemented") } diff --git a/router/pkg/pubsub/kafka/engine_datasource_test.go b/router/pkg/pubsub/kafka/engine_datasource_test.go index aa58bb939e..ca9065c7df 100644 --- a/router/pkg/pubsub/kafka/engine_datasource_test.go +++ b/router/pkg/pubsub/kafka/engine_datasource_test.go @@ -104,7 +104,7 @@ func TestSubscriptionSource_UniqueRequestID(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - source := &SubscriptionSource{ + source := &SubscriptionDataSource{ pubSub: &EngineDataSourceMockAdapter{}, } ctx := &resolve.Context{} @@ -170,7 +170,7 @@ func TestSubscriptionSource_Start(t *testing.T) { mockAdapter := new(EngineDataSourceMockAdapter) tt.mockSetup(mockAdapter) - source := &SubscriptionSource{ + source := &SubscriptionDataSource{ pubSub: mockAdapter, } @@ -245,7 +245,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { mockAdapter := new(EngineDataSourceMockAdapter) tt.mockSetup(mockAdapter) - dataSource := &KafkaPublishDataSource{ + dataSource := &PublishDataSource{ pubSub: mockAdapter, } ctx := context.Background() @@ -270,7 +270,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { func TestKafkaPublishDataSource_LoadWithFiles(t *testing.T) { t.Run("panic on not implemented", func(t *testing.T) { - dataSource := &KafkaPublishDataSource{ + dataSource := &PublishDataSource{ pubSub: &EngineDataSourceMockAdapter{}, } diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go index ff767a7de0..1e86916355 100644 --- a/router/pkg/pubsub/kafka/provider.go +++ b/router/pkg/pubsub/kafka/provider.go @@ -52,7 +52,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta } usedProviders := make(map[string]bool) if kafkaData := in.GetCustomEvents().GetKafka(); kafkaData != nil { - for _, event := range in.CustomEvents.GetKafka() { + for _, event := range kafkaData { if !definedProviders[event.EngineEventConfiguration.ProviderId] { return nil, fmt.Errorf("failed to find Kafka provider with ID %s", event.EngineEventConfiguration.ProviderId) } @@ -75,7 +75,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta } return &PubSubProvider{ - EventConfiguration: in.CustomEvents.GetKafka(), + EventConfiguration: kafkaData, Logger: logger, Providers: providers, }, nil diff --git a/router/pkg/pubsub/kafka/pubsub_datasource.go b/router/pkg/pubsub/kafka/pubsub_datasource.go index fac984411e..f559faa372 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource.go @@ -23,7 +23,7 @@ func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { typeName := c.EventConfiguration.GetEngineEventConfiguration().GetType() switch typeName { case nodev1.EventType_PUBLISH: - dataSource = &KafkaPublishDataSource{ + dataSource = &PublishDataSource{ pubSub: c.KafkaAdapter, } default: @@ -54,7 +54,7 @@ func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, erro } func (c *PubSubDataSource) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { - return &SubscriptionSource{ + return &SubscriptionDataSource{ pubSub: c.KafkaAdapter, }, nil } @@ -75,3 +75,18 @@ func (c *PubSubDataSource) GetResolveDataSourceSubscriptionInput() (string, erro func (c *PubSubDataSource) GetProviderId() string { return c.EventConfiguration.GetEngineEventConfiguration().GetProviderId() } + +type SubscriptionEventConfiguration struct { + ProviderID string `json:"providerId"` + Topics []string `json:"topics"` +} + +type PublishEventConfiguration struct { + ProviderID string `json:"providerId"` + Topic string `json:"topic"` + Data json.RawMessage `json:"data"` +} + +func (s *PublishEventConfiguration) MarshalJSONTemplate() string { + return fmt.Sprintf(`{"topic":"%s", "data": %s, "providerId":"%s"}`, s.Topic, s.Data, s.ProviderID) +} diff --git a/router/pkg/pubsub/nats/engine_datasource.go b/router/pkg/pubsub/nats/engine_datasource.go index c4ef8e1e2f..c63da7ad0f 100644 --- a/router/pkg/pubsub/nats/engine_datasource.go +++ b/router/pkg/pubsub/nats/engine_datasource.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "encoding/json" - "fmt" "io" "github.com/buger/jsonparser" @@ -13,38 +12,6 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -type StreamConfiguration struct { - Consumer string `json:"consumer"` - ConsumerInactiveThreshold int32 `json:"consumerInactiveThreshold"` - StreamName string `json:"streamName"` -} - -type SubscriptionEventConfiguration struct { - ProviderID string `json:"providerId"` - Subjects []string `json:"subjects"` - StreamConfiguration *StreamConfiguration `json:"streamConfiguration,omitempty"` -} - -type PublishAndRequestEventConfiguration struct { - ProviderID string `json:"providerId"` - Subject string `json:"subject"` - Data json.RawMessage `json:"data"` -} - -func (s *PublishAndRequestEventConfiguration) MarshalJSONTemplate() string { - return fmt.Sprintf(`{"subject":"%s", "data": %s, "providerId":"%s"}`, s.Subject, s.Data, s.ProviderID) -} - -type PublishEventConfiguration struct { - ProviderID string `json:"providerId"` - Subject string `json:"subject"` - Data json.RawMessage `json:"data"` -} - -func (s *PublishEventConfiguration) MarshalJSONTemplate() string { - return fmt.Sprintf(`{"subject":"%s", "data": %s, "providerId":"%s"}`, s.Subject, s.Data, s.ProviderID) -} - type SubscriptionSource struct { pubSub AdapterInterface } diff --git a/router/pkg/pubsub/nats/provider.go b/router/pkg/pubsub/nats/provider.go index 22212c735b..e94c3a227d 100644 --- a/router/pkg/pubsub/nats/provider.go +++ b/router/pkg/pubsub/nats/provider.go @@ -118,7 +118,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta definedProviders[provider.ID] = true } usedProviders := make(map[string]bool) - for _, event := range in.CustomEvents.GetNats() { + for _, event := range natsData { if _, found := definedProviders[event.EngineEventConfiguration.ProviderId]; !found { return nil, fmt.Errorf("failed to find Nats provider with ID %s", event.EngineEventConfiguration.ProviderId) } @@ -141,7 +141,7 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta providers[provider.ID] = adapter } return &PubSubProvider{ - EventConfiguration: in.CustomEvents.GetNats(), + EventConfiguration: natsData, Logger: logger, Providers: providers, }, nil diff --git a/router/pkg/pubsub/nats/pubsub_datasource.go b/router/pkg/pubsub/nats/pubsub_datasource.go index 07a2d37542..ca8aae2506 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource.go +++ b/router/pkg/pubsub/nats/pubsub_datasource.go @@ -88,3 +88,35 @@ func (c *PubSubDataSource) GetResolveDataSourceSubscriptionInput() (string, erro func (c *PubSubDataSource) GetProviderId() string { return c.EventConfiguration.GetEngineEventConfiguration().GetProviderId() } + +type StreamConfiguration struct { + Consumer string `json:"consumer"` + ConsumerInactiveThreshold int32 `json:"consumerInactiveThreshold"` + StreamName string `json:"streamName"` +} + +type SubscriptionEventConfiguration struct { + ProviderID string `json:"providerId"` + Subjects []string `json:"subjects"` + StreamConfiguration *StreamConfiguration `json:"streamConfiguration,omitempty"` +} + +type PublishAndRequestEventConfiguration struct { + ProviderID string `json:"providerId"` + Subject string `json:"subject"` + Data json.RawMessage `json:"data"` +} + +func (s *PublishAndRequestEventConfiguration) MarshalJSONTemplate() string { + return fmt.Sprintf(`{"subject":"%s", "data": %s, "providerId":"%s"}`, s.Subject, s.Data, s.ProviderID) +} + +type PublishEventConfiguration struct { + ProviderID string `json:"providerId"` + Subject string `json:"subject"` + Data json.RawMessage `json:"data"` +} + +func (s *PublishEventConfiguration) MarshalJSONTemplate() string { + return fmt.Sprintf(`{"subject":"%s", "data": %s, "providerId":"%s"}`, s.Subject, s.Data, s.ProviderID) +} From 592cbb60539581a1592b649555695199904d6a63 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 11 Apr 2025 19:41:45 +0200 Subject: [PATCH 23/49] feat: implement PubSub provider startup and shutdown on graphServer --- router/core/executor.go | 5 +- router/core/factoryresolver.go | 52 +++++++++++++------ router/core/graph_server.go | 25 +++++++++ router/core/plan_generator.go | 3 +- router/pkg/pubsub/datasource/provider.go | 2 + router/pkg/pubsub/kafka/adapter.go | 1 + .../pubsub/kafka/engine_datasource_test.go | 5 ++ router/pkg/pubsub/kafka/provider.go | 13 +++++ router/pkg/pubsub/kafka/provider_test.go | 5 ++ .../pubsub/kafka/pubsub_datasource_test.go | 5 ++ router/pkg/pubsub/nats/provider.go | 13 +++++ router/pkg/pubsub/pubsub.go | 47 ----------------- 12 files changed, 112 insertions(+), 64 deletions(-) diff --git a/router/core/executor.go b/router/core/executor.go index 63073e06aa..037163a99a 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -18,6 +18,7 @@ import ( nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" ) type ExecutorConfigurationBuilder struct { @@ -30,6 +31,8 @@ type ExecutorConfigurationBuilder struct { transportOptions *TransportOptions subscriptionClientOptions *SubscriptionClientOptions instanceData InstanceData + + addPubSubProviderCallback func(provider datasource.PubSubProvider) } type Executor struct { @@ -207,7 +210,7 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con routerEngineCfg.Execution.EnableSingleFlight, routerEngineCfg.Execution.EnableNetPoll, b.instanceData, - ), b.logger) + ), b.logger, b.addPubSubProviderCallback) // this generates the plan config using the data source factories from the config package planConfig, err := loader.Load(engineConfig, subgraphs, routerEngineCfg) diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index e12d64b5df..6b684ed733 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -11,6 +11,7 @@ import ( "github.com/buger/jsonparser" "github.com/wundergraph/cosmo/router/pkg/pubsub" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" "github.com/wundergraph/cosmo/router/pkg/config" @@ -30,8 +31,9 @@ type Loader struct { ctx context.Context resolver FactoryResolver // includeInfo controls whether additional information like type usage and field usage is included in the plan de - includeInfo bool - logger *zap.Logger + includeInfo bool + logger *zap.Logger + addPubSubProviderCallback func(provider datasource.PubSubProvider) } type InstanceData struct { @@ -159,12 +161,13 @@ func (d *DefaultFactoryResolver) InstanceData() InstanceData { return d.instanceData } -func NewLoader(ctx context.Context, includeInfo bool, resolver FactoryResolver, logger *zap.Logger) *Loader { +func NewLoader(ctx context.Context, includeInfo bool, resolver FactoryResolver, logger *zap.Logger, addPubSubProviderCallback func(provider datasource.PubSubProvider)) *Loader { return &Loader{ - ctx: ctx, - resolver: resolver, - includeInfo: includeInfo, - logger: logger, + ctx: ctx, + resolver: resolver, + includeInfo: includeInfo, + logger: logger, + addPubSubProviderCallback: addPubSubProviderCallback, } } @@ -414,14 +417,33 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod case nodev1.DataSourceKind_PUBSUB: var err error - out, err = pubsub.GetDataSourceFromConfig( - l.ctx, - in, - l.dataSourceMetaData(in), - routerEngineConfig.Events, - l.logger, - l.resolver.InstanceData().hostName, - l.resolver.InstanceData().listenAddress, + var providers []datasource.PubSubProvider + + dsMeta := l.dataSourceMetaData(in) + for _, providerFactory := range pubsub.GetProviderFactories() { + provider, err := providerFactory( + l.ctx, + in, + dsMeta, + routerEngineConfig.Events, + l.logger, + l.resolver.InstanceData().hostName, + l.resolver.InstanceData().listenAddress, + ) + if err != nil { + return nil, err + } + if provider != nil { + providers = append(providers, provider) + l.addPubSubProviderCallback(provider) + } + } + + out, err = plan.NewDataSourceConfiguration( + in.Id, + datasource.NewFactory(l.ctx, routerEngineConfig.Events, providers), + dsMeta, + providers, ) if err != nil { return nil, err diff --git a/router/core/graph_server.go b/router/core/graph_server.go index bf795e9f8e..351cb1a1a5 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -44,6 +44,7 @@ import ( "github.com/wundergraph/cosmo/router/pkg/logging" rmetric "github.com/wundergraph/cosmo/router/pkg/metric" "github.com/wundergraph/cosmo/router/pkg/otel" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/cosmo/router/pkg/statistics" rtrace "github.com/wundergraph/cosmo/router/pkg/trace" ) @@ -85,6 +86,7 @@ type ( otlpEngineMetrics *rmetric.EngineMetrics prometheusEngineMetrics *rmetric.EngineMetrics instanceData InstanceData + pubSubProviders []datasource.PubSubProvider } ) @@ -880,6 +882,12 @@ func (s *graphServer) buildGraphMux(ctx context.Context, SubgraphErrorPropagation: s.subgraphErrorPropagation, } + var pubSubProviders []datasource.PubSubProvider + + addPubSubProviderCallback := func(provider datasource.PubSubProvider) { + pubSubProviders = append(pubSubProviders, provider) + } + ecb := &ExecutorConfigurationBuilder{ introspection: s.introspection, baseURL: s.baseURL, @@ -909,6 +917,7 @@ func (s *graphServer) buildGraphMux(ctx context.Context, LocalhostFallbackInsideDocker: s.localhostFallbackInsideDocker, Logger: s.logger, }, + addPubSubProviderCallback: addPubSubProviderCallback, } executor, err := ecb.Build( @@ -928,6 +937,16 @@ func (s *graphServer) buildGraphMux(ctx context.Context, return nil, fmt.Errorf("failed to build plan configuration: %w", err) } + if len(pubSubProviders) > 0 { + for _, provider := range pubSubProviders { + if err := provider.Startup(ctx); err != nil { + return nil, fmt.Errorf("failed to startup pubsub provider: %w", err) + } + } + + s.pubSubProviders = pubSubProviders + } + operationProcessor := NewOperationProcessor(OperationProcessorOptions{ Executor: executor, MaxOperationSizeInBytes: int64(s.routerTrafficConfig.MaxRequestBodyBytes), @@ -1246,6 +1265,12 @@ func (s *graphServer) Shutdown(ctx context.Context) error { } } + for _, provider := range s.pubSubProviders { + if providerErr := provider.Shutdown(ctx); providerErr != nil { + finalErr = errors.Join(finalErr, providerErr) + } + } + // Shutdown all graphs muxes to release resources // e.g. planner cache s.graphMuxListLock.Lock() diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index c30f108a23..b842fbbcce 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -23,6 +23,7 @@ import ( "go.uber.org/zap" "github.com/wundergraph/cosmo/router/pkg/config" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" @@ -282,7 +283,7 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo streamingClient: http.DefaultClient, subscriptionClient: subscriptionClient, transportOptions: &TransportOptions{SubgraphTransportOptions: NewSubgraphTransportOptions(config.TrafficShapingRules{})}, - }, logger) + }, logger, func(provider datasource.PubSubProvider) {}) // this generates the plan configuration using the data source factories from the config package planConfig, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &routerEngineConfig) diff --git a/router/pkg/pubsub/datasource/provider.go b/router/pkg/pubsub/datasource/provider.go index d496a4e3a0..fa58f3b30e 100644 --- a/router/pkg/pubsub/datasource/provider.go +++ b/router/pkg/pubsub/datasource/provider.go @@ -14,5 +14,7 @@ type ProviderFactory func(ctx context.Context, in *nodev1.DataSourceConfiguratio type ArgumentTemplateCallback func(tpl string) (string, error) type PubSubProvider interface { + Startup(ctx context.Context) error + Shutdown(ctx context.Context) error FindPubSubDataSource(typeName string, fieldName string, extractFn ArgumentTemplateCallback) (PubSubDataSource, error) } diff --git a/router/pkg/pubsub/kafka/adapter.go b/router/pkg/pubsub/kafka/adapter.go index 675c225c9d..4d38085fe5 100644 --- a/router/pkg/pubsub/kafka/adapter.go +++ b/router/pkg/pubsub/kafka/adapter.go @@ -74,6 +74,7 @@ func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (Adapte type AdapterInterface interface { Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error Publish(ctx context.Context, event PublishEventConfiguration) error + Shutdown(ctx context.Context) error } // Adapter is a Kafka pubsub implementation. diff --git a/router/pkg/pubsub/kafka/engine_datasource_test.go b/router/pkg/pubsub/kafka/engine_datasource_test.go index ca9065c7df..9939ce85d4 100644 --- a/router/pkg/pubsub/kafka/engine_datasource_test.go +++ b/router/pkg/pubsub/kafka/engine_datasource_test.go @@ -29,6 +29,11 @@ func (m *EngineDataSourceMockAdapter) Publish(ctx context.Context, event Publish return args.Error(0) } +func (m *EngineDataSourceMockAdapter) Shutdown(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + // MockSubscriptionUpdater implements resolve.SubscriptionUpdater type MockSubscriptionUpdater struct { mock.Mock diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go index 1e86916355..eda78fe824 100644 --- a/router/pkg/pubsub/kafka/provider.go +++ b/router/pkg/pubsub/kafka/provider.go @@ -101,3 +101,16 @@ func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, } return nil, nil } + +func (c *PubSubProvider) Startup(ctx context.Context) error { + return nil +} + +func (c *PubSubProvider) Shutdown(ctx context.Context) error { + for _, provider := range c.Providers { + if err := provider.Shutdown(ctx); err != nil { + return err + } + } + return nil +} diff --git a/router/pkg/pubsub/kafka/provider_test.go b/router/pkg/pubsub/kafka/provider_test.go index de54da2dfa..0fa63a73b8 100644 --- a/router/pkg/pubsub/kafka/provider_test.go +++ b/router/pkg/pubsub/kafka/provider_test.go @@ -31,6 +31,11 @@ func (m *mockAdapter) Publish(ctx context.Context, event PublishEventConfigurati return args.Error(0) } +func (m *mockAdapter) Shutdown(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + func TestBuildKafkaOptions(t *testing.T) { t.Run("basic configuration", func(t *testing.T) { cfg := config.KafkaEventSource{ diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go index b97bbd034b..24f6247a43 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource_test.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -31,6 +31,11 @@ func (m *MockAdapter) Publish(ctx context.Context, event PublishEventConfigurati return args.Error(0) } +func (m *MockAdapter) Shutdown(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + func TestKafkaPubSubDataSource(t *testing.T) { // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ diff --git a/router/pkg/pubsub/nats/provider.go b/router/pkg/pubsub/nats/provider.go index e94c3a227d..eaa4c4d5a3 100644 --- a/router/pkg/pubsub/nats/provider.go +++ b/router/pkg/pubsub/nats/provider.go @@ -149,3 +149,16 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta return nil, nil } + +func (c *PubSubProvider) Startup(ctx context.Context) error { + return nil +} + +func (c *PubSubProvider) Shutdown(ctx context.Context) error { + for _, provider := range c.Providers { + if err := provider.Shutdown(ctx); err != nil { + return err + } + } + return nil +} diff --git a/router/pkg/pubsub/pubsub.go b/router/pkg/pubsub/pubsub.go index 8ae599be28..e70cf66f33 100644 --- a/router/pkg/pubsub/pubsub.go +++ b/router/pkg/pubsub/pubsub.go @@ -1,14 +1,7 @@ package pubsub import ( - "context" - "fmt" - - nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/config" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" - "go.uber.org/zap" "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" @@ -28,43 +21,3 @@ func GetProviderFactories() []datasource.ProviderFactory { nats.GetProvider, }, additionalProviders...) } - -// GetDataSourceFromConfig returns a new plan.DataSource from the given configuration, -// using the registered PubSub providers. -func GetDataSourceFromConfig( - ctx context.Context, - in *nodev1.DataSourceConfiguration, - dsMeta *plan.DataSourceMetadata, - config config.EventsConfiguration, - logger *zap.Logger, - hostName string, - routerListenAddr string, -) (plan.DataSource, error) { - var providers []datasource.PubSubProvider - - for _, providerFactory := range GetProviderFactories() { - provider, err := providerFactory(ctx, in, dsMeta, config, logger, hostName, routerListenAddr) - if err != nil { - return nil, err - } - if provider != nil { - providers = append(providers, provider) - } - } - - if len(providers) == 0 { - return nil, fmt.Errorf("no pubsub data sources found for data source %s", in.Id) - } - - ds, err := plan.NewDataSourceConfiguration( - in.Id, - datasource.NewFactory(ctx, config, providers), - dsMeta, - providers, - ) - if err != nil { - return nil, err - } - - return ds, nil -} From 8295c26eab7fffa5209890ada2c014c3f797d4d7 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 14 Apr 2025 09:50:49 +0200 Subject: [PATCH 24/49] refactor: update Kafka and NATS adapters to use direct client connections and implement Startup method --- router/pkg/pubsub/kafka/adapter.go | 75 ++-- .../pubsub/kafka/engine_datasource_test.go | 44 +- router/pkg/pubsub/kafka/provider.go | 5 + router/pkg/pubsub/kafka/provider_test.go | 5 + .../pubsub/kafka/pubsub_datasource_test.go | 29 +- router/pkg/pubsub/nats/adapter.go | 85 ++-- .../pkg/pubsub/nats/engine_datasource_test.go | 57 +-- router/pkg/pubsub/nats/lazy_client.go | 53 --- router/pkg/pubsub/nats/lazy_client_test.go | 413 ------------------ router/pkg/pubsub/nats/provider.go | 5 + router/pkg/pubsub/nats/provider_test.go | 5 + .../pkg/pubsub/nats/pubsub_datasource_test.go | 35 +- 12 files changed, 129 insertions(+), 682 deletions(-) delete mode 100644 router/pkg/pubsub/nats/lazy_client.go delete mode 100644 router/pkg/pubsub/nats/lazy_client_test.go diff --git a/router/pkg/pubsub/kafka/adapter.go b/router/pkg/pubsub/kafka/adapter.go index 4d38085fe5..7e9d3486c0 100644 --- a/router/pkg/pubsub/kafka/adapter.go +++ b/router/pkg/pubsub/kafka/adapter.go @@ -19,54 +19,18 @@ var ( errClientClosed = errors.New("client closed") ) -type LazyClient struct { - once sync.Once - client *kgo.Client - opts []kgo.Opt -} - -func (c *LazyClient) Connect() (err error) { - c.once.Do(func() { - c.client, err = kgo.NewClient(append(c.opts, - // For observability, we set the client ID to "router" - kgo.ClientID("cosmo.router.producer"))..., - ) - }) - - return -} - -func (c *LazyClient) GetClient() *kgo.Client { - if c.client == nil { - c.Connect() - } - return c.client -} - -func NewLazyClient(opts ...kgo.Opt) *LazyClient { - return &LazyClient{ - opts: opts, - } -} - func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (AdapterInterface, error) { ctx, cancel := context.WithCancel(ctx) if logger == nil { logger = zap.NewNop() } - client := NewLazyClient(append(opts, - // For observability, we set the client ID to "router" - kgo.ClientID("cosmo.router.producer"))..., - ) - return &Adapter{ - ctx: ctx, - logger: logger.With(zap.String("pubsub", "kafka")), - opts: opts, - writeClient: client, - closeWg: sync.WaitGroup{}, - cancel: cancel, + ctx: ctx, + logger: logger.With(zap.String("pubsub", "kafka")), + opts: opts, + closeWg: sync.WaitGroup{}, + cancel: cancel, }, nil } @@ -74,6 +38,7 @@ func NewAdapter(ctx context.Context, logger *zap.Logger, opts []kgo.Opt) (Adapte type AdapterInterface interface { Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error Publish(ctx context.Context, event PublishEventConfiguration) error + Startup(ctx context.Context) error Shutdown(ctx context.Context) error } @@ -86,7 +51,7 @@ type Adapter struct { ctx context.Context opts []kgo.Opt logger *zap.Logger - writeClient *LazyClient + writeClient *kgo.Client closeWg sync.WaitGroup cancel context.CancelFunc } @@ -206,6 +171,10 @@ func (p *Adapter) Publish(ctx context.Context, event PublishEventConfiguration) zap.String("topic", event.Topic), ) + if p.writeClient == nil { + return datasource.NewError("kafka write client not initialized", nil) + } + log.Debug("publish", zap.ByteString("data", event.Data)) var wg sync.WaitGroup @@ -213,7 +182,7 @@ func (p *Adapter) Publish(ctx context.Context, event PublishEventConfiguration) var pErr error - p.writeClient.GetClient().Produce(ctx, &kgo.Record{ + p.writeClient.Produce(ctx, &kgo.Record{ Topic: event.Topic, Value: event.Data, }, func(record *kgo.Record, err error) { @@ -233,14 +202,30 @@ func (p *Adapter) Publish(ctx context.Context, event PublishEventConfiguration) return nil } +func (p *Adapter) Startup(ctx context.Context) (err error) { + p.writeClient, err = kgo.NewClient(append(p.opts, + // For observability, we set the client ID to "router" + kgo.ClientID("cosmo.router.producer"))..., + ) + if err != nil { + return err + } + + return +} + func (p *Adapter) Shutdown(ctx context.Context) error { - err := p.writeClient.GetClient().Flush(ctx) + if p.writeClient == nil { + return nil + } + + err := p.writeClient.Flush(ctx) if err != nil { p.logger.Error("flushing write client", zap.Error(err)) } - p.writeClient.GetClient().Close() + p.writeClient.Close() // Cancel the context to stop all pollers p.cancel() diff --git a/router/pkg/pubsub/kafka/engine_datasource_test.go b/router/pkg/pubsub/kafka/engine_datasource_test.go index 9939ce85d4..b390f315bb 100644 --- a/router/pkg/pubsub/kafka/engine_datasource_test.go +++ b/router/pkg/pubsub/kafka/engine_datasource_test.go @@ -14,26 +14,6 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -// EngineDataSourceMockAdapter is a mock implementation of AdapterInterface for testing -type EngineDataSourceMockAdapter struct { - mock.Mock -} - -func (m *EngineDataSourceMockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { - args := m.Called(ctx, event, updater) - return args.Error(0) -} - -func (m *EngineDataSourceMockAdapter) Publish(ctx context.Context, event PublishEventConfiguration) error { - args := m.Called(ctx, event) - return args.Error(0) -} - -func (m *EngineDataSourceMockAdapter) Shutdown(ctx context.Context) error { - args := m.Called(ctx) - return args.Error(0) -} - // MockSubscriptionUpdater implements resolve.SubscriptionUpdater type MockSubscriptionUpdater struct { mock.Mock @@ -110,7 +90,7 @@ func TestSubscriptionSource_UniqueRequestID(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { source := &SubscriptionDataSource{ - pubSub: &EngineDataSourceMockAdapter{}, + pubSub: &mockAdapter{}, } ctx := &resolve.Context{} input := []byte(tt.input) @@ -137,13 +117,13 @@ func TestSubscriptionSource_Start(t *testing.T) { tests := []struct { name string input string - mockSetup func(*EngineDataSourceMockAdapter) + mockSetup func(*mockAdapter) expectError bool }{ { name: "successful subscription", input: `{"topics":["topic1", "topic2"], "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ ProviderID: "test-provider", Topics: []string{"topic1", "topic2"}, @@ -154,7 +134,7 @@ func TestSubscriptionSource_Start(t *testing.T) { { name: "adapter returns error", input: `{"topics":["topic1"], "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ ProviderID: "test-provider", Topics: []string{"topic1"}, @@ -165,14 +145,14 @@ func TestSubscriptionSource_Start(t *testing.T) { { name: "invalid input json", input: `{"invalid json":`, - mockSetup: func(m *EngineDataSourceMockAdapter) {}, + mockSetup: func(m *mockAdapter) {}, expectError: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - mockAdapter := new(EngineDataSourceMockAdapter) + mockAdapter := new(mockAdapter) tt.mockSetup(mockAdapter) source := &SubscriptionDataSource{ @@ -207,7 +187,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { tests := []struct { name string input string - mockSetup func(*EngineDataSourceMockAdapter) + mockSetup func(*mockAdapter) expectError bool expectedOutput string expectPublished bool @@ -215,7 +195,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { { name: "successful publish", input: `{"topic":"test-topic", "data":{"message":"hello"}, "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishEventConfiguration) bool { return event.ProviderID == "test-provider" && event.Topic == "test-topic" && @@ -229,7 +209,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { { name: "publish error", input: `{"topic":"test-topic", "data":{"message":"hello"}, "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Publish", mock.Anything, mock.Anything).Return(errors.New("publish error")) }, expectError: false, // The Load method doesn't return the publish error directly @@ -239,7 +219,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { { name: "invalid input json", input: `{"invalid json":`, - mockSetup: func(m *EngineDataSourceMockAdapter) {}, + mockSetup: func(m *mockAdapter) {}, expectError: true, expectPublished: false, }, @@ -247,7 +227,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - mockAdapter := new(EngineDataSourceMockAdapter) + mockAdapter := new(mockAdapter) tt.mockSetup(mockAdapter) dataSource := &PublishDataSource{ @@ -276,7 +256,7 @@ func TestKafkaPublishDataSource_Load(t *testing.T) { func TestKafkaPublishDataSource_LoadWithFiles(t *testing.T) { t.Run("panic on not implemented", func(t *testing.T) { dataSource := &PublishDataSource{ - pubSub: &EngineDataSourceMockAdapter{}, + pubSub: &mockAdapter{}, } assert.Panics(t, func() { diff --git a/router/pkg/pubsub/kafka/provider.go b/router/pkg/pubsub/kafka/provider.go index eda78fe824..3d444f7923 100644 --- a/router/pkg/pubsub/kafka/provider.go +++ b/router/pkg/pubsub/kafka/provider.go @@ -103,6 +103,11 @@ func (c *PubSubProvider) FindPubSubDataSource(typeName string, fieldName string, } func (c *PubSubProvider) Startup(ctx context.Context) error { + for _, provider := range c.Providers { + if err := provider.Startup(ctx); err != nil { + return err + } + } return nil } diff --git a/router/pkg/pubsub/kafka/provider_test.go b/router/pkg/pubsub/kafka/provider_test.go index 0fa63a73b8..52621affe6 100644 --- a/router/pkg/pubsub/kafka/provider_test.go +++ b/router/pkg/pubsub/kafka/provider_test.go @@ -31,6 +31,11 @@ func (m *mockAdapter) Publish(ctx context.Context, event PublishEventConfigurati return args.Error(0) } +func (m *mockAdapter) Startup(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + func (m *mockAdapter) Shutdown(ctx context.Context) error { args := m.Called(ctx) return args.Error(0) diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go index 24f6247a43..4388a637a8 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource_test.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -10,32 +10,8 @@ import ( "github.com/stretchr/testify/require" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -// MockAdapter mocks the required functionality from the Adapter for testing -type MockAdapter struct { - mock.Mock -} - -// Ensure MockAdapter implements KafkaAdapterInterface -var _ AdapterInterface = (*MockAdapter)(nil) - -func (m *MockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { - args := m.Called(ctx, event, updater) - return args.Error(0) -} - -func (m *MockAdapter) Publish(ctx context.Context, event PublishEventConfiguration) error { - args := m.Called(ctx, event) - return args.Error(0) -} - -func (m *MockAdapter) Shutdown(ctx context.Context) error { - args := m.Called(ctx) - return args.Error(0) -} - func TestKafkaPubSubDataSource(t *testing.T) { // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ @@ -77,7 +53,7 @@ func TestPubSubDataSourceWithMockAdapter(t *testing.T) { } // Create mock adapter - mockAdapter := new(MockAdapter) + mockAdapter := new(mockAdapter) // Configure mock expectations for Publish mockAdapter.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishEventConfiguration) bool { @@ -124,7 +100,7 @@ func TestPubSubDataSource_GetResolveDataSource_WrongType(t *testing.T) { } // Create mock adapter - mockAdapter := new(MockAdapter) + mockAdapter := new(mockAdapter) // Create the data source with mock adapter pubsub := &PubSubDataSource{ @@ -224,5 +200,4 @@ func TestKafkaPubSubDataSourceMultiTopicSubscription(t *testing.T) { require.Equal(t, 2, len(subscriptionConfig.Topics), "Expected 2 topics in subscription configuration") require.Equal(t, "test-topic-1", subscriptionConfig.Topics[0], "Expected first topic to be 'test-topic-1'") require.Equal(t, "test-topic-2", subscriptionConfig.Topics[1], "Expected second topic to be 'test-topic-2'") - } diff --git a/router/pkg/pubsub/nats/adapter.go b/router/pkg/pubsub/nats/adapter.go index bcdad2f3d1..34a646b158 100644 --- a/router/pkg/pubsub/nats/adapter.go +++ b/router/pkg/pubsub/nats/adapter.go @@ -24,6 +24,8 @@ type AdapterInterface interface { Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error // Request sends a request to the specified subject and writes the response to the given writer Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error + // Startup initializes the adapter + Startup(ctx context.Context) error // Shutdown gracefully shuts down the adapter Shutdown(ctx context.Context) error } @@ -31,11 +33,14 @@ type AdapterInterface interface { // Adapter implements the AdapterInterface for NATS pub/sub type Adapter struct { ctx context.Context - client *LazyClient + client *nats.Conn + js jetstream.JetStream logger *zap.Logger closeWg sync.WaitGroup hostName string routerListenAddr string + url string + opts []nats.Option } // getInstanceIdentifier returns an identifier for the current instance. @@ -73,6 +78,14 @@ func (p *Adapter) Subscribe(ctx context.Context, event SubscriptionEventConfigur zap.Strings("subjects", event.Subjects), ) + if p.client == nil { + return datasource.NewError("nats client not initialized", nil) + } + + if p.js == nil { + return datasource.NewError("nats jetstream not initialized", nil) + } + if event.StreamConfiguration != nil { durableConsumerName, err := p.getDurableConsumerName(event.StreamConfiguration.Consumer, event.Subjects) if err != nil { @@ -86,13 +99,8 @@ func (p *Adapter) Subscribe(ctx context.Context, event SubscriptionEventConfigur if event.StreamConfiguration.ConsumerInactiveThreshold > 0 { consumerConfig.InactiveThreshold = time.Duration(event.StreamConfiguration.ConsumerInactiveThreshold) * time.Second } - js, err := p.client.GetJetStream() - if err != nil { - log.Error("getting jetstream client", zap.Error(err)) - return datasource.NewError("failed to get jetstream client", err) - } - consumer, err := js.CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) + consumer, err := p.js.CreateOrUpdateConsumer(ctx, event.StreamConfiguration.StreamName, consumerConfig) if err != nil { log.Error("creating or updating consumer", zap.Error(err)) return datasource.NewError(fmt.Sprintf(`failed to create or update consumer for stream "%s"`, event.StreamConfiguration.StreamName), err) @@ -140,16 +148,10 @@ func (p *Adapter) Subscribe(ctx context.Context, event SubscriptionEventConfigur return nil } - nc, err := p.client.GetClient() - if err != nil { - log.Error("getting nats client", zap.Error(err)) - return datasource.NewError("failed to get nats client", err) - } - msgChan := make(chan *nats.Msg) subscriptions := make([]*nats.Subscription, len(event.Subjects)) for i, subject := range event.Subjects { - subscription, err := nc.ChanSubscribe(subject, msgChan) + subscription, err := p.client.ChanSubscribe(subject, msgChan) if err != nil { log.Error("subscribing to NATS subject", zap.Error(err), zap.String("subscription_subject", subject)) return datasource.NewError(fmt.Sprintf(`failed to subscribe to NATS subject "%s"`, subject), err) @@ -201,15 +203,13 @@ func (p *Adapter) Publish(_ context.Context, event PublishAndRequestEventConfigu zap.String("subject", event.Subject), ) - log.Debug("publish", zap.ByteString("data", event.Data)) - - nc, err := p.client.GetClient() - if err != nil { - log.Error("getting nats client", zap.Error(err)) - return datasource.NewError("failed to get nats client", err) + if p.client == nil { + return datasource.NewError("nats client not initialized", nil) } - err = nc.Publish(event.Subject, event.Data) + log.Debug("publish", zap.ByteString("data", event.Data)) + + err := p.client.Publish(event.Subject, event.Data) if err != nil { log.Error("publish error", zap.Error(err)) return datasource.NewError(fmt.Sprintf("error publishing to NATS subject %s", event.Subject), err) @@ -225,15 +225,13 @@ func (p *Adapter) Request(ctx context.Context, event PublishAndRequestEventConfi zap.String("subject", event.Subject), ) - log.Debug("request", zap.ByteString("data", event.Data)) - - nc, err := p.client.GetClient() - if err != nil { - log.Error("getting nats client", zap.Error(err)) - return datasource.NewError("failed to get nats client", err) + if p.client == nil { + return datasource.NewError("nats client not initialized", nil) } - msg, err := nc.RequestWithContext(ctx, event.Subject, event.Data) + log.Debug("request", zap.ByteString("data", event.Data)) + + msg, err := p.client.RequestWithContext(ctx, event.Subject, event.Data) if err != nil { log.Error("request error", zap.Error(err)) return datasource.NewError(fmt.Sprintf("error requesting from NATS subject %s", event.Subject), err) @@ -249,21 +247,31 @@ func (p *Adapter) Request(ctx context.Context, event PublishAndRequestEventConfi } func (p *Adapter) flush(ctx context.Context) error { - nc, err := p.client.GetClient() + if p.client == nil { + return nil + } + return p.client.FlushWithContext(ctx) +} + +func (p *Adapter) Startup(ctx context.Context) (err error) { + p.client, err = nats.Connect(p.url, p.opts...) + if err != nil { + return err + } + p.js, err = jetstream.New(p.client) if err != nil { return err } - return nc.FlushWithContext(ctx) + return nil } func (p *Adapter) Shutdown(ctx context.Context) error { - nc, err := p.client.GetClient() - if err != nil { - return nil // Already disconnected or failed to connect + if p.client == nil { + return nil } - if nc.IsClosed() { - return nil + if p.client.IsClosed() { + return nil // Already disconnected or failed to connect } var shutdownErr error @@ -273,7 +281,7 @@ func (p *Adapter) Shutdown(ctx context.Context) error { shutdownErr = errors.Join(shutdownErr, fErr) } - drainErr := nc.Drain() + drainErr := p.client.Drain() if drainErr != nil { shutdownErr = errors.Join(shutdownErr, drainErr) } @@ -293,14 +301,13 @@ func NewAdapter(ctx context.Context, logger *zap.Logger, url string, opts []nats logger = zap.NewNop() } - client := NewLazyClient(url, opts...) - return &Adapter{ ctx: ctx, - client: client, logger: logger.With(zap.String("pubsub", "nats")), closeWg: sync.WaitGroup{}, hostName: hostName, routerListenAddr: routerListenAddr, + url: url, + opts: opts, }, nil } diff --git a/router/pkg/pubsub/nats/engine_datasource_test.go b/router/pkg/pubsub/nats/engine_datasource_test.go index 4fb1ae767d..4a4dfdd20a 100644 --- a/router/pkg/pubsub/nats/engine_datasource_test.go +++ b/router/pkg/pubsub/nats/engine_datasource_test.go @@ -15,31 +15,6 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -// EngineDataSourceMockAdapter is a mock implementation of AdapterInterface for testing -type EngineDataSourceMockAdapter struct { - mock.Mock -} - -func (m *EngineDataSourceMockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { - args := m.Called(ctx, event, updater) - return args.Error(0) -} - -func (m *EngineDataSourceMockAdapter) Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error { - args := m.Called(ctx, event) - return args.Error(0) -} - -func (m *EngineDataSourceMockAdapter) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { - args := m.Called(ctx, event, w) - return args.Error(0) -} - -func (m *EngineDataSourceMockAdapter) Shutdown(ctx context.Context) error { - args := m.Called(ctx) - return args.Error(0) -} - // MockSubscriptionUpdater implements resolve.SubscriptionUpdater type MockSubscriptionUpdater struct { mock.Mock @@ -150,7 +125,7 @@ func TestSubscriptionSource_UniqueRequestID(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { source := &SubscriptionSource{ - pubSub: &EngineDataSourceMockAdapter{}, + pubSub: &mockAdapter{}, } ctx := &resolve.Context{} input := []byte(tt.input) @@ -177,13 +152,13 @@ func TestSubscriptionSource_Start(t *testing.T) { tests := []struct { name string input string - mockSetup func(*EngineDataSourceMockAdapter) + mockSetup func(*mockAdapter) expectError bool }{ { name: "successful subscription", input: `{"subjects":["subject1", "subject2"], "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ ProviderID: "test-provider", Subjects: []string{"subject1", "subject2"}, @@ -194,7 +169,7 @@ func TestSubscriptionSource_Start(t *testing.T) { { name: "adapter returns error", input: `{"subjects":["subject1"], "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Subscribe", mock.Anything, SubscriptionEventConfiguration{ ProviderID: "test-provider", Subjects: []string{"subject1"}, @@ -205,14 +180,14 @@ func TestSubscriptionSource_Start(t *testing.T) { { name: "invalid input json", input: `{"invalid json":`, - mockSetup: func(m *EngineDataSourceMockAdapter) {}, + mockSetup: func(m *mockAdapter) {}, expectError: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - mockAdapter := new(EngineDataSourceMockAdapter) + mockAdapter := new(mockAdapter) tt.mockSetup(mockAdapter) source := &SubscriptionSource{ @@ -247,7 +222,7 @@ func TestNatsPublishDataSource_Load(t *testing.T) { tests := []struct { name string input string - mockSetup func(*EngineDataSourceMockAdapter) + mockSetup func(*mockAdapter) expectError bool expectedOutput string expectPublished bool @@ -255,7 +230,7 @@ func TestNatsPublishDataSource_Load(t *testing.T) { { name: "successful publish", input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { return event.ProviderID == "test-provider" && event.Subject == "test-subject" && @@ -269,7 +244,7 @@ func TestNatsPublishDataSource_Load(t *testing.T) { { name: "publish error", input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Publish", mock.Anything, mock.Anything).Return(errors.New("publish error")) }, expectError: false, // The Load method doesn't return the publish error directly @@ -279,7 +254,7 @@ func TestNatsPublishDataSource_Load(t *testing.T) { { name: "invalid input json", input: `{"invalid json":`, - mockSetup: func(m *EngineDataSourceMockAdapter) {}, + mockSetup: func(m *mockAdapter) {}, expectError: true, expectPublished: false, }, @@ -287,7 +262,7 @@ func TestNatsPublishDataSource_Load(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - mockAdapter := new(EngineDataSourceMockAdapter) + mockAdapter := new(mockAdapter) tt.mockSetup(mockAdapter) dataSource := &NatsPublishDataSource{ @@ -326,14 +301,14 @@ func TestNatsRequestDataSource_Load(t *testing.T) { tests := []struct { name string input string - mockSetup func(*EngineDataSourceMockAdapter) + mockSetup func(*mockAdapter) expectError bool expectedOutput string }{ { name: "successful request", input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Request", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { return event.ProviderID == "test-provider" && event.Subject == "test-subject" && @@ -350,7 +325,7 @@ func TestNatsRequestDataSource_Load(t *testing.T) { { name: "request error", input: `{"subject":"test-subject", "data":{"message":"hello"}, "providerId":"test-provider"}`, - mockSetup: func(m *EngineDataSourceMockAdapter) { + mockSetup: func(m *mockAdapter) { m.On("Request", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("request error")) }, expectError: true, @@ -359,7 +334,7 @@ func TestNatsRequestDataSource_Load(t *testing.T) { { name: "invalid input json", input: `{"invalid json":`, - mockSetup: func(m *EngineDataSourceMockAdapter) {}, + mockSetup: func(m *mockAdapter) {}, expectError: true, expectedOutput: "", }, @@ -367,7 +342,7 @@ func TestNatsRequestDataSource_Load(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - mockAdapter := new(EngineDataSourceMockAdapter) + mockAdapter := new(mockAdapter) tt.mockSetup(mockAdapter) dataSource := &NatsRequestDataSource{ diff --git a/router/pkg/pubsub/nats/lazy_client.go b/router/pkg/pubsub/nats/lazy_client.go deleted file mode 100644 index a5abc82bda..0000000000 --- a/router/pkg/pubsub/nats/lazy_client.go +++ /dev/null @@ -1,53 +0,0 @@ -package nats - -import ( - "sync" - - "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" -) - -type LazyClient struct { - once sync.Once - url string - opts []nats.Option - client *nats.Conn - js jetstream.JetStream - err error -} - -func (c *LazyClient) Connect(opts ...nats.Option) error { - c.once.Do(func() { - c.client, c.err = nats.Connect(c.url, opts...) - if c.err != nil { - return - } - c.js, c.err = jetstream.New(c.client) - }) - return c.err -} - -func (c *LazyClient) GetClient() (*nats.Conn, error) { - if c.client == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.client, c.err -} - -func (c *LazyClient) GetJetStream() (jetstream.JetStream, error) { - if c.js == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.js, c.err -} - -func NewLazyClient(url string, opts ...nats.Option) *LazyClient { - return &LazyClient{ - url: url, - opts: opts, - } -} diff --git a/router/pkg/pubsub/nats/lazy_client_test.go b/router/pkg/pubsub/nats/lazy_client_test.go deleted file mode 100644 index 03b63b27a8..0000000000 --- a/router/pkg/pubsub/nats/lazy_client_test.go +++ /dev/null @@ -1,413 +0,0 @@ -package nats - -import ( - "errors" - "sync" - "testing" - "time" - - "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// Define interfaces for our testing -type natsConnector interface { - Connect(url string, opts ...nats.Option) (*nats.Conn, error) -} - -type jetStreamCreator interface { - Create(nc *nats.Conn) (jetstream.JetStream, error) -} - -// mockNatsConnector implements natsConnector for testing -type mockNatsConnector struct { - mockConn *nats.Conn - mockErr error - callCount int - lastOpts []nats.Option -} - -func (m *mockNatsConnector) Connect(url string, opts ...nats.Option) (*nats.Conn, error) { - m.callCount++ - m.lastOpts = opts - return m.mockConn, m.mockErr -} - -// threadSafeConnector is a specialized version for thread safety testing -type threadSafeConnector struct { - mockConn *nats.Conn - mockErr error - callCount int - countMutex sync.Mutex - sleepTime time.Duration -} - -func (t *threadSafeConnector) Connect(url string, opts ...nats.Option) (*nats.Conn, error) { - // Simulate a slow connection - time.Sleep(t.sleepTime) - t.countMutex.Lock() - t.callCount++ - t.countMutex.Unlock() - return t.mockConn, t.mockErr -} - -// mockJetStreamCreator implements jetStreamCreator for testing -type mockJetStreamCreator struct { - mockJS jetstream.JetStream - mockErr error - callCount int -} - -func (m *mockJetStreamCreator) Create(nc *nats.Conn) (jetstream.JetStream, error) { - m.callCount++ - return m.mockJS, m.mockErr -} - -// testLazyClient is a wrapper for LazyClient that allows us to inject mocks -type testLazyClient struct { - *LazyClient - connector natsConnector - jetStreamCreator jetStreamCreator -} - -// newTestLazyClient creates a LazyClient with mocked dependencies -func newTestLazyClient(url string, connector natsConnector, creator jetStreamCreator, opts ...nats.Option) *testLazyClient { - c := &testLazyClient{ - LazyClient: NewLazyClient(url, opts...), - connector: connector, - jetStreamCreator: creator, - } - - return c -} - -// Connect overrides LazyClient.Connect to use our mocks -func (c *testLazyClient) Connect(opts ...nats.Option) error { - c.once.Do(func() { - // If no options are provided, use the ones stored during initialization - optionsToUse := opts - if len(optionsToUse) == 0 { - optionsToUse = c.opts - } - c.client, c.err = c.connector.Connect(c.url, optionsToUse...) - if c.err != nil { - return - } - c.js, c.err = c.jetStreamCreator.Create(c.client) - }) - return c.err -} - -// GetClient overrides LazyClient.GetClient to ensure we use our Connect method -func (c *testLazyClient) GetClient() (*nats.Conn, error) { - if c.client == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.client, c.err -} - -// GetJetStream overrides LazyClient.GetJetStream to ensure we use our Connect method -func (c *testLazyClient) GetJetStream() (jetstream.JetStream, error) { - if c.js == nil { - if err := c.Connect(c.opts...); err != nil { - return nil, err - } - } - return c.js, c.err -} - -func TestNewLazyClient(t *testing.T) { - url := "nats://localhost:4222" - opts := []nats.Option{nats.Name("test-client")} - - client := NewLazyClient(url, opts...) - - assert.Equal(t, url, client.url) - assert.Equal(t, opts, client.opts) - assert.Nil(t, client.client) - assert.Nil(t, client.js) - assert.Nil(t, client.err) -} - -func TestLazyClient_Connect_Success(t *testing.T) { - // Create mocks - connector := &mockNatsConnector{ - mockConn: &nats.Conn{}, - } - - jsCreator := &mockJetStreamCreator{ - mockJS: &struct{ jetstream.JetStream }{}, - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // First call should connect - err := client.Connect() - require.NoError(t, err) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") - assert.Equal(t, connector.mockConn, client.client) - assert.Equal(t, jsCreator.mockJS, client.js) - - // Second call should not connect again due to sync.Once - err = client.Connect() - require.NoError(t, err) - assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should still be called only once") -} - -func TestLazyClient_Connect_Error(t *testing.T) { - // Create mocks with connection error - connector := &mockNatsConnector{ - mockErr: errors.New("mock connect error"), - } - - jsCreator := &mockJetStreamCreator{ - mockJS: &struct{ jetstream.JetStream }{}, - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // Connect should return an error - err := client.Connect() - require.Error(t, err) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should not be called") - assert.Nil(t, client.client) - assert.Nil(t, client.js) - assert.NotNil(t, client.err) - assert.Equal(t, "mock connect error", client.err.Error()) - - // Second call should not connect again and return the same error - connector.mockErr = errors.New("different error") // Should be ignored due to sync.Once - err2 := client.Connect() - require.Error(t, err2) - assert.Equal(t, err, err2, "Should return the same error") - assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") - assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should still not be called") -} - -func TestLazyClient_Connect_JetStreamError(t *testing.T) { - // Create mocks with jetstream error - connector := &mockNatsConnector{ - mockConn: &nats.Conn{}, - } - - jsCreator := &mockJetStreamCreator{ - mockErr: errors.New("mock jetstream error"), - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // Connect should return the jetstream error - err := client.Connect() - require.Error(t, err) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") - assert.NotNil(t, client.client) // NATS connection was successful - assert.Nil(t, client.js) // JetStream failed - assert.NotNil(t, client.err) - assert.Equal(t, "mock jetstream error", client.err.Error()) - - // Second call should not connect again and return the same error - jsCreator.mockErr = errors.New("different error") // Should be ignored due to sync.Once - err2 := client.Connect() - require.Error(t, err2) - assert.Equal(t, err, err2, "Should return the same error") - assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should still be called only once") -} - -func TestLazyClient_GetClient(t *testing.T) { - t.Run("with successful connection", func(t *testing.T) { - // Create mocks - connector := &mockNatsConnector{ - mockConn: &nats.Conn{}, - } - - jsCreator := &mockJetStreamCreator{ - mockJS: &struct{ jetstream.JetStream }{}, - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // First call to GetClient should connect - conn, err := client.GetClient() - require.NoError(t, err) - assert.Equal(t, connector.mockConn, conn) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - - // Second call should not connect again - conn2, err := client.GetClient() - require.NoError(t, err) - assert.Equal(t, conn, conn2) - assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") - }) - - t.Run("with failed connection", func(t *testing.T) { - // Create mocks with connection error - connector := &mockNatsConnector{ - mockErr: errors.New("mock connect error"), - } - - jsCreator := &mockJetStreamCreator{} - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // GetClient should attempt to connect and return an error - conn, err := client.GetClient() - require.Error(t, err) - assert.Nil(t, conn) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should not be called") - assert.Equal(t, "mock connect error", err.Error()) - }) -} - -func TestLazyClient_GetJetStream(t *testing.T) { - t.Run("with successful connection", func(t *testing.T) { - // Create mocks - connector := &mockNatsConnector{ - mockConn: &nats.Conn{}, - } - - jsCreator := &mockJetStreamCreator{ - mockJS: &struct{ jetstream.JetStream }{}, - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // First call to GetJetStream should connect - js, err := client.GetJetStream() - require.NoError(t, err) - assert.Equal(t, jsCreator.mockJS, js) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") - - // Second call should not connect again - js2, err := client.GetJetStream() - require.NoError(t, err) - assert.Equal(t, js, js2) - assert.Equal(t, 1, connector.callCount, "Connect should still be called only once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should still be called only once") - }) - - t.Run("with failed connection", func(t *testing.T) { - // Create mocks with connection error - connector := &mockNatsConnector{ - mockErr: errors.New("mock connect error"), - } - - jsCreator := &mockJetStreamCreator{} - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // GetJetStream should attempt to connect and return an error - js, err := client.GetJetStream() - require.Error(t, err) - assert.Nil(t, js) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 0, jsCreator.callCount, "JetStream.Create should not be called") - assert.Equal(t, "mock connect error", err.Error()) - }) - - t.Run("with jetstream creation failure", func(t *testing.T) { - // Create mocks with jetstream error - connector := &mockNatsConnector{ - mockConn: &nats.Conn{}, - } - - jsCreator := &mockJetStreamCreator{ - mockErr: errors.New("mock jetstream error"), - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // GetJetStream should create the connection but fail on jetstream - js, err := client.GetJetStream() - require.Error(t, err) - assert.Nil(t, js) - assert.Equal(t, 1, connector.callCount, "Connect should be called once") - assert.Equal(t, 1, jsCreator.callCount, "JetStream.Create should be called once") - assert.Equal(t, "mock jetstream error", err.Error()) - }) -} - -func TestLazyClient_WithOptions(t *testing.T) { - // Create mocks that will track options - connector := &mockNatsConnector{ - mockConn: &nats.Conn{}, - } - - jsCreator := &mockJetStreamCreator{ - mockJS: &struct{ jetstream.JetStream }{}, - } - - // Create options that we can verify are passed to Connect - option1 := nats.Name("test-client") - option2 := nats.NoEcho() - - // Create test client with options - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator, option1, option2) - - // Call Connect to trigger the connection - err := client.Connect() - require.NoError(t, err) - - // Verify that the options were passed - they're stored in connector.lastOpts - require.Len(t, connector.lastOpts, 2, "Options should be passed to Connect") - - // Now we can call GetClient() and verify it reuses the connection - prevCallCount := connector.callCount - _, err = client.GetClient() - require.NoError(t, err) - - // Verify Connect wasn't called again - assert.Equal(t, prevCallCount, connector.callCount, "GetClient should reuse the existing connection") -} - -func TestLazyClient_ThreadSafety(t *testing.T) { - // This test checks that LazyClient's sync.Once protection works as expected - // by calling Connect from multiple goroutines simultaneously - - // Create a thread-safe connector that simulates slow connections - connector := &threadSafeConnector{ - mockConn: &nats.Conn{}, - sleepTime: 10 * time.Millisecond, - } - - jsCreator := &mockJetStreamCreator{ - mockJS: &struct{ jetstream.JetStream }{}, - } - - // Create test client - client := newTestLazyClient("nats://localhost:4222", connector, jsCreator) - - // Launch multiple goroutines that call Connect simultaneously - var wg sync.WaitGroup - for i := 0; i < 10; i++ { - wg.Add(1) - go func() { - defer wg.Done() - err := client.Connect() - assert.NoError(t, err) - }() - } - - wg.Wait() - - // Connect should have been called exactly once - assert.Equal(t, 1, connector.callCount, "Connect should have been called exactly once") -} diff --git a/router/pkg/pubsub/nats/provider.go b/router/pkg/pubsub/nats/provider.go index eaa4c4d5a3..a4e9af1248 100644 --- a/router/pkg/pubsub/nats/provider.go +++ b/router/pkg/pubsub/nats/provider.go @@ -151,6 +151,11 @@ func GetProvider(ctx context.Context, in *nodev1.DataSourceConfiguration, dsMeta } func (c *PubSubProvider) Startup(ctx context.Context) error { + for _, provider := range c.Providers { + if err := provider.Startup(ctx); err != nil { + return err + } + } return nil } diff --git a/router/pkg/pubsub/nats/provider_test.go b/router/pkg/pubsub/nats/provider_test.go index fd2255bf60..9a4ba9f03a 100644 --- a/router/pkg/pubsub/nats/provider_test.go +++ b/router/pkg/pubsub/nats/provider_test.go @@ -36,6 +36,11 @@ func (m *mockAdapter) Request(ctx context.Context, event PublishAndRequestEventC return args.Error(0) } +func (m *mockAdapter) Startup(ctx context.Context) error { + args := m.Called(ctx) + return args.Error(0) +} + func (m *mockAdapter) Shutdown(ctx context.Context) error { args := m.Called(ctx) return args.Error(0) diff --git a/router/pkg/pubsub/nats/pubsub_datasource_test.go b/router/pkg/pubsub/nats/pubsub_datasource_test.go index 43fb75fef8..d14a8721a6 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource_test.go +++ b/router/pkg/pubsub/nats/pubsub_datasource_test.go @@ -11,37 +11,8 @@ import ( "github.com/stretchr/testify/require" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -// MockAdapter mocks the required functionality from the Adapter for testing -type MockAdapter struct { - mock.Mock -} - -// Ensure MockAdapter implements AdapterInterface -var _ AdapterInterface = (*MockAdapter)(nil) - -func (m *MockAdapter) Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error { - args := m.Called(ctx, event, updater) - return args.Error(0) -} - -func (m *MockAdapter) Publish(ctx context.Context, event PublishAndRequestEventConfiguration) error { - args := m.Called(ctx, event) - return args.Error(0) -} - -func (m *MockAdapter) Request(ctx context.Context, event PublishAndRequestEventConfiguration, w io.Writer) error { - args := m.Called(ctx, event, w) - return args.Error(0) -} - -func (m *MockAdapter) Shutdown(ctx context.Context) error { - args := m.Called(ctx) - return args.Error(0) -} - func TestNatsPubSubDataSource(t *testing.T) { // Create event configuration with required fields engineEventConfig := &nodev1.EngineEventConfiguration{ @@ -82,7 +53,7 @@ func TestPubSubDataSourceWithMockAdapter(t *testing.T) { } // Create mock adapter - mockAdapter := new(MockAdapter) + mockAdapter := new(mockAdapter) // Configure mock expectations for Publish mockAdapter.On("Publish", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { @@ -128,7 +99,7 @@ func TestPubSubDataSource_GetResolveDataSource_WrongType(t *testing.T) { } // Create mock adapter - mockAdapter := new(MockAdapter) + mockAdapter := new(mockAdapter) // Create the data source with mock adapter pubsub := &PubSubDataSource{ @@ -279,7 +250,7 @@ func TestPubSubDataSource_RequestDataSource(t *testing.T) { } // Create mock adapter - mockAdapter := new(MockAdapter) + mockAdapter := new(mockAdapter) // Configure mock expectations for Request mockAdapter.On("Request", mock.Anything, mock.MatchedBy(func(event PublishAndRequestEventConfiguration) bool { From 0a6e8a01baede997824cc32d027d5fc3c5a7cc1d Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 14 Apr 2025 11:33:15 +0200 Subject: [PATCH 25/49] refactor: enhance NATS adapter with flush timeout and update test expectations for log entries --- router-tests/events/nats_events_test.go | 2 +- router-tests/structured_logging_test.go | 6 ++++-- router-tests/testenv/testenv.go | 4 ++++ router/core/factoryresolver.go | 3 ++- router/pkg/pubsub/nats/adapter.go | 8 ++++++++ 5 files changed, 19 insertions(+), 4 deletions(-) diff --git a/router-tests/events/nats_events_test.go b/router-tests/events/nats_events_test.go index 46fbf310ad..a8d06e3b22 100644 --- a/router-tests/events/nats_events_test.go +++ b/router-tests/events/nats_events_test.go @@ -136,7 +136,7 @@ func TestNatsEvents(t *testing.T) { natsLogs := xEnv.Observer().FilterMessageSnippet("Nats").All() require.Len(t, natsLogs, 4) providerIDFields := xEnv.Observer().FilterField(zap.String("provider_id", "my-nats")).All() - require.Len(t, providerIDFields, 1) + require.Len(t, providerIDFields, 3) }) }) diff --git a/router-tests/structured_logging_test.go b/router-tests/structured_logging_test.go index b6213327ba..67c87d5199 100644 --- a/router-tests/structured_logging_test.go +++ b/router-tests/structured_logging_test.go @@ -164,11 +164,13 @@ func TestRouterStartLogs(t *testing.T) { }, }, func(t *testing.T, xEnv *testenv.Environment) { logEntries := xEnv.Observer().All() - require.Len(t, logEntries, 11) + require.Len(t, logEntries, 15) natsLogs := xEnv.Observer().FilterMessageSnippet("Nats Event source enabled").All() require.Len(t, natsLogs, 4) + natsConnectedLogs := xEnv.Observer().FilterMessageSnippet("NATS connection established").All() + require.Len(t, natsConnectedLogs, 4) providerIDFields := xEnv.Observer().FilterField(zap.String("provider_id", "default")).All() - require.Len(t, providerIDFields, 1) + require.Len(t, providerIDFields, 3) kafkaLogs := xEnv.Observer().FilterMessageSnippet("Kafka Event source enabled").All() require.Len(t, kafkaLogs, 2) playgroundLog := xEnv.Observer().FilterMessage("Serving GraphQL playground") diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index f0a2003ba5..f974863b24 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -2098,6 +2098,10 @@ func subgraphOptions(ctx context.Context, t testing.TB, logger *zap.Logger, nats for _, sourceName := range demoNatsProviders { adapter, err := pubsubNats.NewAdapter(ctx, logger, natsData.Params[0].Url, natsData.Params[0].Opts, "hostname", "listenaddr") require.NoError(t, err) + require.NoError(t, adapter.Startup(ctx)) + t.Cleanup(func() { + require.NoError(t, adapter.Shutdown(context.Background())) + }) natsPubSubByProviderID[sourceName] = adapter } diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 6b684ed733..262751de01 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -420,7 +420,8 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod var providers []datasource.PubSubProvider dsMeta := l.dataSourceMetaData(in) - for _, providerFactory := range pubsub.GetProviderFactories() { + providersFactories := pubsub.GetProviderFactories() + for _, providerFactory := range providersFactories { provider, err := providerFactory( l.ctx, in, diff --git a/router/pkg/pubsub/nats/adapter.go b/router/pkg/pubsub/nats/adapter.go index 34a646b158..a355cb676a 100644 --- a/router/pkg/pubsub/nats/adapter.go +++ b/router/pkg/pubsub/nats/adapter.go @@ -41,6 +41,7 @@ type Adapter struct { routerListenAddr string url string opts []nats.Option + flushTimeout time.Duration } // getInstanceIdentifier returns an identifier for the current instance. @@ -250,6 +251,12 @@ func (p *Adapter) flush(ctx context.Context) error { if p.client == nil { return nil } + _, ok := ctx.Deadline() + if !ok { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, p.flushTimeout) + defer cancel() + } return p.client.FlushWithContext(ctx) } @@ -309,5 +316,6 @@ func NewAdapter(ctx context.Context, logger *zap.Logger, url string, opts []nats routerListenAddr: routerListenAddr, url: url, opts: opts, + flushTimeout: 10 * time.Second, }, nil } From 3aeba040957ba9e981e1807466a8e910a628bd1d Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 14 Apr 2025 11:55:21 +0200 Subject: [PATCH 26/49] test: log unexpected errors during WebSocket message read in tests --- router-tests/websocket_test.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/router-tests/websocket_test.go b/router-tests/websocket_test.go index b47d7a756f..b6a1b59066 100644 --- a/router-tests/websocket_test.go +++ b/router-tests/websocket_test.go @@ -2160,6 +2160,9 @@ func expectConnectAndReadCurrentTime(t *testing.T, xEnv *testenv.Environment) { err = testenv.WSReadJSON(t, conn, &msg) require.NoError(t, err) require.Equal(t, "1", msg.ID) + if msg.Type == "error" { + t.Logf("unexpected error on read: %s", string(msg.Payload)) + } require.Equal(t, "next", msg.Type) err = json.Unmarshal(msg.Payload, &payload) require.NoError(t, err) @@ -2169,6 +2172,9 @@ func expectConnectAndReadCurrentTime(t *testing.T, xEnv *testenv.Environment) { err = testenv.WSReadJSON(t, conn, &msg) require.NoError(t, err) require.Equal(t, "1", msg.ID) + if msg.Type == "error" { + t.Logf("unexpected error on read: %s", string(msg.Payload)) + } require.Equal(t, "next", msg.Type) err = json.Unmarshal(msg.Payload, &payload) require.NoError(t, err) @@ -2189,6 +2195,9 @@ func expectConnectAndReadCurrentTime(t *testing.T, xEnv *testenv.Environment) { err = testenv.WSReadJSON(t, conn, &complete) require.NoError(t, err) require.Equal(t, "1", complete.ID) + if complete.Type == "error" { + t.Logf("unexpected error on read: %s", string(complete.Payload)) + } require.Equal(t, "complete", complete.Type) err = conn.SetReadDeadline(time.Now().Add(1 * time.Second)) From f5e8e907268d3849aab0eaccc91ab5e3e318d221 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Mon, 14 Apr 2025 12:38:29 +0200 Subject: [PATCH 27/49] feat: add custom JSON unmarshalling for WebSocketMessage to improve error handling --- router-tests/testenv/testenv.go | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index f974863b24..b9a0e6d8da 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -1545,6 +1545,18 @@ type WebSocketMessage struct { Payload json.RawMessage `json:"payload,omitempty"` } +func (t *WebSocketMessage) UnmarshalJSON(data []byte) error { + err := json.Unmarshal(data, (*struct { + ID string `json:"id,omitempty"` + Type string `json:"type"` + Payload json.RawMessage `json:"payload,omitempty"` + })(t)) + if err != nil { + return fmt.Errorf("failed to unmarshal WebSocketMessage: %w (data: %s)", err, string(data)) + } + return nil +} + type GraphQLResponse struct { Data json.RawMessage `json:"data,omitempty"` Errors []GraphQLError `json:"errors,omitempty"` From ad96e3dc6d2a4735190358cfe57b1382790003f8 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Tue, 15 Apr 2025 09:31:23 +0200 Subject: [PATCH 28/49] refactor: replace custom JSON unmarshalling in WebSocketMessage with standard unmarshalling and improve error logging during message reads --- router-tests/testenv/testenv.go | 31 ++++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index b9a0e6d8da..4664bba11a 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -1545,18 +1545,6 @@ type WebSocketMessage struct { Payload json.RawMessage `json:"payload,omitempty"` } -func (t *WebSocketMessage) UnmarshalJSON(data []byte) error { - err := json.Unmarshal(data, (*struct { - ID string `json:"id,omitempty"` - Type string `json:"type"` - Payload json.RawMessage `json:"payload,omitempty"` - })(t)) - if err != nil { - return fmt.Errorf("failed to unmarshal WebSocketMessage: %w (data: %s)", err, string(data)) - } - return nil -} - type GraphQLResponse struct { Data json.RawMessage `json:"data,omitempty"` Errors []GraphQLError `json:"errors,omitempty"` @@ -1624,8 +1612,14 @@ func (e *Environment) InitGraphQLWebSocketConnection(header http.Header, query u }) require.NoError(e.t, err) var ack WebSocketMessage - err = conn.ReadJSON(&ack) - require.NoError(e.t, err) + _, payload, err := conn.ReadMessage() + if err != nil { + require.NoError(e.t, err) + } + if err := json.Unmarshal(payload, &ack); err != nil { + e.t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) + require.NoError(e.t, err) + } require.Equal(e.t, "connection_ack", ack.Type) return conn } @@ -1986,7 +1980,14 @@ func WSReadJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) { return err } - err = conn.ReadJSON(v) + _, payload, err := conn.ReadMessage() + if err != nil { + require.NoError(t, err) + } + if err := json.Unmarshal(payload, &v); err != nil { + t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) + require.NoError(t, err) + } // Reset the deadline to prevent future operations from timing out if resetErr := conn.SetReadDeadline(time.Time{}); resetErr != nil { From 85adb377bb1e5d49e5073e49d239dd929102a2e2 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 16 Apr 2025 08:40:44 +0200 Subject: [PATCH 29/49] fix: adapter base interface --- router/pkg/pubsub/README.md | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/router/pkg/pubsub/README.md b/router/pkg/pubsub/README.md index 76e9efcbc2..6862d311f8 100644 --- a/router/pkg/pubsub/README.md +++ b/router/pkg/pubsub/README.md @@ -30,6 +30,8 @@ The Adapter contains the logic that is actually calling the provider, usually it type AdapterInterface interface { Subscribe(ctx context.Context, event SubscriptionEventConfiguration, updater resolve.SubscriptionUpdater) error Publish(ctx context.Context, event PublishEventConfiguration) error + Startup(ctx context.Context) error + Shutdown(ctx context.Context) error } ``` @@ -69,4 +71,17 @@ To complete the `PubSubDataSource` implementation you should also add the engine So you have to implement the SubscriptionDataSource, a structure that implements all the methods needed by the interface `resolve.SubscriptionDataSource`, like the [kafka implementation](./kafka/engine_datasource.go). -And also, you have to implement the DataSource, a structure that implements all the methods needed by the interface `resolve.DataSource`, like `PublishDataSource` in the [kafka implementation](./kafka/pubsub_datasource.go). \ No newline at end of file +And also, you have to implement the DataSource, a structure that implements all the methods needed by the interface `resolve.DataSource`, like `PublishDataSource` in the [kafka implementation](./kafka/pubsub_datasource.go). + +## How to use the new PubSub Provider + +After you have implemented all the above, you can use your PubSub Provider by adding the following to your router config: + +```yaml +pubsub: + providers: + - name: provider-name + type: new-provider +``` + +But to use it in the schema you will have to work in the [composition](../../../composition) folder. \ No newline at end of file From 1bb2dfd60b5fd0b3d1461468227f5115fc8eb87f Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 16 Apr 2025 10:21:05 +0200 Subject: [PATCH 30/49] docs: update README for PubSub Provider usage section --- router/pkg/pubsub/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/router/pkg/pubsub/README.md b/router/pkg/pubsub/README.md index 6862d311f8..ebcc636424 100644 --- a/router/pkg/pubsub/README.md +++ b/router/pkg/pubsub/README.md @@ -73,7 +73,7 @@ So you have to implement the SubscriptionDataSource, a structure that implements And also, you have to implement the DataSource, a structure that implements all the methods needed by the interface `resolve.DataSource`, like `PublishDataSource` in the [kafka implementation](./kafka/pubsub_datasource.go). -## How to use the new PubSub Provider +# How to use the new PubSub Provider After you have implemented all the above, you can use your PubSub Provider by adding the following to your router config: From 43977a0dec1d0624fce940cc938a2e1bb2a99bd1 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 23 Apr 2025 08:56:17 +0200 Subject: [PATCH 31/49] refactor: remove unused EnginePubSubProviders struct and related import from graph_server.go --- router/core/graph_server.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/router/core/graph_server.go b/router/core/graph_server.go index c75a717b74..917f9088be 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -14,7 +14,6 @@ import ( "github.com/cespare/xxhash/v2" "github.com/wundergraph/cosmo/router/pkg/execution_config" rtrace "github.com/wundergraph/cosmo/router/pkg/trace" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/pubsub_datasource" otelmetric "go.opentelemetry.io/otel/metric" oteltrace "go.opentelemetry.io/otel/trace" @@ -61,11 +60,6 @@ type ( HealthChecks() health.Checker } - EnginePubSubProviders struct { - nats map[string]pubsub_datasource.NatsPubSub - kafka map[string]pubsub_datasource.KafkaPubSub - } - // graphServer is the swappable implementation of a Graph instance which is an HTTP mux with middlewares. // Everytime a schema is updated, the old graph server is shutdown and a new graph server is created. // For feature flags, a graphql server has multiple mux and is dynamically switched based on the feature flag header or cookie. From 13f8d33c7aa1e10aad63c4077df84aa7f49a5895 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 23 Apr 2025 09:22:42 +0200 Subject: [PATCH 32/49] refactor: update MoodHandler and NewSchema to include GetPubSubName parameter --- demo/pkg/subgraphs/mood/mood.go | 3 ++- demo/pkg/subgraphs/subgraphs.go | 4 ++-- router-tests/mcp_test.go | 2 ++ 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/demo/pkg/subgraphs/mood/mood.go b/demo/pkg/subgraphs/mood/mood.go index 8bce0d004b..efbc307be6 100644 --- a/demo/pkg/subgraphs/mood/mood.go +++ b/demo/pkg/subgraphs/mood/mood.go @@ -8,8 +8,9 @@ import ( "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood/subgraph/generated" ) -func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface) graphql.ExecutableSchema { +func NewSchema(natsPubSubByProviderID map[string]nats.AdapterInterface, getPubSubName func(string) string) graphql.ExecutableSchema { return generated.NewExecutableSchema(generated.Config{Resolvers: &subgraph.Resolver{ NatsPubSubByProviderID: natsPubSubByProviderID, + GetPubSubName: getPubSubName, }}) } diff --git a/demo/pkg/subgraphs/subgraphs.go b/demo/pkg/subgraphs/subgraphs.go index c5e91ab034..a66199e643 100644 --- a/demo/pkg/subgraphs/subgraphs.go +++ b/demo/pkg/subgraphs/subgraphs.go @@ -194,7 +194,7 @@ func AvailabilityHandler(opts *SubgraphOptions) http.Handler { } func MoodHandler(opts *SubgraphOptions) http.Handler { - return subgraphHandler(mood.NewSchema(opts.NatsPubSubByProviderID)) + return subgraphHandler(mood.NewSchema(opts.NatsPubSubByProviderID, opts.GetPubSubName)) } func CountriesHandler(opts *SubgraphOptions) http.Handler { @@ -260,7 +260,7 @@ func New(ctx context.Context, config *Config) (*Subgraphs, error) { if srv := newServer("availability", config.EnableDebug, config.Ports.Availability, availability.NewSchema(natsPubSubByProviderID, config.GetPubSubName)); srv != nil { servers = append(servers, srv) } - if srv := newServer("mood", config.EnableDebug, config.Ports.Mood, mood.NewSchema(natsPubSubByProviderID)); srv != nil { + if srv := newServer("mood", config.EnableDebug, config.Ports.Mood, mood.NewSchema(natsPubSubByProviderID, config.GetPubSubName)); srv != nil { servers = append(servers, srv) } if srv := newServer("countries", config.EnableDebug, config.Ports.Countries, countries.NewSchema(natsPubSubByProviderID)); srv != nil { diff --git a/router-tests/mcp_test.go b/router-tests/mcp_test.go index 1fea3e2660..9a9e078b68 100644 --- a/router-tests/mcp_test.go +++ b/router-tests/mcp_test.go @@ -179,6 +179,7 @@ func TestMCP(t *testing.T) { t.Run("Execute Query", func(t *testing.T) { t.Run("Execute operation of type query with valid input", func(t *testing.T) { testenv.Run(t, &testenv.Config{ + EnableNats: true, MCP: config.MCPConfiguration{ Enabled: true, }, @@ -235,6 +236,7 @@ func TestMCP(t *testing.T) { t.Run("Execute Mutation", func(t *testing.T) { t.Run("Execute operation of type mutation with valid input", func(t *testing.T) { testenv.Run(t, &testenv.Config{ + EnableNats: true, MCP: config.MCPConfiguration{ Enabled: true, }, From 5f957398ba9715d7468bf3ce8f2d11c9cf296d86 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 23 Apr 2025 09:30:44 +0200 Subject: [PATCH 33/49] refactor: enhance NATS publishing logic in UpdateMood resolver to handle missing providers --- demo/cmd/mood/main.go | 7 ++-- .../mood/subgraph/schema.resolvers.go | 32 ++++++++++++------- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/demo/cmd/mood/main.go b/demo/cmd/mood/main.go index 555de06584..96f7658c93 100644 --- a/demo/cmd/mood/main.go +++ b/demo/cmd/mood/main.go @@ -3,11 +3,12 @@ package main import ( "context" "fmt" - "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood" "log" "net/http" "os" + "github.com/wundergraph/cosmo/demo/pkg/subgraphs/mood" + "github.com/99designs/gqlgen/graphql" "github.com/99designs/gqlgen/graphql/handler/debug" "github.com/99designs/gqlgen/graphql/playground" @@ -31,7 +32,9 @@ func main() { port = defaultPort } - srv := subgraphs.NewDemoServer(mood.NewSchema(nil)) + srv := subgraphs.NewDemoServer(mood.NewSchema(nil, func(name string) string { + return name + })) srv.Use(&debug.Tracer{}) srv.Use(otelgqlgen.Middleware(otelgqlgen.WithCreateSpanFromFields(func(ctx *graphql.FieldContext) bool { diff --git a/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go b/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go index 3508b7d36d..8ab7c73941 100644 --- a/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go +++ b/demo/pkg/subgraphs/mood/subgraph/schema.resolvers.go @@ -18,21 +18,29 @@ func (r *mutationResolver) UpdateMood(ctx context.Context, employeeID int, mood storage.Set(employeeID, mood) myNatsTopic := r.GetPubSubName(fmt.Sprintf("employeeUpdated.%d", employeeID)) payload := fmt.Sprintf(`{"id":%d,"__typename": "Employee"}`, employeeID) - err := r.NatsPubSubByProviderID["default"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ - Subject: myNatsTopic, - Data: []byte(payload), - }) - if err != nil { - return nil, err + if r.NatsPubSubByProviderID["default"] != nil { + err := r.NatsPubSubByProviderID["default"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ + Subject: myNatsTopic, + Data: []byte(payload), + }) + if err != nil { + return nil, err + } + } else { + return nil, fmt.Errorf("no nats pubsub default provider found") } defaultTopic := r.GetPubSubName(fmt.Sprintf("employeeUpdatedMyNats.%d", employeeID)) - err = r.NatsPubSubByProviderID["my-nats"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ - Subject: defaultTopic, - Data: []byte(payload), - }) - if err != nil { - return nil, err + if r.NatsPubSubByProviderID["my-nats"] != nil { + err := r.NatsPubSubByProviderID["my-nats"].Publish(ctx, nats.PublishAndRequestEventConfiguration{ + Subject: defaultTopic, + Data: []byte(payload), + }) + if err != nil { + return nil, err + } + } else { + return nil, fmt.Errorf("no nats pubsub my-nats provider found") } return &model.Employee{ID: employeeID, CurrentMood: mood}, nil From 01e12bc143487e91ac06ab291b26b6aa8a005433 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 7 May 2025 14:11:13 +0200 Subject: [PATCH 34/49] chore: pass up providers instead of passing down a callback, small fixes --- router/core/executor.go | 11 +++++- router/core/factoryresolver.go | 39 ++++++++++--------- router/core/graph_server.go | 12 ++---- router/core/plan_generator.go | 7 ++-- router/pkg/plan_generator/plan_generator.go | 2 + router/pkg/pubsub/datasource/datasource.go | 10 ++--- router/pkg/pubsub/datasource/planner.go | 16 ++++---- router/pkg/pubsub/datasource/test_utils.go | 10 ++--- .../{utils/utils.go => eventdata/build.go} | 2 +- router/pkg/pubsub/kafka/pubsub_datasource.go | 10 ++--- .../pubsub/kafka/pubsub_datasource_test.go | 12 +++--- router/pkg/pubsub/nats/pubsub_datasource.go | 10 ++--- .../pkg/pubsub/nats/pubsub_datasource_test.go | 18 ++++----- 13 files changed, 82 insertions(+), 77 deletions(-) rename router/pkg/pubsub/{utils/utils.go => eventdata/build.go} (98%) diff --git a/router/core/executor.go b/router/core/executor.go index a4f4cb6738..ee334c2adc 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -34,7 +34,7 @@ type ExecutorConfigurationBuilder struct { subscriptionClientOptions *SubscriptionClientOptions instanceData InstanceData - addPubSubProviderCallback func(provider datasource.PubSubProvider) + providers []datasource.PubSubProvider } type Executor struct { @@ -213,7 +213,7 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con routerEngineCfg.Execution.EnableSingleFlight, routerEngineCfg.Execution.EnableNetPoll, b.instanceData, - ), b.logger, b.addPubSubProviderCallback) + ), b.logger) // this generates the plan config using the data source factories from the config package planConfig, err := loader.Load(engineConfig, subgraphs, routerEngineCfg) @@ -234,5 +234,12 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con planConfig.MinifySubgraphOperations = routerEngineCfg.Execution.MinifySubgraphOperations planConfig.EnableOperationNamePropagation = routerEngineCfg.Execution.EnableSubgraphFetchOperationName + + b.providers = loader.GetProviders() + return planConfig, nil } + +func (b *ExecutorConfigurationBuilder) GetProviders() []datasource.PubSubProvider { + return b.providers +} diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index dd524189c0..61cc7652ee 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -31,14 +31,14 @@ type Loader struct { ctx context.Context resolver FactoryResolver // includeInfo controls whether additional information like type usage and field usage is included in the plan de - includeInfo bool - logger *zap.Logger - addPubSubProviderCallback func(provider datasource.PubSubProvider) + includeInfo bool + logger *zap.Logger + providers []datasource.PubSubProvider } type InstanceData struct { - hostName string - listenAddress string + HostName string + ListenAddress string } type FactoryResolver interface { @@ -175,13 +175,13 @@ func (d *DefaultFactoryResolver) InstanceData() InstanceData { return d.instanceData } -func NewLoader(ctx context.Context, includeInfo bool, resolver FactoryResolver, logger *zap.Logger, addPubSubProviderCallback func(provider datasource.PubSubProvider)) *Loader { +func NewLoader(ctx context.Context, includeInfo bool, resolver FactoryResolver, logger *zap.Logger) *Loader { return &Loader{ - ctx: ctx, - resolver: resolver, - includeInfo: includeInfo, - logger: logger, - addPubSubProviderCallback: addPubSubProviderCallback, + ctx: ctx, + resolver: resolver, + includeInfo: includeInfo, + logger: logger, + providers: []datasource.PubSubProvider{}, } } @@ -194,6 +194,10 @@ func (l *Loader) LoadInternedString(engineConfig *nodev1.EngineConfiguration, st return s, nil } +func (l *Loader) GetProviders() []datasource.PubSubProvider { + return l.providers +} + type RouterEngineConfiguration struct { Execution config.EngineExecutionConfiguration Headers *config.HeaderRules @@ -431,8 +435,6 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod case nodev1.DataSourceKind_PUBSUB: var err error - var providers []datasource.PubSubProvider - dsMeta := l.dataSourceMetaData(in) providersFactories := pubsub.GetProviderFactories() for _, providerFactory := range providersFactories { @@ -442,23 +444,22 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod dsMeta, routerEngineConfig.Events, l.logger, - l.resolver.InstanceData().hostName, - l.resolver.InstanceData().listenAddress, + l.resolver.InstanceData().HostName, + l.resolver.InstanceData().ListenAddress, ) if err != nil { return nil, err } if provider != nil { - providers = append(providers, provider) - l.addPubSubProviderCallback(provider) + l.providers = append(l.providers, provider) } } out, err = plan.NewDataSourceConfiguration( in.Id, - datasource.NewFactory(l.ctx, routerEngineConfig.Events, providers), + datasource.NewFactory(l.ctx, routerEngineConfig.Events, l.providers), dsMeta, - providers, + l.providers, ) if err != nil { return nil, err diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 5981cca590..93d7b0c586 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -125,8 +125,8 @@ func newGraphServer(ctx context.Context, r *Router, routerConfig *nodev1.RouterC inFlightRequests: &atomic.Uint64{}, graphMuxList: make([]*graphMux, 0, 1), instanceData: InstanceData{ - hostName: r.hostName, - listenAddress: r.listenAddr, + HostName: r.hostName, + ListenAddress: r.listenAddr, }, storageProviders: &r.storageProviders, } @@ -922,12 +922,6 @@ func (s *graphServer) buildGraphMux(ctx context.Context, SubgraphErrorPropagation: s.subgraphErrorPropagation, } - var pubSubProviders []datasource.PubSubProvider - - addPubSubProviderCallback := func(provider datasource.PubSubProvider) { - pubSubProviders = append(pubSubProviders, provider) - } - // map[string]*http.Transport cannot be coerced into map[string]http.RoundTripper, unfortunately subgraphTippers := map[string]http.RoundTripper{} for subgraph, subgraphTransport := range s.subgraphTransports { @@ -966,7 +960,6 @@ func (s *graphServer) buildGraphMux(ctx context.Context, LocalhostFallbackInsideDocker: s.localhostFallbackInsideDocker, Logger: s.logger, }, - addPubSubProviderCallback: addPubSubProviderCallback, } executor, err := ecb.Build( @@ -986,6 +979,7 @@ func (s *graphServer) buildGraphMux(ctx context.Context, return nil, fmt.Errorf("failed to build plan configuration: %w", err) } + pubSubProviders := ecb.GetProviders() if len(pubSubProviders) > 0 { for _, provider := range pubSubProviders { if err := provider.Startup(ctx); err != nil { diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index 63aa5882ee..4420f70327 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -10,7 +10,6 @@ import ( log "github.com/jensneuse/abstractlogger" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/cosmo/router/pkg/pubsub/kafka" "github.com/wundergraph/cosmo/router/pkg/pubsub/nats" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" @@ -275,13 +274,15 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo graphql_datasource.WithNetPollConfiguration(netPollConfig), ) - ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + loader := NewLoader(ctx, false, &DefaultFactoryResolver{ engineCtx: ctx, httpClient: http.DefaultClient, streamingClient: http.DefaultClient, subscriptionClient: subscriptionClient, - }, logger, func(provider datasource.PubSubProvider) {}) + }, logger) // this generates the plan configuration using the data source factories from the config package planConfig, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &routerEngineConfig) diff --git a/router/pkg/plan_generator/plan_generator.go b/router/pkg/plan_generator/plan_generator.go index 530b453ea7..8b14858431 100644 --- a/router/pkg/plan_generator/plan_generator.go +++ b/router/pkg/plan_generator/plan_generator.go @@ -116,6 +116,8 @@ func PlanGenerator(ctx context.Context, cfg QueryPlanConfig) error { defer wg.Done() planner, err := pg.GetPlanner() if err != nil { + // if we fail to get the planner, we need to cancel the context to stop the other goroutines + // and return here to stop the current goroutine cancelError(fmt.Errorf("failed to get planner: %v", err)) return } diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index 458a44e709..5885a2cc2b 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -6,9 +6,9 @@ import ( ) type PubSubDataSource interface { - GetResolveDataSource() (resolve.DataSource, error) - GetResolveDataSourceInput(event []byte) (string, error) - GetEngineEventConfiguration() *nodev1.EngineEventConfiguration - GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) - GetResolveDataSourceSubscriptionInput() (string, error) + ResolveDataSource() (resolve.DataSource, error) + ResolveDataSourceInput(event []byte) (string, error) + EngineEventConfiguration() *nodev1.EngineEventConfiguration + ResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) + ResolveDataSourceSubscriptionInput() (string, error) } diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index 9f9ae3326a..81490ac162 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -4,7 +4,7 @@ import ( "fmt" "strings" - "github.com/wundergraph/cosmo/router/pkg/pubsub/utils" + "github.com/wundergraph/cosmo/router/pkg/pubsub/eventdata" "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" @@ -55,19 +55,19 @@ func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { var dataSource resolve.DataSource - dataSource, err := p.pubSubDataSource.GetResolveDataSource() + dataSource, err := p.pubSubDataSource.ResolveDataSource() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get data source: %w", err)) return resolve.FetchConfiguration{} } - event, err := utils.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) + event, err := eventdata.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) return resolve.FetchConfiguration{} } - input, err := p.pubSubDataSource.GetResolveDataSourceInput(event) + input, err := p.pubSubDataSource.ResolveDataSourceInput(event) if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) return resolve.FetchConfiguration{} @@ -78,7 +78,7 @@ func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { Variables: p.variables, DataSource: dataSource, PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.pubSubDataSource.GetEngineEventConfiguration().GetFieldName()}, + MergePath: []string{p.pubSubDataSource.EngineEventConfiguration().GetFieldName()}, }, } } @@ -89,13 +89,13 @@ func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { return plan.SubscriptionConfiguration{} } - dataSource, err := p.pubSubDataSource.GetResolveDataSourceSubscription() + dataSource, err := p.pubSubDataSource.ResolveDataSourceSubscription() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription: %w", err)) return plan.SubscriptionConfiguration{} } - input, err := p.pubSubDataSource.GetResolveDataSourceSubscriptionInput() + input, err := p.pubSubDataSource.ResolveDataSourceSubscriptionInput() if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source subscription input: %w", err)) return plan.SubscriptionConfiguration{} @@ -106,7 +106,7 @@ func (p *Planner) ConfigureSubscription() plan.SubscriptionConfiguration { Variables: p.variables, DataSource: dataSource, PostProcessing: resolve.PostProcessingConfiguration{ - MergePath: []string{p.pubSubDataSource.GetEngineEventConfiguration().GetFieldName()}, + MergePath: []string{p.pubSubDataSource.EngineEventConfiguration().GetFieldName()}, }, } } diff --git a/router/pkg/pubsub/datasource/test_utils.go b/router/pkg/pubsub/datasource/test_utils.go index 33500536ce..40b53c743e 100644 --- a/router/pkg/pubsub/datasource/test_utils.go +++ b/router/pkg/pubsub/datasource/test_utils.go @@ -12,17 +12,17 @@ import ( // This function can be used by other packages to test their PubSubDataSource implementations func VerifyPubSubDataSourceImplementation(t *testing.T, pubSub PubSubDataSource) { // Test GetEngineEventConfiguration - engineEventConfig := pubSub.GetEngineEventConfiguration() + engineEventConfig := pubSub.EngineEventConfiguration() require.NotNil(t, engineEventConfig, "Expected non-nil EngineEventConfiguration") // Test GetResolveDataSource - dataSource, err := pubSub.GetResolveDataSource() + dataSource, err := pubSub.ResolveDataSource() require.NoError(t, err, "Expected no error from GetResolveDataSource") require.NotNil(t, dataSource, "Expected non-nil DataSource") // Test GetResolveDataSourceInput with sample event data testEvent := []byte(`{"test":"data"}`) - input, err := pubSub.GetResolveDataSourceInput(testEvent) + input, err := pubSub.ResolveDataSourceInput(testEvent) require.NoError(t, err, "Expected no error from GetResolveDataSourceInput") assert.NotEmpty(t, input, "Expected non-empty input") @@ -32,12 +32,12 @@ func VerifyPubSubDataSourceImplementation(t *testing.T, pubSub PubSubDataSource) assert.NoError(t, err, "Expected valid JSON from GetResolveDataSourceInput") // Test GetResolveDataSourceSubscription - subscription, err := pubSub.GetResolveDataSourceSubscription() + subscription, err := pubSub.ResolveDataSourceSubscription() require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscription") require.NotNil(t, subscription, "Expected non-nil SubscriptionDataSource") // Test GetResolveDataSourceSubscriptionInput - subscriptionInput, err := pubSub.GetResolveDataSourceSubscriptionInput() + subscriptionInput, err := pubSub.ResolveDataSourceSubscriptionInput() require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") assert.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") diff --git a/router/pkg/pubsub/utils/utils.go b/router/pkg/pubsub/eventdata/build.go similarity index 98% rename from router/pkg/pubsub/utils/utils.go rename to router/pkg/pubsub/eventdata/build.go index c51775a37f..5d1b980849 100644 --- a/router/pkg/pubsub/utils/utils.go +++ b/router/pkg/pubsub/eventdata/build.go @@ -1,4 +1,4 @@ -package utils +package eventdata import ( "bytes" diff --git a/router/pkg/pubsub/kafka/pubsub_datasource.go b/router/pkg/pubsub/kafka/pubsub_datasource.go index f559faa372..1f749ab474 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource.go @@ -13,11 +13,11 @@ type PubSubDataSource struct { KafkaAdapter AdapterInterface } -func (c *PubSubDataSource) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { +func (c *PubSubDataSource) EngineEventConfiguration() *nodev1.EngineEventConfiguration { return c.EventConfiguration.GetEngineEventConfiguration() } -func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { +func (c *PubSubDataSource) ResolveDataSource() (resolve.DataSource, error) { var dataSource resolve.DataSource typeName := c.EventConfiguration.GetEngineEventConfiguration().GetType() @@ -33,7 +33,7 @@ func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { return dataSource, nil } -func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, error) { +func (c *PubSubDataSource) ResolveDataSourceInput(event []byte) (string, error) { topics := c.EventConfiguration.GetTopics() if len(topics) != 1 { @@ -53,13 +53,13 @@ func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, erro return evtCfg.MarshalJSONTemplate(), nil } -func (c *PubSubDataSource) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { +func (c *PubSubDataSource) ResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { return &SubscriptionDataSource{ pubSub: c.KafkaAdapter, }, nil } -func (c *PubSubDataSource) GetResolveDataSourceSubscriptionInput() (string, error) { +func (c *PubSubDataSource) ResolveDataSourceSubscriptionInput() (string, error) { providerId := c.GetProviderId() evtCfg := SubscriptionEventConfiguration{ ProviderID: providerId, diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go index 4388a637a8..cd780730a0 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource_test.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -67,11 +67,11 @@ func TestPubSubDataSourceWithMockAdapter(t *testing.T) { } // Get the data source - ds, err := pubsub.GetResolveDataSource() + ds, err := pubsub.ResolveDataSource() require.NoError(t, err) // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.NoError(t, err) // Call Load on the data source @@ -109,7 +109,7 @@ func TestPubSubDataSource_GetResolveDataSource_WrongType(t *testing.T) { } // Get the data source - ds, err := pubsub.GetResolveDataSource() + ds, err := pubsub.ResolveDataSource() require.Error(t, err) require.Nil(t, ds) } @@ -135,7 +135,7 @@ func TestPubSubDataSource_GetResolveDataSourceInput_MultipleTopics(t *testing.T) } // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.Error(t, err) require.Empty(t, input) } @@ -161,7 +161,7 @@ func TestPubSubDataSource_GetResolveDataSourceInput_NoTopics(t *testing.T) { } // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.Error(t, err) require.Empty(t, input) } @@ -189,7 +189,7 @@ func TestKafkaPubSubDataSourceMultiTopicSubscription(t *testing.T) { } // Test GetResolveDataSourceSubscriptionInput - subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + subscriptionInput, err := pubsub.ResolveDataSourceSubscriptionInput() require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") diff --git a/router/pkg/pubsub/nats/pubsub_datasource.go b/router/pkg/pubsub/nats/pubsub_datasource.go index ca8aae2506..ffbac39db6 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource.go +++ b/router/pkg/pubsub/nats/pubsub_datasource.go @@ -14,11 +14,11 @@ type PubSubDataSource struct { NatsAdapter AdapterInterface } -func (c *PubSubDataSource) GetEngineEventConfiguration() *nodev1.EngineEventConfiguration { +func (c *PubSubDataSource) EngineEventConfiguration() *nodev1.EngineEventConfiguration { return c.EventConfiguration.GetEngineEventConfiguration() } -func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { +func (c *PubSubDataSource) ResolveDataSource() (resolve.DataSource, error) { var dataSource resolve.DataSource typeName := c.EventConfiguration.GetEngineEventConfiguration().GetType() @@ -38,7 +38,7 @@ func (c *PubSubDataSource) GetResolveDataSource() (resolve.DataSource, error) { return dataSource, nil } -func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, error) { +func (c *PubSubDataSource) ResolveDataSourceInput(event []byte) (string, error) { subjects := c.EventConfiguration.GetSubjects() if len(subjects) != 1 { @@ -58,13 +58,13 @@ func (c *PubSubDataSource) GetResolveDataSourceInput(event []byte) (string, erro return evtCfg.MarshalJSONTemplate(), nil } -func (c *PubSubDataSource) GetResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { +func (c *PubSubDataSource) ResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) { return &SubscriptionSource{ pubSub: c.NatsAdapter, }, nil } -func (c *PubSubDataSource) GetResolveDataSourceSubscriptionInput() (string, error) { +func (c *PubSubDataSource) ResolveDataSourceSubscriptionInput() (string, error) { providerId := c.GetProviderId() evtCfg := SubscriptionEventConfiguration{ diff --git a/router/pkg/pubsub/nats/pubsub_datasource_test.go b/router/pkg/pubsub/nats/pubsub_datasource_test.go index d14a8721a6..64525e9d7a 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource_test.go +++ b/router/pkg/pubsub/nats/pubsub_datasource_test.go @@ -67,11 +67,11 @@ func TestPubSubDataSourceWithMockAdapter(t *testing.T) { } // Get the data source - ds, err := pubsub.GetResolveDataSource() + ds, err := pubsub.ResolveDataSource() require.NoError(t, err) // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.NoError(t, err) // Call Load on the data source @@ -108,7 +108,7 @@ func TestPubSubDataSource_GetResolveDataSource_WrongType(t *testing.T) { } // Get the data source - ds, err := pubsub.GetResolveDataSource() + ds, err := pubsub.ResolveDataSource() require.Error(t, err) require.Nil(t, ds) } @@ -133,7 +133,7 @@ func TestPubSubDataSource_GetResolveDataSourceInput_MultipleSubjects(t *testing. } // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.Error(t, err) require.Empty(t, input) } @@ -158,7 +158,7 @@ func TestPubSubDataSource_GetResolveDataSourceInput_NoSubjects(t *testing.T) { } // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.Error(t, err) require.Empty(t, input) } @@ -183,7 +183,7 @@ func TestNatsPubSubDataSourceMultiSubjectSubscription(t *testing.T) { } // Test GetResolveDataSourceSubscriptionInput - subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + subscriptionInput, err := pubsub.ResolveDataSourceSubscriptionInput() require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") @@ -221,7 +221,7 @@ func TestNatsPubSubDataSourceWithStreamConfiguration(t *testing.T) { } // Test GetResolveDataSourceSubscriptionInput with stream configuration - subscriptionInput, err := pubsub.GetResolveDataSourceSubscriptionInput() + subscriptionInput, err := pubsub.ResolveDataSourceSubscriptionInput() require.NoError(t, err, "Expected no error from GetResolveDataSourceSubscriptionInput") require.NotEmpty(t, subscriptionInput, "Expected non-empty subscription input") @@ -267,12 +267,12 @@ func TestPubSubDataSource_RequestDataSource(t *testing.T) { } // Get the data source - ds, err := pubsub.GetResolveDataSource() + ds, err := pubsub.ResolveDataSource() require.NoError(t, err) require.NotNil(t, ds) // Get the input - input, err := pubsub.GetResolveDataSourceInput([]byte(`{"test":"data"}`)) + input, err := pubsub.ResolveDataSourceInput([]byte(`{"test":"data"}`)) require.NoError(t, err) // Call Load on the data source From b5c73577d3f20820fe154dc7a4833ff7f93c1876 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 7 May 2025 14:13:31 +0200 Subject: [PATCH 35/49] fix: rename NatParams to NatsParams for consistency in pubsub.go --- router-tests/testenv/pubsub.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/router-tests/testenv/pubsub.go b/router-tests/testenv/pubsub.go index f141045ad2..53075005b1 100644 --- a/router-tests/testenv/pubsub.go +++ b/router-tests/testenv/pubsub.go @@ -8,20 +8,20 @@ import ( nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" ) -type NatParams struct { +type NatsParams struct { Opts []nats.Option Url string } type NatsData struct { Connections []*nats.Conn - Params []*NatParams + Params []*NatsParams } func setupNatsClients(t testing.TB) (*NatsData, error) { natsData := &NatsData{} for range demoNatsProviders { - param := &NatParams{ + param := &NatsParams{ Url: nats.DefaultURL, Opts: []nats.Option{ nats.MaxReconnects(10), From 8eb72eb356609d86fc810c3a4c81676fef1d7b17 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 7 May 2025 14:26:04 +0200 Subject: [PATCH 36/49] refactor: replace datasource package with pubsubtest for common PubSubDataSource tests --- router/pkg/pubsub/kafka/pubsub_datasource_test.go | 4 ++-- router/pkg/pubsub/nats/pubsub_datasource_test.go | 4 ++-- .../{datasource/test_utils.go => pubsubtest/pubsubtest.go} | 5 +++-- 3 files changed, 7 insertions(+), 6 deletions(-) rename router/pkg/pubsub/{datasource/test_utils.go => pubsubtest/pubsubtest.go} (91%) diff --git a/router/pkg/pubsub/kafka/pubsub_datasource_test.go b/router/pkg/pubsub/kafka/pubsub_datasource_test.go index cd780730a0..cd10a1eadb 100644 --- a/router/pkg/pubsub/kafka/pubsub_datasource_test.go +++ b/router/pkg/pubsub/kafka/pubsub_datasource_test.go @@ -9,7 +9,7 @@ import ( "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/pubsubtest" ) func TestKafkaPubSubDataSource(t *testing.T) { @@ -34,7 +34,7 @@ func TestKafkaPubSubDataSource(t *testing.T) { } // Run the standard test suite - datasource.VerifyPubSubDataSourceImplementation(t, pubsub) + pubsubtest.VerifyPubSubDataSourceImplementation(t, pubsub) } // TestPubSubDataSourceWithMockAdapter tests the PubSubDataSource with a mocked adapter diff --git a/router/pkg/pubsub/nats/pubsub_datasource_test.go b/router/pkg/pubsub/nats/pubsub_datasource_test.go index 64525e9d7a..e48639b0c9 100644 --- a/router/pkg/pubsub/nats/pubsub_datasource_test.go +++ b/router/pkg/pubsub/nats/pubsub_datasource_test.go @@ -10,7 +10,7 @@ import ( "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + "github.com/wundergraph/cosmo/router/pkg/pubsub/pubsubtest" ) func TestNatsPubSubDataSource(t *testing.T) { @@ -35,7 +35,7 @@ func TestNatsPubSubDataSource(t *testing.T) { } // Run the standard test suite - datasource.VerifyPubSubDataSourceImplementation(t, pubsub) + pubsubtest.VerifyPubSubDataSourceImplementation(t, pubsub) } func TestPubSubDataSourceWithMockAdapter(t *testing.T) { diff --git a/router/pkg/pubsub/datasource/test_utils.go b/router/pkg/pubsub/pubsubtest/pubsubtest.go similarity index 91% rename from router/pkg/pubsub/datasource/test_utils.go rename to router/pkg/pubsub/pubsubtest/pubsubtest.go index 40b53c743e..09ad3b25b0 100644 --- a/router/pkg/pubsub/datasource/test_utils.go +++ b/router/pkg/pubsub/pubsubtest/pubsubtest.go @@ -1,4 +1,4 @@ -package datasource +package pubsubtest import ( "encoding/json" @@ -6,11 +6,12 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" ) // VerifyPubSubDataSourceImplementation is a common test function to verify any PubSubDataSource implementation // This function can be used by other packages to test their PubSubDataSource implementations -func VerifyPubSubDataSourceImplementation(t *testing.T, pubSub PubSubDataSource) { +func VerifyPubSubDataSourceImplementation(t *testing.T, pubSub datasource.PubSubDataSource) { // Test GetEngineEventConfiguration engineEventConfig := pubSub.EngineEventConfiguration() require.NotNil(t, engineEventConfig, "Expected non-nil EngineEventConfiguration") From 3148ce30c83c041458639801bf9b3c0cb03eaf1f Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 7 May 2025 16:29:28 +0200 Subject: [PATCH 37/49] docs: enhance PubSubDataSource interface documentation with detailed method descriptions and implementation guidelines --- router/pkg/pubsub/datasource/datasource.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/router/pkg/pubsub/datasource/datasource.go b/router/pkg/pubsub/datasource/datasource.go index 5885a2cc2b..7107a9efe1 100644 --- a/router/pkg/pubsub/datasource/datasource.go +++ b/router/pkg/pubsub/datasource/datasource.go @@ -5,10 +5,26 @@ import ( "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) +// PubSubDataSource is the interface that all pubsub data sources must implement. +// It serves three main purposes: +// 1. Resolving the data source and subscription data source +// 2. Generating the appropriate input for these data sources +// 3. Providing access to the engine event configuration +// +// For detailed implementation guidelines, see: +// https://github.com/wundergraph/cosmo/blob/main/router/pkg/pubsub/README.md type PubSubDataSource interface { + // ResolveDataSource returns the engine DataSource implementation that contains + // methods which will be called by the Planner when resolving a field ResolveDataSource() (resolve.DataSource, error) + // ResolveDataSourceInput build the input that will be passed to the engine DataSource ResolveDataSourceInput(event []byte) (string, error) + // EngineEventConfiguration get the engine event configuration, contains the provider id, type, type name and field name EngineEventConfiguration() *nodev1.EngineEventConfiguration + // ResolveDataSourceSubscription returns the engine SubscriptionDataSource implementation + // that contains methods to start a subscription, which will be called by the Planner + // when a subscription is initiated ResolveDataSourceSubscription() (resolve.SubscriptionDataSource, error) + // ResolveDataSourceSubscriptionInput build the input that will be passed to the engine SubscriptionDataSource ResolveDataSourceSubscriptionInput() (string, error) } From 02c0a45230011e9d8ede794f493b603a087b499e Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 7 May 2025 18:15:10 +0200 Subject: [PATCH 38/49] fix: improve pubsub provider startup and shutdown processes with timeout handling --- router/core/graph_server.go | 90 +++++++++++++++++++++++++++++++------ 1 file changed, 77 insertions(+), 13 deletions(-) diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 93d7b0c586..27c292d035 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -979,15 +979,9 @@ func (s *graphServer) buildGraphMux(ctx context.Context, return nil, fmt.Errorf("failed to build plan configuration: %w", err) } - pubSubProviders := ecb.GetProviders() - if len(pubSubProviders) > 0 { - for _, provider := range pubSubProviders { - if err := provider.Startup(ctx); err != nil { - return nil, fmt.Errorf("failed to startup pubsub provider: %w", err) - } - } - - s.pubSubProviders = pubSubProviders + s.pubSubProviders = ecb.GetProviders() + if pubSubStartupErr := s.startupPubSubProviders(ctx); pubSubStartupErr != nil { + return nil, pubSubStartupErr } operationProcessor := NewOperationProcessor(OperationProcessorOptions{ @@ -1316,10 +1310,8 @@ func (s *graphServer) Shutdown(ctx context.Context) error { } } - for _, provider := range s.pubSubProviders { - if providerErr := provider.Shutdown(ctx); providerErr != nil { - finalErr = errors.Join(finalErr, providerErr) - } + if err := s.shutdownPubSubProviders(ctx); err != nil { + finalErr = errors.Join(finalErr, err) } // Shutdown all graphs muxes to release resources @@ -1341,6 +1333,78 @@ func (s *graphServer) Shutdown(ctx context.Context) error { return finalErr } +// startupPubSubProviders starts all pubsub providers +// It returns an error if any of the providers fail to start +// or if some providers takes to long to start +func (s *graphServer) startupPubSubProviders(ctx context.Context) error { + // Default timeout for pubsub provider startup + const defaultStartupTimeout = 5 * time.Second + + cancellableCtx, cancel := context.WithCancel(ctx) + defer cancel() + + for _, provider := range s.pubSubProviders { + startupDone := make(chan error, 1) + + go func(p datasource.PubSubProvider) { + startupDone <- p.Startup(cancellableCtx) + }(provider) + + timer := time.NewTimer(defaultStartupTimeout) + defer timer.Stop() + + select { + case err := <-startupDone: + if err != nil { + return fmt.Errorf("failed to startup pubsub provider within allowed time: %s", err.Error()) + } + case <-timer.C: + return fmt.Errorf("pubsub provider startup timed out after %s", defaultStartupTimeout) + } + } + + return nil +} + +// shutdownPubSubProviders shuts down all pubsub providers +// It returns an error if any of the providers fail to shutdown +// or if some providers takes to long to shutdown +func (s *graphServer) shutdownPubSubProviders(ctx context.Context) error { + // Default timeout for pubsub provider shutdown + const defaultShutdownTimeout = 5 * time.Second + cancellableCtx, cancel := context.WithCancel(ctx) + defer cancel() + + var finalErrs []error + var wg sync.WaitGroup + for _, provider := range s.pubSubProviders { + shutdownDone := make(chan error, 1) + + wg.Add(1) + go func(p datasource.PubSubProvider) { + shutdownDone <- p.Shutdown(cancellableCtx) + }(provider) + + timer := time.NewTimer(defaultShutdownTimeout) + defer timer.Stop() + + select { + case err := <-shutdownDone: + wg.Done() + if err != nil { + finalErrs = append(finalErrs, fmt.Errorf("failed to shutdown pubsub provider within allowed time: %s", err.Error())) + } + case <-timer.C: + wg.Done() + finalErrs = append(finalErrs, fmt.Errorf("pubsub provider shutdown timed out after %s", defaultShutdownTimeout)) + } + } + + wg.Wait() + + return errors.Join(finalErrs...) +} + func configureSubgraphOverwrites( engineConfig *nodev1.EngineConfiguration, configSubgraphs []*nodev1.Subgraph, From 711c37dc34489d62d781e49df2bfb13d59154848 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Thu, 8 May 2025 10:44:02 +0200 Subject: [PATCH 39/49] feat: add tests for Kafka and NATS startup/shutdown behavior with incorrect broker/URLs --- router-tests/events/kafka_events_test.go | 25 +++++++++++ router-tests/events/nats_events_test.go | 31 ++++++++++++++ router-tests/testenv/pubsub.go | 2 +- router-tests/testenv/testenv.go | 16 +++---- router-tests/testenv/waitinglistener.go | 54 ++++++++++++++++++++++++ 5 files changed, 119 insertions(+), 9 deletions(-) create mode 100644 router-tests/testenv/waitinglistener.go diff --git a/router-tests/events/kafka_events_test.go b/router-tests/events/kafka_events_test.go index 958722f0ef..06374cea2b 100644 --- a/router-tests/events/kafka_events_test.go +++ b/router-tests/events/kafka_events_test.go @@ -7,6 +7,7 @@ import ( "encoding/json" "fmt" "net/http" + "strconv" "sync/atomic" "testing" "time" @@ -1060,6 +1061,30 @@ func TestKafkaEvents(t *testing.T) { require.Equal(t, `{"employeeID":3,"update":{"name":"name test"}}`, string(records[0].Value)) }) }) + + t.Run("kafka startup and shutdown with wrong broker should not stop router from starting indefinitely", func(t *testing.T) { + t.Parallel() + + listener := testenv.NewWaitingListener(t, time.Second*10) + listener.Start() + defer listener.Close() + + // kafka client is lazy and will not connect to the broker until the first message is produced + // so the router will start even if the kafka connection fails + errRouter := testenv.RunWithError(t, &testenv.Config{ + RouterConfigJSONTemplate: testenv.ConfigWithEdfsKafkaJSONTemplate, + EnableKafka: true, + ModifyEventsConfiguration: func(config *config.EventsConfiguration) { + for i := range config.Providers.Kafka { + config.Providers.Kafka[i].Brokers = []string{"localhost:" + strconv.Itoa(listener.Port())} + } + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + t.Log("should be called") + }) + + assert.NoError(t, errRouter) + }) } func TestFlakyKafkaEvents(t *testing.T) { diff --git a/router-tests/events/nats_events_test.go b/router-tests/events/nats_events_test.go index 989a9a4106..09bca3c995 100644 --- a/router-tests/events/nats_events_test.go +++ b/router-tests/events/nats_events_test.go @@ -8,6 +8,7 @@ import ( "io" "net/http" "net/url" + "strconv" "sync/atomic" "testing" "time" @@ -1807,6 +1808,36 @@ func TestNatsEvents(t *testing.T) { assert.Eventually(t, completed.Load, NatsWaitTimeout, time.Millisecond*100) }) }) + + t.Run("NATS startup and shutdown with wrong URLs should not stop router from starting indefinitely", func(t *testing.T) { + t.Parallel() + + listener := testenv.NewWaitingListener(t, time.Second*10) + listener.Start() + defer listener.Close() + + // kafka client is lazy and will not connect to the broker until the first message is produced + // so the router will start even if the kafka connection fails + errRouter := testenv.RunWithError(t, &testenv.Config{ + RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate, + EnableNats: false, + ModifyEventsConfiguration: func(cfg *config.EventsConfiguration) { + url := "nats://127.0.0.1:" + strconv.Itoa(listener.Port()) + natsEventSources := make([]config.NatsEventSource, len(testenv.DemoNatsProviders)) + for _, sourceName := range testenv.DemoNatsProviders { + natsEventSources = append(natsEventSources, config.NatsEventSource{ + ID: sourceName, + URL: url, + }) + } + cfg.Providers.Nats = natsEventSources + }, + }, func(t *testing.T, xEnv *testenv.Environment) { + assert.Fail(t, "Should not be called") + }) + + assert.Error(t, errRouter) + }) } func TestFlakyNatsEvents(t *testing.T) { diff --git a/router-tests/testenv/pubsub.go b/router-tests/testenv/pubsub.go index 53075005b1..4011c231d4 100644 --- a/router-tests/testenv/pubsub.go +++ b/router-tests/testenv/pubsub.go @@ -20,7 +20,7 @@ type NatsData struct { func setupNatsClients(t testing.TB) (*NatsData, error) { natsData := &NatsData{} - for range demoNatsProviders { + for range DemoNatsProviders { param := &NatsParams{ Url: nats.DefaultURL, Opts: []nats.Option{ diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index 5adf63ccde..73094b17fa 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -80,8 +80,8 @@ var ( ConfigWithEdfsKafkaJSONTemplate string //go:embed testdata/configWithEdfsNats.json ConfigWithEdfsNatsJSONTemplate string - demoNatsProviders = []string{natsDefaultSourceName, myNatsProviderID} - demoKafkaProviders = []string{myKafkaProviderID} + DemoNatsProviders = []string{natsDefaultSourceName, myNatsProviderID} + DemoKafkaProviders = []string{myKafkaProviderID} ) func init() { @@ -843,18 +843,18 @@ func configureRouter(listenerAddr string, testConfig *Config, routerConfig *node testConfig.ModifySubgraphErrorPropagation(&cfg.SubgraphErrorPropagation) } - natsEventSources := make([]config.NatsEventSource, len(demoNatsProviders)) - kafkaEventSources := make([]config.KafkaEventSource, len(demoKafkaProviders)) + natsEventSources := make([]config.NatsEventSource, len(DemoNatsProviders)) + kafkaEventSources := make([]config.KafkaEventSource, len(DemoKafkaProviders)) if natsData != nil { - for _, sourceName := range demoNatsProviders { + for _, sourceName := range DemoNatsProviders { natsEventSources = append(natsEventSources, config.NatsEventSource{ ID: sourceName, URL: nats.DefaultURL, }) } } - for _, sourceName := range demoKafkaProviders { + for _, sourceName := range DemoKafkaProviders { kafkaEventSources = append(kafkaEventSources, config.KafkaEventSource{ ID: sourceName, Brokers: testConfig.KafkaSeeds, @@ -2250,8 +2250,8 @@ func subgraphOptions(ctx context.Context, t testing.TB, logger *zap.Logger, nats GetPubSubName: pubSubName, } } - natsPubSubByProviderID := make(map[string]pubsubNats.AdapterInterface, len(demoNatsProviders)) - for _, sourceName := range demoNatsProviders { + natsPubSubByProviderID := make(map[string]pubsubNats.AdapterInterface, len(DemoNatsProviders)) + for _, sourceName := range DemoNatsProviders { adapter, err := pubsubNats.NewAdapter(ctx, logger, natsData.Params[0].Url, natsData.Params[0].Opts, "hostname", "listenaddr") require.NoError(t, err) require.NoError(t, adapter.Startup(ctx)) diff --git a/router-tests/testenv/waitinglistener.go b/router-tests/testenv/waitinglistener.go new file mode 100644 index 0000000000..25aa32c80f --- /dev/null +++ b/router-tests/testenv/waitinglistener.go @@ -0,0 +1,54 @@ +package testenv + +import ( + "context" + "net" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type WaitingListener struct { + cancel context.CancelFunc + listener *net.Listener + waitTime time.Duration + port int +} + +func (l *WaitingListener) Close() error { + l.cancel() + return (*l.listener).Close() +} + +func (l *WaitingListener) Start() { + go func() { + for { + conn, err := (*l.listener).Accept() + if err != nil { + return + } + time.Sleep(l.waitTime) + conn.Close() + } + }() +} + +func (l *WaitingListener) Port() int { + return l.port +} + +func NewWaitingListener(t *testing.T, waitTime time.Duration) (wl *WaitingListener) { + ctx, cancel := context.WithCancel(context.Background()) + var lc net.ListenConfig + listener, err := lc.Listen(ctx, "tcp", "127.0.0.1:0") + require.NoError(t, err) + + wl = &WaitingListener{ + cancel: cancel, + listener: &listener, + waitTime: waitTime, + port: listener.Addr().(*net.TCPAddr).Port, + } + return wl +} From 25e15aa3578ce827696ce2437b7b728987cad408 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Thu, 8 May 2025 11:29:24 +0200 Subject: [PATCH 40/49] refactor: extract JSON reading and checking logic into a reusable function for WebSocket messages --- router-tests/testenv/testenv.go | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index 73094b17fa..18588d5dfc 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -1743,6 +1743,18 @@ func (e *Environment) GraphQLWebsocketDialWithRetry(header http.Header, query ur return nil, nil, err } +func ReadAndCheckJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) { + _, payload, err := conn.ReadMessage() + if err != nil { + return err + } + if err := json.Unmarshal(payload, &v); err != nil { + t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) + return err + } + return nil +} + func (e *Environment) InitGraphQLWebSocketConnection(header http.Header, query url.Values, initialPayload json.RawMessage) *websocket.Conn { conn, _, err := e.GraphQLWebsocketDialWithRetry(header, query) require.NoError(e.t, err) @@ -1755,14 +1767,7 @@ func (e *Environment) InitGraphQLWebSocketConnection(header http.Header, query u }) require.NoError(e.t, err) var ack WebSocketMessage - _, payload, err := conn.ReadMessage() - if err != nil { - require.NoError(e.t, err) - } - if err := json.Unmarshal(payload, &ack); err != nil { - e.t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) - require.NoError(e.t, err) - } + require.NoError(e.t, ReadAndCheckJSON(e.t, conn, &ack)) require.Equal(e.t, "connection_ack", ack.Type) return conn } @@ -2123,14 +2128,7 @@ func WSReadJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) { return err } - _, payload, err := conn.ReadMessage() - if err != nil { - require.NoError(t, err) - } - if err := json.Unmarshal(payload, &v); err != nil { - t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) - require.NoError(t, err) - } + require.NoError(t, ReadAndCheckJSON(t, conn, v)) // Reset the deadline to prevent future operations from timing out if resetErr := conn.SetReadDeadline(time.Time{}); resetErr != nil { From 18745214de19dd08fead0227e957965bd6bb64e0 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Thu, 8 May 2025 11:30:20 +0200 Subject: [PATCH 41/49] refactor: move ReadAndCheckJSON function to improve code organization and avoid duplication --- router-tests/testenv/testenv.go | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/router-tests/testenv/testenv.go b/router-tests/testenv/testenv.go index 18588d5dfc..a77178ef5b 100644 --- a/router-tests/testenv/testenv.go +++ b/router-tests/testenv/testenv.go @@ -1149,6 +1149,18 @@ func gqlURL(srv *httptest.Server) string { return path } +func ReadAndCheckJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) { + _, payload, err := conn.ReadMessage() + if err != nil { + return err + } + if err := json.Unmarshal(payload, &v); err != nil { + t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) + return err + } + return nil +} + type Environment struct { t testing.TB cfg *Config @@ -1743,18 +1755,6 @@ func (e *Environment) GraphQLWebsocketDialWithRetry(header http.Header, query ur return nil, nil, err } -func ReadAndCheckJSON(t testing.TB, conn *websocket.Conn, v interface{}) (err error) { - _, payload, err := conn.ReadMessage() - if err != nil { - return err - } - if err := json.Unmarshal(payload, &v); err != nil { - t.Logf("Failed to decode WebSocket message. Raw payload: %s", string(payload)) - return err - } - return nil -} - func (e *Environment) InitGraphQLWebSocketConnection(header http.Header, query url.Values, initialPayload json.RawMessage) *websocket.Conn { conn, _, err := e.GraphQLWebsocketDialWithRetry(header, query) require.NoError(e.t, err) From 92b91d7ad07ca2f11fa291976f14b8e3f4188477 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 08:55:29 +0200 Subject: [PATCH 42/49] refactor: remove outdated comment regarding Kafka client connection behavior in NATS events test --- router-tests/events/nats_events_test.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/router-tests/events/nats_events_test.go b/router-tests/events/nats_events_test.go index 09bca3c995..531bcee18a 100644 --- a/router-tests/events/nats_events_test.go +++ b/router-tests/events/nats_events_test.go @@ -1816,8 +1816,6 @@ func TestNatsEvents(t *testing.T) { listener.Start() defer listener.Close() - // kafka client is lazy and will not connect to the broker until the first message is produced - // so the router will start even if the kafka connection fails errRouter := testenv.RunWithError(t, &testenv.Config{ RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate, EnableNats: false, From 802659a7f16e5a5baf036a767e14db3434ed1e9b Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 09:12:35 +0200 Subject: [PATCH 43/49] refactor: rename datasource import to pubsub_datasource for clarity in executor.go --- router/core/executor.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/router/core/executor.go b/router/core/executor.go index ee334c2adc..ed66fe2018 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -18,7 +18,7 @@ import ( nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" "github.com/wundergraph/cosmo/router/pkg/config" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + pubsub_datasource "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" ) type ExecutorConfigurationBuilder struct { @@ -34,7 +34,7 @@ type ExecutorConfigurationBuilder struct { subscriptionClientOptions *SubscriptionClientOptions instanceData InstanceData - providers []datasource.PubSubProvider + providers []pubsub_datasource.PubSubProvider } type Executor struct { @@ -240,6 +240,6 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con return planConfig, nil } -func (b *ExecutorConfigurationBuilder) GetProviders() []datasource.PubSubProvider { +func (b *ExecutorConfigurationBuilder) GetProviders() []pubsub_datasource.PubSubProvider { return b.providers } From 86fada30c42e2832532b855d5ed4f86c17724eca Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 09:14:27 +0200 Subject: [PATCH 44/49] fix: instanceData not set in the structure --- router/core/factoryresolver.go | 1 + 1 file changed, 1 insertion(+) diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 61cc7652ee..5dcb11f9d4 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -156,6 +156,7 @@ func NewDefaultFactoryResolver( httpClient: defaultHTTPClient, subgraphHTTPClients: subgraphHTTPClients, + instanceData: instanceData, } } From ee87b59a4b6355bd568f9d2a5fbd34ed708c00c1 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 10:03:37 +0200 Subject: [PATCH 45/49] refactor: streamline pubsub provider startup and shutdown with unified timeout handling --- router/core/graph_server.go | 75 +++++++++++++------------------------ 1 file changed, 27 insertions(+), 48 deletions(-) diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 27c292d035..092c714a20 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -31,6 +31,7 @@ import ( "go.uber.org/zap" "go.uber.org/zap/zapcore" "golang.org/x/exp/maps" + "golang.org/x/sync/errgroup" "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/common" nodev1 "github.com/wundergraph/cosmo/router/gen/proto/wg/cosmo/node/v1" @@ -1340,30 +1341,9 @@ func (s *graphServer) startupPubSubProviders(ctx context.Context) error { // Default timeout for pubsub provider startup const defaultStartupTimeout = 5 * time.Second - cancellableCtx, cancel := context.WithCancel(ctx) - defer cancel() - - for _, provider := range s.pubSubProviders { - startupDone := make(chan error, 1) - - go func(p datasource.PubSubProvider) { - startupDone <- p.Startup(cancellableCtx) - }(provider) - - timer := time.NewTimer(defaultStartupTimeout) - defer timer.Stop() - - select { - case err := <-startupDone: - if err != nil { - return fmt.Errorf("failed to startup pubsub provider within allowed time: %s", err.Error()) - } - case <-timer.C: - return fmt.Errorf("pubsub provider startup timed out after %s", defaultStartupTimeout) - } - } - - return nil + return s.providersActionWithTimeout(ctx, func(ctx context.Context, provider datasource.PubSubProvider) error { + return provider.Startup(ctx) + }, defaultStartupTimeout, "pubsub provider startup timed out") } // shutdownPubSubProviders shuts down all pubsub providers @@ -1372,37 +1352,36 @@ func (s *graphServer) startupPubSubProviders(ctx context.Context) error { func (s *graphServer) shutdownPubSubProviders(ctx context.Context) error { // Default timeout for pubsub provider shutdown const defaultShutdownTimeout = 5 * time.Second - cancellableCtx, cancel := context.WithCancel(ctx) - defer cancel() - var finalErrs []error - var wg sync.WaitGroup - for _, provider := range s.pubSubProviders { - shutdownDone := make(chan error, 1) + return s.providersActionWithTimeout(ctx, func(ctx context.Context, provider datasource.PubSubProvider) error { + return provider.Shutdown(ctx) + }, defaultShutdownTimeout, "pubsub provider shutdown timed out") +} - wg.Add(1) - go func(p datasource.PubSubProvider) { - shutdownDone <- p.Shutdown(cancellableCtx) - }(provider) +func (s *graphServer) providersActionWithTimeout(ctx context.Context, action func(ctx context.Context, provider datasource.PubSubProvider) error, timeout time.Duration, timeoutMessage string) error { + cancellableCtx, cancel := context.WithCancel(ctx) + defer cancel() - timer := time.NewTimer(defaultShutdownTimeout) - defer timer.Stop() + timer := time.NewTimer(timeout) + defer timer.Stop() - select { - case err := <-shutdownDone: - wg.Done() - if err != nil { - finalErrs = append(finalErrs, fmt.Errorf("failed to shutdown pubsub provider within allowed time: %s", err.Error())) + providersGroup := new(errgroup.Group) + for _, provider := range s.pubSubProviders { + providersGroup.Go(func() error { + actionDone := make(chan error, 1) + go func() { + actionDone <- action(cancellableCtx, provider) + }() + select { + case err := <-actionDone: + return err + case <-timer.C: + return fmt.Errorf(timeoutMessage, timeout) } - case <-timer.C: - wg.Done() - finalErrs = append(finalErrs, fmt.Errorf("pubsub provider shutdown timed out after %s", defaultShutdownTimeout)) - } + }) } - wg.Wait() - - return errors.Join(finalErrs...) + return providersGroup.Wait() } func configureSubgraphOverwrites( From 39517e819f96e8121215111dc139133138712d42 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 10:31:45 +0200 Subject: [PATCH 46/49] fix: replace timeout error formatting with errors.New for consistency in providersActionWithTimeout --- router/core/graph_server.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/router/core/graph_server.go b/router/core/graph_server.go index 092c714a20..d6cda32a13 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -1376,7 +1376,7 @@ func (s *graphServer) providersActionWithTimeout(ctx context.Context, action fun case err := <-actionDone: return err case <-timer.C: - return fmt.Errorf(timeoutMessage, timeout) + return errors.New(timeoutMessage) } }) } From 914f21c50140dbf051d1f266eb5fdf9d39df0584 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 11:44:44 +0200 Subject: [PATCH 47/49] refactor: update ExecutorConfigurationBuilder and Loader to return PubSubProviders in error handling --- router/core/executor.go | 34 +++++++++++----------------- router/core/factoryresolver.go | 41 ++++++++++++++++------------------ router/core/graph_server.go | 4 ++-- router/core/plan_generator.go | 2 +- 4 files changed, 35 insertions(+), 46 deletions(-) diff --git a/router/core/executor.go b/router/core/executor.go index ed66fe2018..0958ee147a 100644 --- a/router/core/executor.go +++ b/router/core/executor.go @@ -33,8 +33,6 @@ type ExecutorConfigurationBuilder struct { subscriptionClientOptions *SubscriptionClientOptions instanceData InstanceData - - providers []pubsub_datasource.PubSubProvider } type Executor struct { @@ -60,10 +58,10 @@ type ExecutorBuildOptions struct { InstanceData InstanceData } -func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *ExecutorBuildOptions) (*Executor, error) { - planConfig, err := b.buildPlannerConfiguration(ctx, opts.EngineConfig, opts.Subgraphs, opts.RouterEngineConfig) +func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *ExecutorBuildOptions) (*Executor, []pubsub_datasource.PubSubProvider, error) { + planConfig, providers, err := b.buildPlannerConfiguration(ctx, opts.EngineConfig, opts.Subgraphs, opts.RouterEngineConfig) if err != nil { - return nil, fmt.Errorf("failed to build planner configuration: %w", err) + return nil, nil, fmt.Errorf("failed to build planner configuration: %w", err) } options := resolve.ResolverOptions{ @@ -127,7 +125,7 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor routerSchemaDefinition, report = astparser.ParseGraphqlDocumentString(opts.EngineConfig.GraphqlSchema) if report.HasErrors() { - return nil, fmt.Errorf("failed to parse graphql schema from engine config: %w", report) + return nil, providers, fmt.Errorf("failed to parse graphql schema from engine config: %w", report) } // we need to merge the base schema, it contains the __schema and __type queries, // as well as built-in scalars like Int, String, etc... @@ -135,7 +133,7 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor // the engine needs to have them defined, otherwise it cannot resolve such fields err = asttransform.MergeDefinitionWithBaseSchema(&routerSchemaDefinition) if err != nil { - return nil, fmt.Errorf("failed to merge graphql schema with base schema: %w", err) + return nil, providers, fmt.Errorf("failed to merge graphql schema with base schema: %w", err) } if clientSchemaStr := opts.EngineConfig.GetGraphqlClientSchema(); clientSchemaStr != "" { @@ -144,11 +142,11 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor clientSchema, report := astparser.ParseGraphqlDocumentString(clientSchemaStr) if report.HasErrors() { - return nil, fmt.Errorf("failed to parse graphql client schema from engine config: %w", report) + return nil, providers, fmt.Errorf("failed to parse graphql client schema from engine config: %w", report) } err = asttransform.MergeDefinitionWithBaseSchema(&clientSchema) if err != nil { - return nil, fmt.Errorf("failed to merge graphql client schema with base schema: %w", err) + return nil, providers, fmt.Errorf("failed to merge graphql client schema with base schema: %w", err) } clientSchemaDefinition = &clientSchema } else { @@ -164,7 +162,7 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor // datasource is attached to Query.__schema, Query.__type, __Type.fields and __Type.enumValues fields introspectionFactory, err := introspection_datasource.NewIntrospectionConfigFactory(clientSchemaDefinition) if err != nil { - return nil, fmt.Errorf("failed to create introspection config factory: %w", err) + return nil, providers, fmt.Errorf("failed to create introspection config factory: %w", err) } fieldConfigs := introspectionFactory.BuildFieldConfigurations() // we need to add these fields to the config @@ -195,10 +193,10 @@ func (b *ExecutorConfigurationBuilder) Build(ctx context.Context, opts *Executor Resolver: resolver, RenameTypeNames: renameTypeNames, TrackUsageInfo: b.trackUsageInfo, - }, nil + }, providers, nil } -func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineCfg *RouterEngineConfiguration) (*plan.Configuration, error) { +func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Context, engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineCfg *RouterEngineConfiguration) (*plan.Configuration, []pubsub_datasource.PubSubProvider, error) { // this loader is used to take the engine config and create a plan config // the plan config is what the engine uses to turn a GraphQL Request into an execution plan // the plan config is stateful as it carries connection pools and other things @@ -216,9 +214,9 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con ), b.logger) // this generates the plan config using the data source factories from the config package - planConfig, err := loader.Load(engineConfig, subgraphs, routerEngineCfg) + planConfig, providers, err := loader.Load(engineConfig, subgraphs, routerEngineCfg) if err != nil { - return nil, fmt.Errorf("failed to load configuration: %w", err) + return nil, nil, fmt.Errorf("failed to load configuration: %w", err) } debug := &routerEngineCfg.Execution.Debug planConfig.Debug = plan.DebugConfiguration{ @@ -235,11 +233,5 @@ func (b *ExecutorConfigurationBuilder) buildPlannerConfiguration(ctx context.Con planConfig.EnableOperationNamePropagation = routerEngineCfg.Execution.EnableSubgraphFetchOperationName - b.providers = loader.GetProviders() - - return planConfig, nil -} - -func (b *ExecutorConfigurationBuilder) GetProviders() []pubsub_datasource.PubSubProvider { - return b.providers + return planConfig, providers, nil } diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index 5dcb11f9d4..feb340890d 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -12,6 +12,7 @@ import ( "github.com/buger/jsonparser" "github.com/wundergraph/cosmo/router/pkg/pubsub" "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" + pubsub_datasource "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" "github.com/wundergraph/cosmo/router/pkg/config" @@ -33,7 +34,6 @@ type Loader struct { // includeInfo controls whether additional information like type usage and field usage is included in the plan de includeInfo bool logger *zap.Logger - providers []datasource.PubSubProvider } type InstanceData struct { @@ -182,7 +182,6 @@ func NewLoader(ctx context.Context, includeInfo bool, resolver FactoryResolver, resolver: resolver, includeInfo: includeInfo, logger: logger, - providers: []datasource.PubSubProvider{}, } } @@ -195,10 +194,6 @@ func (l *Loader) LoadInternedString(engineConfig *nodev1.EngineConfiguration, st return s, nil } -func (l *Loader) GetProviders() []datasource.PubSubProvider { - return l.providers -} - type RouterEngineConfiguration struct { Execution config.EngineExecutionConfiguration Headers *config.HeaderRules @@ -259,7 +254,7 @@ func mapProtoFilterToPlanFilter(input *nodev1.SubscriptionFilterCondition, outpu return nil } -func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineConfig *RouterEngineConfiguration) (*plan.Configuration, error) { +func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nodev1.Subgraph, routerEngineConfig *RouterEngineConfiguration) (*plan.Configuration, []pubsub_datasource.PubSubProvider, error) { var outConfig plan.Configuration // attach field usage information to the plan outConfig.DefaultFlushIntervalMillis = engineConfig.DefaultFlushInterval @@ -294,6 +289,8 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod }) } + var providers []pubsub_datasource.PubSubProvider + for _, in := range engineConfig.DatasourceConfigurations { var out plan.DataSource @@ -301,7 +298,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod case nodev1.DataSourceKind_STATIC: factory, err := l.resolver.ResolveStaticFactory() if err != nil { - return nil, err + return nil, providers, err } out, err = plan.NewDataSourceConfiguration[staticdatasource.Configuration]( @@ -313,7 +310,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod }, ) if err != nil { - return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) + return nil, providers, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) } case nodev1.DataSourceKind_GRAPHQL: @@ -341,7 +338,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod graphqlSchema, err := l.LoadInternedString(engineConfig, in.CustomGraphql.GetUpstreamSchema()) if err != nil { - return nil, fmt.Errorf("could not load GraphQL schema for data source %s: %w", in.Id, err) + return nil, providers, fmt.Errorf("could not load GraphQL schema for data source %s: %w", in.Id, err) } var subscriptionUseSSE bool @@ -380,7 +377,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod dataSourceRules := FetchURLRules(routerEngineConfig.Headers, subgraphs, subscriptionUrl) forwardedClientHeaders, forwardedClientRegexps, err := PropagatedHeaders(dataSourceRules) if err != nil { - return nil, fmt.Errorf("error parsing header rules for data source %s: %w", in.Id, err) + return nil, providers, fmt.Errorf("error parsing header rules for data source %s: %w", in.Id, err) } schemaConfiguration, err := graphql_datasource.NewSchemaConfiguration( @@ -391,7 +388,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod }, ) if err != nil { - return nil, fmt.Errorf("error creating schema configuration for data source %s: %w", in.Id, err) + return nil, providers, fmt.Errorf("error creating schema configuration for data source %s: %w", in.Id, err) } customConfiguration, err := graphql_datasource.NewConfiguration(graphql_datasource.ConfigurationInput{ @@ -412,14 +409,14 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod CustomScalarTypeFields: customScalarTypeFields, }) if err != nil { - return nil, fmt.Errorf("error creating custom configuration for data source %s: %w", in.Id, err) + return nil, providers, fmt.Errorf("error creating custom configuration for data source %s: %w", in.Id, err) } dataSourceName := l.subgraphName(subgraphs, in.Id) factory, err := l.resolver.ResolveGraphqlFactory(dataSourceName) if err != nil { - return nil, err + return nil, providers, err } out, err = plan.NewDataSourceConfigurationWithName[graphql_datasource.Configuration]( @@ -430,7 +427,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod customConfiguration, ) if err != nil { - return nil, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) + return nil, providers, fmt.Errorf("error creating data source configuration for data source %s: %w", in.Id, err) } case nodev1.DataSourceKind_PUBSUB: @@ -449,29 +446,29 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod l.resolver.InstanceData().ListenAddress, ) if err != nil { - return nil, err + return nil, providers, err } if provider != nil { - l.providers = append(l.providers, provider) + providers = append(providers, provider) } } out, err = plan.NewDataSourceConfiguration( in.Id, - datasource.NewFactory(l.ctx, routerEngineConfig.Events, l.providers), + datasource.NewFactory(l.ctx, routerEngineConfig.Events, providers), dsMeta, - l.providers, + providers, ) if err != nil { - return nil, err + return nil, providers, err } default: - return nil, fmt.Errorf("unknown data source type %q", in.Kind) + return nil, providers, fmt.Errorf("unknown data source type %q", in.Kind) } outConfig.DataSources = append(outConfig.DataSources, out) } - return &outConfig, nil + return &outConfig, providers, nil } func (l *Loader) subgraphName(subgraphs []*nodev1.Subgraph, dataSourceID string) string { diff --git a/router/core/graph_server.go b/router/core/graph_server.go index d6cda32a13..ca7e7c97a7 100644 --- a/router/core/graph_server.go +++ b/router/core/graph_server.go @@ -963,7 +963,7 @@ func (s *graphServer) buildGraphMux(ctx context.Context, }, } - executor, err := ecb.Build( + executor, providers, err := ecb.Build( ctx, &ExecutorBuildOptions{ EngineConfig: engineConfig, @@ -980,7 +980,7 @@ func (s *graphServer) buildGraphMux(ctx context.Context, return nil, fmt.Errorf("failed to build plan configuration: %w", err) } - s.pubSubProviders = ecb.GetProviders() + s.pubSubProviders = providers if pubSubStartupErr := s.startupPubSubProviders(ctx); pubSubStartupErr != nil { return nil, pubSubStartupErr } diff --git a/router/core/plan_generator.go b/router/core/plan_generator.go index 4420f70327..8d5f4f0e03 100644 --- a/router/core/plan_generator.go +++ b/router/core/plan_generator.go @@ -285,7 +285,7 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo }, logger) // this generates the plan configuration using the data source factories from the config package - planConfig, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &routerEngineConfig) + planConfig, _, err := loader.Load(routerConfig.GetEngineConfig(), routerConfig.GetSubgraphs(), &routerEngineConfig) if err != nil { return fmt.Errorf("failed to load configuration: %w", err) } From ddb437d2f0ff241463886b371194a1b03e833477 Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Fri, 9 May 2025 11:59:44 +0200 Subject: [PATCH 48/49] refactor: update BuildEventDataBytes to use operation from ast.Document instead of plan.Visitor --- router/pkg/pubsub/datasource/planner.go | 2 +- router/pkg/pubsub/eventdata/build.go | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/router/pkg/pubsub/datasource/planner.go b/router/pkg/pubsub/datasource/planner.go index 81490ac162..9b5fe1a84a 100644 --- a/router/pkg/pubsub/datasource/planner.go +++ b/router/pkg/pubsub/datasource/planner.go @@ -61,7 +61,7 @@ func (p *Planner) ConfigureFetch() resolve.FetchConfiguration { return resolve.FetchConfiguration{} } - event, err := eventdata.BuildEventDataBytes(p.rootFieldRef, p.visitor, &p.variables) + event, err := eventdata.BuildEventDataBytes(p.rootFieldRef, p.visitor.Operation, &p.variables) if err != nil { p.visitor.Walker.StopWithInternalErr(fmt.Errorf("failed to get resolve data source input: %w", err)) return resolve.FetchConfiguration{} diff --git a/router/pkg/pubsub/eventdata/build.go b/router/pkg/pubsub/eventdata/build.go index 5d1b980849..23d6acf933 100644 --- a/router/pkg/pubsub/eventdata/build.go +++ b/router/pkg/pubsub/eventdata/build.go @@ -4,27 +4,27 @@ import ( "bytes" "encoding/json" - "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" ) -func BuildEventDataBytes(ref int, visitor *plan.Visitor, variables *resolve.Variables) ([]byte, error) { +func BuildEventDataBytes(ref int, operation *ast.Document, variables *resolve.Variables) ([]byte, error) { // Collect the field arguments for fetch based operations - fieldArgs := visitor.Operation.FieldArguments(ref) + fieldArgs := operation.FieldArguments(ref) var dataBuffer bytes.Buffer dataBuffer.WriteByte('{') for i, arg := range fieldArgs { if i > 0 { dataBuffer.WriteByte(',') } - argValue := visitor.Operation.ArgumentValue(arg) - variableName := visitor.Operation.VariableValueNameBytes(argValue.Ref) + argValue := operation.ArgumentValue(arg) + variableName := operation.VariableValueNameBytes(argValue.Ref) contextVariable := &resolve.ContextVariable{ Path: []string{string(variableName)}, Renderer: resolve.NewPlainVariableRenderer(), } variablePlaceHolder, _ := variables.AddVariable(contextVariable) - argumentName := visitor.Operation.ArgumentNameString(arg) + argumentName := operation.ArgumentNameString(arg) escapedKey, err := json.Marshal(argumentName) if err != nil { return nil, err From 49db63a622612c76f415bfd940929d0ca77b520e Mon Sep 17 00:00:00 2001 From: Alessandro Pagnin Date: Wed, 28 May 2025 10:57:24 +0200 Subject: [PATCH 49/49] refactor: update datasource import to use pubsub_datasource in factoryresolver.go --- router/core/factoryresolver.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/router/core/factoryresolver.go b/router/core/factoryresolver.go index acdc0d6cb7..1145d708fd 100644 --- a/router/core/factoryresolver.go +++ b/router/core/factoryresolver.go @@ -11,7 +11,6 @@ import ( "github.com/buger/jsonparser" "github.com/wundergraph/cosmo/router/pkg/pubsub" - "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" pubsub_datasource "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource" "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/argument_templates" "google.golang.org/grpc" @@ -477,7 +476,7 @@ func (l *Loader) Load(engineConfig *nodev1.EngineConfiguration, subgraphs []*nod out, err = plan.NewDataSourceConfiguration( in.Id, - datasource.NewFactory(l.ctx, routerEngineConfig.Events, providers), + pubsub_datasource.NewFactory(l.ctx, routerEngineConfig.Events, providers), dsMeta, providers, )