diff --git a/spanner/client.go b/spanner/client.go index 3190db17c7d3..00fefee852f2 100644 --- a/spanner/client.go +++ b/spanner/client.go @@ -68,6 +68,10 @@ const ( requestsCompressionHeader = "x-response-encoding" + // endToEndTracingHeader is the name of the metadata header if client + // has opted-in for the creation of trace spans on the Spanner layer. + endToEndTracingHeader = "x-goog-spanner-end-to-end-tracing" + // numChannels is the default value for NumChannels of client. numChannels = 4 ) @@ -337,6 +341,14 @@ type ClientConfig struct { DirectedReadOptions *sppb.DirectedReadOptions OpenTelemetryMeterProvider metric.MeterProvider + + // EnableEndToEndTracing indicates whether end to end tracing is enabled or not. If + // it is enabled, trace spans will be created at Spanner layer. Enabling end to end + // tracing requires OpenTelemetry to be set up. Simply enabling this option won't + // generate traces at Spanner layer. + // + // Default: false + EnableEndToEndTracing bool } type openTelemetryConfig struct { @@ -452,6 +464,13 @@ func newClientWithConfig(ctx context.Context, database string, config ClientConf if config.Compression == gzip.Name { md.Append(requestsCompressionHeader, gzip.Name) } + // Append end to end tracing header if SPANNER_ENABLE_END_TO_END_TRACING + // environment variable has been set or client has passed the opt-in + // option in ClientConfig. + endToEndTracingEnvironmentVariable := os.Getenv("SPANNER_ENABLE_END_TO_END_TRACING") + if config.EnableEndToEndTracing || endToEndTracingEnvironmentVariable == "true" { + md.Append(endToEndTracingHeader, "true") + } // Create a session client. sc := newSessionClient(pool, database, config.UserAgent, sessionLabels, config.DatabaseRole, config.DisableRouteToLeader, md, config.BatchTimeout, config.Logger, config.CallOptions) diff --git a/spanner/client_test.go b/spanner/client_test.go index b7965dc03073..82b2a9a2d493 100644 --- a/spanner/client_test.go +++ b/spanner/client_test.go @@ -4387,6 +4387,63 @@ func TestClient_WithGRPCConnectionPoolAndNumChannels_Misconfigured(t *testing.T) } } +func TestClient_EndToEndTracingHeader(t *testing.T) { + tests := []struct { + name string + endToEndTracingEnv string + enableEndToEndTracing bool + wantEndToEndTracing bool + }{ + { + name: "when end-to-end tracing is enabled via config", + enableEndToEndTracing: true, + wantEndToEndTracing: true, + endToEndTracingEnv: "false", + }, + { + name: "when end-to-end tracing is enabled via env", + enableEndToEndTracing: false, + wantEndToEndTracing: true, + endToEndTracingEnv: "true", + }, + { + name: "when end-to-end tracing is disabled", + enableEndToEndTracing: false, + wantEndToEndTracing: false, + endToEndTracingEnv: "false", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("SPANNER_ENABLE_END_TO_END_TRACING", tt.endToEndTracingEnv) + + server, opts, teardown := NewMockedSpannerInMemTestServer(t) + defer teardown() + config := ClientConfig{} + if tt.enableEndToEndTracing { + config.EnableEndToEndTracing = true + } + + client, err := makeClientWithConfig(context.Background(), "projects/p/instances/i/databases/d", config, server.ServerAddress, opts...) + if err != nil { + t.Fatalf("failed to get a client: %v", err) + } + + gotEndToEndTracing := false + for _, val := range client.sc.md.Get(endToEndTracingHeader) { + if val == "true" { + gotEndToEndTracing = true + } + } + + if gotEndToEndTracing != tt.wantEndToEndTracing { + t.Fatalf("mismatch in client configuration for property EnableEndToEndTracing: got %v, want %v", gotEndToEndTracing, tt.wantEndToEndTracing) + } + }) + } +} + func TestClient_WithCustomBatchTimeout(t *testing.T) { t.Parallel()