diff --git a/plugin/storage/grpc/shared/grpc_client.go b/plugin/storage/grpc/shared/grpc_client.go index ee0e910003b..5bb8b860007 100644 --- a/plugin/storage/grpc/shared/grpc_client.go +++ b/plugin/storage/grpc/shared/grpc_client.go @@ -18,13 +18,18 @@ import ( "github.com/jaegertracing/jaeger/model" "github.com/jaegertracing/jaeger/pkg/bearertoken" _ "github.com/jaegertracing/jaeger/pkg/gogocodec" // force gogo codec registration + "github.com/jaegertracing/jaeger/pkg/tenancy" "github.com/jaegertracing/jaeger/proto-gen/storage_v1" "github.com/jaegertracing/jaeger/storage/dependencystore" "github.com/jaegertracing/jaeger/storage/spanstore" ) -// BearerTokenKey is the key name for the bearer token context value. -const BearerTokenKey = "bearer.token" +const ( + // BearerTokenKey is the key name for the bearer token context value. + BearerTokenKey = "bearer.token" + // TenantKey is the key name for the x-tenant context value. + TenantKey = "x.tenant" +) var ( _ StoragePlugin = (*GRPCClient)(nil) @@ -32,7 +37,7 @@ var ( _ PluginCapabilities = (*GRPCClient)(nil) // upgradeContext composites several steps of upgrading context - upgradeContext = composeContextUpgradeFuncs(upgradeContextWithBearerToken) + upgradeContext = composeContextUpgradeFuncs(upgradeContextWithBearerToken, upgradeContextWithXTenant) ) // GRPCClient implements shared.StoragePlugin and reads/writes spans and dependencies @@ -88,6 +93,22 @@ func upgradeContextWithBearerToken(ctx context.Context) context.Context { return ctx } +// upgradeContextWithXTenant turns the context into a gRPC outgoing context with x tenant +// in the request metadata, if the original context has x-tenant attached. +// Otherwise returns original context. +func upgradeContextWithXTenant(ctx context.Context) context.Context { + tenant := tenancy.GetTenant(ctx) + if tenant != "" { + md, ok := metadata.FromOutgoingContext(ctx) + if !ok { + md = metadata.New(nil) + } + md.Set(TenantKey, tenant) + return metadata.NewOutgoingContext(ctx, md) + } + return ctx +} + // DependencyReader implements shared.StoragePlugin. func (c *GRPCClient) DependencyReader() dependencystore.Reader { return c diff --git a/plugin/storage/grpc/shared/grpc_client_test.go b/plugin/storage/grpc/shared/grpc_client_test.go index 51e56df6bb5..da20ff1ae9f 100644 --- a/plugin/storage/grpc/shared/grpc_client_test.go +++ b/plugin/storage/grpc/shared/grpc_client_test.go @@ -20,6 +20,7 @@ import ( "github.com/jaegertracing/jaeger/model" "github.com/jaegertracing/jaeger/pkg/bearertoken" + "github.com/jaegertracing/jaeger/pkg/tenancy" "github.com/jaegertracing/jaeger/proto-gen/storage_v1" grpcMocks "github.com/jaegertracing/jaeger/proto-gen/storage_v1/mocks" "github.com/jaegertracing/jaeger/storage/spanstore" @@ -132,6 +133,22 @@ func TestContextUpgradeWithoutToken(t *testing.T) { assert.Falsef(t, ok, "Expected no metadata in context") } +func TestContextUpgradeWithTenant(t *testing.T) { + testTenantValue := "test-tenant-val-1" + ctx := tenancy.WithTenant(context.Background(), testTenantValue) + upgradedCtx := upgradeContextWithXTenant(ctx) + md, ok := metadata.FromOutgoingContext(upgradedCtx) + assert.Truef(t, ok, "Expected metadata in context") + xTenantFromMetadata := md.Get(TenantKey) + assert.Equal(t, []string{testTenantValue}, xTenantFromMetadata) +} + +func TestContextUpgradeWithoutTenant(t *testing.T) { + upgradedTenant := upgradeContextWithXTenant(context.Background()) + _, ok := metadata.FromOutgoingContext(upgradedTenant) + assert.Falsef(t, ok, "Expected no metadata in context") +} + func TestGRPCClientGetServices(t *testing.T) { withGRPCClient(func(r *grpcClientTest) { r.spanReader.On("GetServices", mock.Anything, &storage_v1.GetServicesRequest{}).