Make Headers Setter context default value configurable (open-telemetr…

…y#35151)
kyo-ke · Sep 12, 2024 · 650620a · 650620a
1 parent 21208c0
commit 650620a
Show file tree

Hide file tree

Showing 4 changed files with 93 additions and 8 deletions.
diff --git a/extension/headerssetterextension/config.go b/extension/headerssetterextension/config.go
@@ -19,10 +19,11 @@ type Config struct {
 }
 
 type HeaderConfig struct {
-	Action      ActionValue `mapstructure:"action"`
-	Key         *string     `mapstructure:"key"`
-	Value       *string     `mapstructure:"value"`
-	FromContext *string     `mapstructure:"from_context"`
+	Action       ActionValue `mapstructure:"action"`
+	Key          *string     `mapstructure:"key"`
+	Value        *string     `mapstructure:"value"`
+	FromContext  *string     `mapstructure:"from_context"`
+	DefaultValue *string     `mapstructure:"default_value"`
 }
 
 // ActionValue is the enum to capture the four types of actions to perform on a header

diff --git a/extension/headerssetterextension/extension.go b/extension/headerssetterextension/extension.go
@@ -30,13 +30,20 @@ func newHeadersSetterExtension(cfg *Config, logger *zap.Logger) (auth.Client, er
 	headers := make([]Header, 0, len(cfg.HeadersConfig))
 	for _, header := range cfg.HeadersConfig {
 		var s source.Source
+
 		if header.Value != nil {
 			s = &source.StaticSource{
 				Value: *header.Value,
 			}
 		} else if header.FromContext != nil {
+			DefaultValue := ""
+			if header.DefaultValue != nil {
+				DefaultValue = *header.DefaultValue
+			}
+
 			s = &source.ContextSource{
-				Key: *header.FromContext,
+				Key:          *header.FromContext,
+				DefaultValue: DefaultValue,
 			}
 		}
 

diff --git a/extension/headerssetterextension/extension_test.go b/extension/headerssetterextension/extension_test.go
@@ -218,6 +218,82 @@ var (
 				"header_name": "",
 			},
 		},
+		// Default Value test - INSERT
+		{
+			// when "TenantID" exists in the context, use the value as header_name
+			cfg: &Config{
+				HeadersConfig: []HeaderConfig{
+					{
+						Key:         &header,
+						Action:      INSERT,
+						FromContext: stringp("TenantID"),
+					},
+				},
+			},
+			metadata: client.NewMetadata(
+				map[string][]string{"TenantID": {"12345"}},
+			),
+			expectedHeaders: map[string]string{
+				"header_name": "12345",
+			},
+		},
+		{
+			// even when "TenantID" doesn't exist in the context, use the DefaultValue
+			cfg: &Config{
+				HeadersConfig: []HeaderConfig{
+					{
+						Key:          &header,
+						Action:       INSERT,
+						FromContext:  stringp("TenantID"),
+						DefaultValue: stringp("hoge"),
+					},
+				},
+			},
+			metadata: client.NewMetadata(
+				map[string][]string{"foo": {"bar"}},
+			),
+			expectedHeaders: map[string]string{
+				"header_name": "hoge",
+			},
+		},
+		{
+			// when DefaultValue and "TenantID" exist in the context, the context takes the precedence
+			cfg: &Config{
+				HeadersConfig: []HeaderConfig{
+					{
+						Key:          &header,
+						Action:       INSERT,
+						FromContext:  stringp("TenantID"),
+						DefaultValue: stringp("hoge"),
+					},
+				},
+			},
+			metadata: client.NewMetadata(
+				map[string][]string{"TenantID": {"12345"}},
+			),
+			expectedHeaders: map[string]string{
+				"header_name": "12345",
+			},
+		},
+		{
+			// when "TenantID" and DefaultValue don't exist, set the header = ""
+			cfg: &Config{
+				HeadersConfig: []HeaderConfig{
+					{
+						Key:         &header,
+						Action:      INSERT,
+						FromContext: stringp("TenantID"),
+					},
+				},
+			},
+			metadata: client.NewMetadata(
+				map[string][]string{"foo": {"bar"}},
+			),
+			expectedHeaders: map[string]string{
+				"header_name": "",
+			},
+		},
+		// Default Value test - UPSERT ... should behave as same as INSERT
 	}
 )
 

diff --git a/extension/headerssetterextension/internal/source/context.go b/extension/headerssetterextension/internal/source/context.go
@@ -13,19 +13,20 @@ import (
 var _ Source = (*ContextSource)(nil)
 
 type ContextSource struct {
-	Key string
+	Key          string
+	DefaultValue string
 }
 
 func (ts *ContextSource) Get(ctx context.Context) (string, error) {
 	cl := client.FromContext(ctx)
 	ss := cl.Metadata.Get(ts.Key)
 
 	if len(ss) == 0 {
-		return "", nil
+		return ts.DefaultValue, nil
 	}
 
 	if len(ss) > 1 {
-		return "", fmt.Errorf("%d source keys found in the context, can't determine which one to use", len(ss))
+		return ts.DefaultValue, fmt.Errorf("%d source keys found in the context, can't determine which one to use", len(ss))
 	}
 
 	return ss[0], nil