chore: improve

wundergraph · Jan 21, 2024 · b0f67d5 · b0f67d5
1 parent a31e215
commit b0f67d5
Show file tree

Hide file tree

Showing 3 changed files with 165 additions and 60 deletions.
diff --git a/aws-lambda-router/cmd/main.go b/aws-lambda-router/cmd/main.go
@@ -7,10 +7,7 @@ import (
 	"github.com/akrylysov/algnhsa"
 	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/wundergraph/cosmo/aws-lambda-router/internal"
-	"github.com/wundergraph/cosmo/router/core"
 	"github.com/wundergraph/cosmo/router/pkg/logging"
-	"github.com/wundergraph/cosmo/router/pkg/metric"
-	"github.com/wundergraph/cosmo/router/pkg/trace"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"net/http"
@@ -27,6 +24,8 @@ var (
 	defaultSampleRate = 0.2 // 20% of requests will be sampled
 	enableTelemetry   = os.Getenv("DISABLE_TELEMETRY") != "true"
 	stage             = os.Getenv("STAGE")
+	graphApiToken     = os.Getenv("GRAPH_API_TOKEN")
+	httpPort          = os.Getenv("HTTP_PORT")
 )
 
 func main() {
@@ -42,61 +41,16 @@ func main() {
 		}
 	}()
 
-	httpPort := os.Getenv("HTTP_PORT")
-
-	routerConfig, err := core.SerializeConfigFromFile(routerConfigPath)
-	if err != nil {
-		logger.Fatal("Could not read router config", zap.Error(err), zap.String("path", routerConfigPath))
-	}
-
-	routerOpts := []core.Option{
-		core.WithLogger(logger),
-		core.WithPlayground(true),
-		core.WithIntrospection(true),
-		core.WithStaticRouterConfig(routerConfig),
-		core.WithAwsLambdaRuntime(),
-		core.WithGraphApiToken(os.Getenv("GRAPH_API_TOKEN")),
-	}
-
-	if httpPort != "" {
-		routerOpts = append(routerOpts, core.WithListenerAddr(":"+httpPort))
-	}
-
-	if enableTelemetry {
-		routerOpts = append(routerOpts,
-			core.WithGraphQLMetrics(&core.GraphQLMetricsConfig{
-				Enabled:           true,
-				CollectorEndpoint: "https://cosmo-metrics.wundergraph.com",
-			}),
-			core.WithMetrics(&metric.Config{
-				Name:    telemetryServiceName,
-				Version: internal.Version,
-				OpenTelemetry: metric.OpenTelemetry{
-					Enabled: true,
-				},
-			}),
-			core.WithTracing(&trace.Config{
-				Enabled: true,
-				Name:    telemetryServiceName,
-				Version: internal.Version,
-				Sampler: defaultSampleRate,
-				Propagators: []trace.Propagator{
-					trace.PropagatorTraceContext,
-				},
-			}),
-		)
-	}
-
-	if stage != "" {
-		routerOpts = append(routerOpts,
-			core.WithGraphQLWebURL(fmt.Sprintf("/%s%s", os.Getenv("STAGE"), "/graphql")),
-		)
-	}
-
-	r, err := core.NewRouter(routerOpts...)
-	if err != nil {
-		logger.Fatal("Could not create router", zap.Error(err))
-	}
+	r := internal.NewRouter(
+		internal.WithGraphApiToken(graphApiToken),
+		internal.WithLogger(logger),
+		internal.WithRouterConfigPath(routerConfigPath),
+		internal.WithTelemetryServiceName(telemetryServiceName),
+		internal.WithStage(stage),
+		internal.WithTraceSampleRate(defaultSampleRate),
+		internal.WithEnableTelemetry(enableTelemetry),
+		internal.WithHttpPort(httpPort),
+	)
 
 	svr, err := r.NewServer(ctx)
 	if err != nil {

diff --git a/aws-lambda-router/internal/router.go b/aws-lambda-router/internal/router.go
@@ -0,0 +1,148 @@
+package internal
+
+import (
+	"fmt"
+	"github.com/wundergraph/cosmo/router/core"
+	"github.com/wundergraph/cosmo/router/pkg/metric"
+	"github.com/wundergraph/cosmo/router/pkg/trace"
+	"go.uber.org/zap"
+)
+
+type Option func(*RouterConfig)
+
+type RouterConfig struct {
+	RouterConfigPath     string
+	TelemetryServiceName string
+	RouterOpts           []core.Option
+	GraphApiToken        string
+	HttpPort             string
+	EnableTelemetry      bool
+	Stage                string
+	TraceSampleRate      float64
+	Logger               *zap.Logger
+}
+
+func NewRouter(opts ...Option) *core.Router {
+
+	rc := &RouterConfig{}
+
+	for _, opt := range opts {
+		opt(rc)
+	}
+
+	if rc.Logger == nil {
+		rc.Logger = zap.NewNop()
+	}
+
+	logger := rc.Logger
+
+	routerConfig, err := core.SerializeConfigFromFile(rc.RouterConfigPath)
+	if err != nil {
+		logger.Fatal("Could not read router config", zap.Error(err), zap.String("path", rc.RouterConfigPath))
+	}
+
+	routerOpts := []core.Option{
+		core.WithLogger(logger),
+		core.WithPlayground(true),
+		core.WithIntrospection(true),
+		core.WithStaticRouterConfig(routerConfig),
+		core.WithAwsLambdaRuntime(),
+		core.WithGraphApiToken(rc.GraphApiToken),
+	}
+
+	if rc.HttpPort != "" {
+		routerOpts = append(routerOpts, core.WithListenerAddr(":"+rc.HttpPort))
+	}
+
+	if rc.EnableTelemetry {
+		routerOpts = append(routerOpts,
+			core.WithGraphQLMetrics(&core.GraphQLMetricsConfig{
+				Enabled:           true,
+				CollectorEndpoint: "https://cosmo-metrics.wundergraph.com",
+			}),
+			core.WithMetrics(&metric.Config{
+				Name:    rc.TelemetryServiceName,
+				Version: Version,
+				OpenTelemetry: metric.OpenTelemetry{
+					Enabled: true,
+				},
+			}),
+			core.WithTracing(&trace.Config{
+				Enabled: true,
+				Name:    rc.TelemetryServiceName,
+				Version: Version,
+				Sampler: rc.TraceSampleRate,
+				Propagators: []trace.Propagator{
+					trace.PropagatorTraceContext,
+				},
+			}),
+		)
+	}
+
+	if rc.Stage != "" {
+		routerOpts = append(routerOpts,
+			core.WithGraphQLWebURL(fmt.Sprintf("/%s%s", rc.Stage, "/graphql")),
+		)
+	}
+
+	r, err := core.NewRouter(routerOpts...)
+	if err != nil {
+		logger.Fatal("Could not create router", zap.Error(err))
+	}
+
+	return r
+}
+
+func WithRouterConfigPath(path string) Option {
+	return func(r *RouterConfig) {
+		r.RouterConfigPath = path
+	}
+}
+
+func WithTelemetryServiceName(name string) Option {
+	return func(r *RouterConfig) {
+		r.TelemetryServiceName = name
+	}
+}
+
+func WithRouterOpts(opts ...core.Option) Option {
+	return func(r *RouterConfig) {
+		r.RouterOpts = append(r.RouterOpts, opts...)
+	}
+}
+
+func WithGraphApiToken(token string) Option {
+	return func(r *RouterConfig) {
+		r.GraphApiToken = token
+	}
+}
+
+func WithHttpPort(port string) Option {
+	return func(r *RouterConfig) {
+		r.HttpPort = port
+	}
+}
+
+func WithEnableTelemetry(enable bool) Option {
+	return func(r *RouterConfig) {
+		r.EnableTelemetry = enable
+	}
+}
+
+func WithStage(stage string) Option {
+	return func(r *RouterConfig) {
+		r.Stage = stage
+	}
+}
+
+func WithTraceSampleRate(rate float64) Option {
+	return func(r *RouterConfig) {
+		r.TraceSampleRate = rate
+	}
+}
+
+func WithLogger(logger *zap.Logger) Option {
+	return func(r *RouterConfig) {
+		r.Logger = logger
+	}
+}
diff --git a/aws-lambda-router/internal/router_test.go b/aws-lambda-router/internal/router_test.go
@@ -12,10 +12,13 @@ import (
 )
 
 func TestHandler(t *testing.T) {
-	zapLogger, err := zap.NewProduction()
+	logger, err := zap.NewProduction()
 	require.NoError(t, err)
 
-	r, err := NewRouter(zapLogger, "../router.json")
+	r := NewRouter(
+		WithLogger(logger),
+		WithRouterConfigPath("../router.json"),
+	)
 	require.NoError(t, err)
 
 	svr, err := r.NewServer(context.Background())