storage: propagate tenant to grpc backend

Signed-off-by: Benedikt Bongartz <[email protected]>

plugin/storage/grpc/shared/grpc_client.go

@@ -18,21 +18,26 @@ 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)
 	_ ArchiveStoragePlugin = (*GRPCClient)(nil)
 	_ 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

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{}).