diff --git a/client/client.go b/client/client.go index 4b8215ec..cb999112 100644 --- a/client/client.go +++ b/client/client.go @@ -48,6 +48,7 @@ const ( daprPortEnvVarName = "DAPR_GRPC_PORT" /* #nosec */ daprGRPCEndpointEnvVarName = "DAPR_GRPC_ENDPOINT" traceparentKey = "traceparent" + baggageHeader = "baggage" apiTokenKey = "dapr-api-token" /* #nosec */ apiTokenEnvVarName = "DAPR_API_TOKEN" /* #nosec */ clientDefaultTimeoutSeconds = 5 @@ -199,6 +200,9 @@ type Client interface { // WithTraceID adds existing trace ID to the outgoing context. WithTraceID(ctx context.Context, id string) context.Context + // WithBaggage adds baggage information to the outgoing context + WithBaggage(ctx context.Context, baggage map[string]string) context.Context + // WithAuthToken sets Dapr API token on the instantiated client. WithAuthToken(token string) @@ -506,6 +510,17 @@ func (c *GRPCClient) WithTraceID(ctx context.Context, id string) context.Context return metadata.NewOutgoingContext(ctx, md) } +// WithBaggage adds baggage information to the outgoing context +func (c *GRPCClient) WithBaggage(ctx context.Context, baggage map[string]string) context.Context { + baggageValues := make([]string, 0, len(baggage)) + for key, value := range baggage { + baggageValues = append(baggageValues, key+"="+value) + } + + baggageString := strings.Join(baggageValues, ",") + return metadata.AppendToOutgoingContext(ctx, baggageHeader, baggageString) +} + // Shutdown the sidecar. func (c *GRPCClient) Shutdown(ctx context.Context) error { _, err := c.protoClient.Shutdown(ctx, &pb.ShutdownRequest{}) diff --git a/client/client_test.go b/client/client_test.go index 9b41e503..67d463d1 100644 --- a/client/client_test.go +++ b/client/client_test.go @@ -30,6 +30,7 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" "google.golang.org/grpc/test/bufconn" "google.golang.org/protobuf/types/known/anypb" "google.golang.org/protobuf/types/known/emptypb" @@ -101,6 +102,19 @@ func TestNewClient(t *testing.T) { _ = testClient.WithTraceID(t.Context(), "test") }) + t.Run("new client with baggage", func(t *testing.T) { + baggage := map[string]string{ + "key1": "value1", + "key2": "value2", + } + + ctx := testClient.WithBaggage(t.Context(), baggage) + metadata, ok := metadata.FromOutgoingContext(ctx) + require.True(t, ok) + baggageString := metadata.Get(baggageHeader) + require.Equal(t, []string{"key1=value1,key2=value2"}, baggageString) + }) + t.Run("new socket client closed with token", func(t *testing.T) { t.Setenv(apiTokenEnvVarName, "test") c, err := NewClientWithSocket(testSocket) @@ -123,6 +137,23 @@ func TestNewClient(t *testing.T) { ctx := c.WithTraceID(t.Context(), "") _ = c.WithTraceID(ctx, "test") }) + + t.Run("new socket client with baggage", func(t *testing.T) { + c, err := NewClientWithSocket(testSocket) + require.NoError(t, err) + defer c.Close() + + baggage := map[string]string{ + "key1": "value1", + "key2": "value2", + } + + ctx := c.WithBaggage(t.Context(), baggage) + metadata, ok := metadata.FromOutgoingContext(ctx) + require.True(t, ok) + baggageString := metadata.Get(baggageHeader) + require.Equal(t, []string{"key1=value1,key2=value2"}, baggageString) + }) } func TestShutdown(t *testing.T) {