refactor: inject trace ID as a default header (#50)

cmd/influx/main.go

@@ -116,7 +116,6 @@ func newCli(ctx *cli.Context) (*internal.CLI, error) {
 
 	return &internal.CLI{
 		StdIO:            stdio.TerminalStdio,
-		TraceId:          ctx.String(traceIdFlag),
 		PrintAsJSON:      ctx.Bool(printJsonFlag),
 		HideTableHeaders: ctx.Bool(hideHeadersFlag),
 		ActiveConfig:     activeConfig,
@@ -154,6 +153,15 @@ func newApiClient(ctx *cli.Context, cli *internal.CLI, injectToken bool) (*api.A
 	if injectToken {
 		apiConfig.DefaultHeader["Authorization"] = fmt.Sprintf("Token %s", cfg.Token)
 	}
+	if ctx.IsSet(traceIdFlag) {
+		// NOTE: This is circumventing our codegen. If the header we use for tracing ever changes,
+		// we'll need to manually update the string here to match.
+		//
+		// The alternative is to pass the trace ID to the business logic for every CLI command, and
+		// use codegen'd logic to set the header on every HTTP request. Early versions of the CLI
+		// used that technique, and we found it to be error-prone and easy to forget during testing.
+		apiConfig.DefaultHeader["Zap-Trace-Span"] = ctx.String(traceIdFlag)
+	}
 	apiConfig.Debug = ctx.Bool(httpDebugFlag)
 
 	return api.NewAPIClient(apiConfig), nil

internal/cli.go

@@ -12,7 +12,6 @@ import (
 type CLI struct {
 	StdIO stdio.StdIO
 
-	TraceId          string
 	HideTableHeaders bool
 	PrintAsJSON      bool
 

internal/ping.go

@@ -9,9 +9,6 @@ import (
 // Ping checks the health of a remote InfluxDB instance.
 func (c *CLI) Ping(ctx context.Context, client api.HealthApi) error {
 	req := client.GetHealth(ctx)
-	if c.TraceId != "" {
-		req = req.ZapTraceSpan(c.TraceId)
-	}
 	if _, _, err := client.GetHealthExecute(req); err != nil {
 		return err
 	}

internal/ping_test.go

@@ -17,7 +17,6 @@ func Test_PingSuccess(t *testing.T) {
 
 	client := &mock.HealthApi{
 		GetHealthExecuteFn: func(req api.ApiGetHealthRequest) (api.HealthCheck, *http.Response, error) {
-			require.Nil(t, req.GetZapTraceSpan())
 			return api.HealthCheck{Status: api.HEALTHCHECKSTATUS_PASS}, nil, nil
 		},
 	}
@@ -29,25 +28,6 @@ func Test_PingSuccess(t *testing.T) {
 	require.Equal(t, "OK\n", stdio.Stdout())
 }
 
-func Test_PingSuccessWithTracing(t *testing.T) {
-	t.Parallel()
-
-	traceId := "trace-id"
-	client := &mock.HealthApi{
-		GetHealthExecuteFn: func(req api.ApiGetHealthRequest) (api.HealthCheck, *http.Response, error) {
-			require.NotNil(t, req.GetZapTraceSpan())
-			require.Equal(t, traceId, *req.GetZapTraceSpan())
-			return api.HealthCheck{Status: api.HEALTHCHECKSTATUS_PASS}, nil, nil
-		},
-	}
-
-	stdio := mock.NewMockStdio(nil, true)
-	cli := &internal.CLI{TraceId: traceId, StdIO: stdio}
-
-	require.NoError(t, cli.Ping(context.Background(), client))
-	require.Equal(t, "OK\n", stdio.Stdout())
-}
-
 func Test_PingFailedRequest(t *testing.T) {
 	t.Parallel()
 

internal/setup.go

@@ -45,9 +45,6 @@ func (c *CLI) Setup(ctx context.Context, client api.SetupApi, params *SetupParam
 
 	// Check if setup is even allowed.
 	checkReq := client.GetSetup(ctx)
-	if c.TraceId != "" {
-		checkReq = checkReq.ZapTraceSpan(c.TraceId)
-	}
 	checkResp, _, err := client.GetSetupExecute(checkReq)
 	if err != nil {
 		return fmt.Errorf("failed to check if already set up: %w", err)
@@ -62,9 +59,6 @@ func (c *CLI) Setup(ctx context.Context, client api.SetupApi, params *SetupParam
 		return err
 	}
 	setupReq := client.PostSetup(ctx).OnboardingRequest(setupBody)
-	if c.TraceId != "" {
-		setupReq = setupReq.ZapTraceSpan(c.TraceId)
-	}
 	resp, _, err := client.PostSetupExecute(setupReq)
 	if err != nil {
 		return fmt.Errorf("failed to setup instance: %w", err)

internal/setup_test.go

@@ -152,69 +152,6 @@ func Test_SetupSuccessNoninteractive(t *testing.T) {
 	require.Regexp(t, fmt.Sprintf("%s\\s+%s\\s+%s", params.Username, params.Org, params.Bucket), data)
 }
 
-func Test_SetupSuccessNoninteractiveWithTracing(t *testing.T) {
-	t.Parallel()
-
-	traceId := "trace-id"
-	retentionSecs := int64(duration.Week.Seconds())
-	params := internal.SetupParams{
-		Username:   "user",
-		Password:   "mysecretpassword",
-		AuthToken:  "mytoken",
-		Org:        "org",
-		Bucket:     "bucket",
-		Retention:  fmt.Sprintf("%ds", retentionSecs),
-		Force:      true,
-		ConfigName: "my-config",
-	}
-	resp := api.OnboardingResponse{
-		Auth:   &api.Authorization{Token: &params.AuthToken},
-		Org:    &api.Organization{Name: params.Org},
-		User:   &api.UserResponse{Name: params.Username},
-		Bucket: &api.Bucket{Name: params.Bucket},
-	}
-	client := &mock.SetupApi{
-		GetSetupExecuteFn: func(req api.ApiGetSetupRequest) (api.InlineResponse200, *http.Response, error) {
-			require.Equal(t, traceId, *req.GetZapTraceSpan())
-			return api.InlineResponse200{Allowed: api.PtrBool(true)}, nil, nil
-		},
-		PostSetupExecuteFn: func(req api.ApiPostSetupRequest) (api.OnboardingResponse, *http.Response, error) {
-			require.Equal(t, traceId, *req.GetZapTraceSpan())
-			body := req.GetOnboardingRequest()
-			require.Equal(t, params.Username, body.Username)
-			require.Equal(t, params.Password, *body.Password)
-			require.Equal(t, params.AuthToken, *body.Token)
-			require.Equal(t, params.Org, body.Org)
-			require.Equal(t, params.Bucket, body.Bucket)
-			require.Equal(t, retentionSecs, *body.RetentionPeriodSeconds)
-			return resp, nil, nil
-		},
-	}
-
-	host := "fake-host"
-	configSvc := &mock.ConfigService{
-		ListConfigsFn: func() (config.Configs, error) {
-			return nil, nil
-		},
-		CreateConfigFn: func(cfg config.Config) (config.Config, error) {
-			require.Equal(t, params.ConfigName, cfg.Name)
-			require.Equal(t, params.AuthToken, cfg.Token)
-			require.Equal(t, host, cfg.Host)
-			require.Equal(t, params.Org, cfg.Org)
-			return cfg, nil
-		},
-	}
-	stdio := mock.NewMockStdio(nil, true)
-	cli := &internal.CLI{ConfigService: configSvc, ActiveConfig: config.Config{Host: host}, StdIO: stdio, TraceId: traceId}
-	require.NoError(t, cli.Setup(context.Background(), client, &params))
-
-	outLines := strings.Split(strings.TrimSpace(stdio.Stdout()), "\n")
-	require.Len(t, outLines, 2)
-	header, data := outLines[0], outLines[1]
-	require.Regexp(t, "User\\s+Organization\\s+Bucket", header)
-	require.Regexp(t, fmt.Sprintf("%s\\s+%s\\s+%s", params.Username, params.Org, params.Bucket), data)
-}
-
 func Test_SetupSuccessInteractive(t *testing.T) {
 	t.Parallel()
 

internal/write.go

@@ -56,9 +56,6 @@ func (c *CLI) Write(ctx context.Context, clients *WriteClients, params *WritePar
 
 	writeBatch := func(batch []byte) error {
 		req := clients.Client.PostWrite(ctx).Body(batch).ContentEncoding("gzip").Precision(params.Precision)
-		if c.TraceId != "" {
-			req = req.ZapTraceSpan(c.TraceId)
-		}
 		if params.BucketID != "" {
 			req = req.Bucket(params.BucketID)
 		} else {

internal/write_test.go

@@ -109,7 +109,7 @@ func TestWriteByNames(t *testing.T) {
 		BucketName: "my-bucket",
 		Precision:  api.WRITEPRECISION_US,
 	}
-	cli := internal.CLI{TraceId: "my-trace-id", ActiveConfig: config.Config{Org: "my-default-org"}}
+	cli := internal.CLI{ActiveConfig: config.Config{Org: "my-default-org"}}
 
 	var writtenLines []string
 	client := mock.WriteApi{
@@ -118,7 +118,6 @@ func TestWriteByNames(t *testing.T) {
 			require.Equal(t, params.OrgName, *req.GetOrg())
 			require.Equal(t, params.BucketName, *req.GetBucket())
 			require.Equal(t, params.Precision, *req.GetPrecision())
-			require.Equal(t, cli.TraceId, *req.GetZapTraceSpan())
 
 			// Make sure the body is properly marked for compression, and record what was sent.
 			require.Equal(t, "gzip", *req.GetContentEncoding())