[cmd/telemetrygen] Support boolean values in attributes (open-telemet…

…ry#34558) **Description:** Support boolean values in attributes **Link to tracking Issue:** Resolves open-telemetry#18928 **Testing:** <Describe what testing was performed and which tests were added.> **Documentation:** <Describe the documentation added.> --------- Signed-off-by: Murphy Chen <[email protected]>
abhishek-at-cloudwerx · Aug 9, 2024 · 148f5e3 · 148f5e3
1 parent 18a47f3
commit 148f5e3
Show file tree

Hide file tree

Showing 6 changed files with 91 additions and 20 deletions.
diff --git a/.chloggen/telemetrygen-attribute-support-boolean.yaml b/.chloggen/telemetrygen-attribute-support-boolean.yaml
@@ -0,0 +1,27 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: enhancement
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: telemetrygen
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: Support boolean values in `--telemetry-attributes` and `--otlp-attributes` flag
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [18928]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext:
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: []
diff --git a/cmd/telemetrygen/internal/common/config.go b/cmd/telemetrygen/internal/common/config.go
@@ -5,6 +5,7 @@ package common
 
 import (
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
@@ -22,7 +23,7 @@ const (
 	defaultHTTPEndpoint = "localhost:4318"
 )
 
-type KeyValue map[string]string
+type KeyValue map[string]any
 
 var _ pflag.Value = (*KeyValue)(nil)
 
@@ -36,6 +37,14 @@ func (v *KeyValue) Set(s string) error {
 		return errFormatOTLPAttributes
 	}
 	val := kv[1]
+	if val == "true" {
+		(*v)[kv[0]] = true
+		return nil
+	}
+	if val == "false" {
+		(*v)[kv[0]] = false
+		return nil
+	}
 	if len(val) < 2 || !strings.HasPrefix(val, "\"") || !strings.HasSuffix(val, "\"") {
 		return errDoubleQuotesOTLPAttributes
 	}
@@ -45,7 +54,7 @@ func (v *KeyValue) Set(s string) error {
 }
 
 func (v *KeyValue) Type() string {
-	return "map[string]string"
+	return "map[string]any"
 }
 
 type Config struct {
@@ -94,8 +103,13 @@ func (c *Config) GetAttributes() []attribute.KeyValue {
 	var attributes []attribute.KeyValue
 
 	if len(c.ResourceAttributes) > 0 {
-		for k, v := range c.ResourceAttributes {
-			attributes = append(attributes, attribute.String(k, v))
+		for k, t := range c.ResourceAttributes {
+			switch v := t.(type) {
+			case string:
+				attributes = append(attributes, attribute.String(k, v))
+			case bool:
+				attributes = append(attributes, attribute.Bool(k, v))
+			}
 		}
 	}
 	return attributes
@@ -105,13 +119,33 @@ func (c *Config) GetTelemetryAttributes() []attribute.KeyValue {
 	var attributes []attribute.KeyValue
 
 	if len(c.TelemetryAttributes) > 0 {
-		for k, v := range c.TelemetryAttributes {
-			attributes = append(attributes, attribute.String(k, v))
+		for k, t := range c.TelemetryAttributes {
+			switch v := t.(type) {
+			case string:
+				attributes = append(attributes, attribute.String(k, v))
+			case bool:
+				attributes = append(attributes, attribute.Bool(k, v))
+			}
 		}
 	}
 	return attributes
 }
 
+func (c *Config) GetHeaders() map[string]string {
+	m := make(map[string]string, len(c.Headers))
+
+	for k, t := range c.Headers {
+		switch v := t.(type) {
+		case bool:
+			m[k] = strconv.FormatBool(v)
+		case string:
+			m[k] = v
+		}
+	}
+
+	return m
+}
+
 // CommonFlags registers common config flags.
 func (c *Config) CommonFlags(fs *pflag.FlagSet) {
 	fs.IntVar(&c.WorkerCount, "workers", 1, "Number of workers (goroutines) to run")
@@ -125,21 +159,23 @@ func (c *Config) CommonFlags(fs *pflag.FlagSet) {
 	fs.BoolVar(&c.UseHTTP, "otlp-http", false, "Whether to use HTTP exporter rather than a gRPC one")
 
 	// custom headers
-	c.Headers = make(map[string]string)
+	c.Headers = make(KeyValue)
 	fs.Var(&c.Headers, "otlp-header", "Custom header to be passed along with each OTLP request. The value is expected in the format key=\"value\". "+
 		"Note you may need to escape the quotes when using the tool from a cli. "+
 		`Flag may be repeated to set multiple headers (e.g --otlp-header key1=\"value1\" --otlp-header key2=\"value2\")`)
 
 	// custom resource attributes
-	c.ResourceAttributes = make(map[string]string)
+	c.ResourceAttributes = make(KeyValue)
 	fs.Var(&c.ResourceAttributes, "otlp-attributes", "Custom resource attributes to use. The value is expected in the format key=\"value\". "+
+		"You can use key=true or key=false. to set boolean attribute."+
 		"Note you may need to escape the quotes when using the tool from a cli. "+
-		`Flag may be repeated to set multiple attributes (e.g --otlp-attributes key1=\"value1\" --otlp-attributes key2=\"value2\")`)
+		`Flag may be repeated to set multiple attributes (e.g --otlp-attributes key1=\"value1\" --otlp-attributes key2=\"value2\" --telemetry-attributes key3=true)`)
 
-	c.TelemetryAttributes = make(map[string]string)
+	c.TelemetryAttributes = make(KeyValue)
 	fs.Var(&c.TelemetryAttributes, "telemetry-attributes", "Custom telemetry attributes to use. The value is expected in the format key=\"value\". "+
+		"You can use key=true or key=false. to set boolean attribute."+
 		"Note you may need to escape the quotes when using the tool from a cli. "+
-		`Flag may be repeated to set multiple attributes (e.g --telemetry-attributes key1=\"value1\" --telemetry-attributes key2=\"value2\")`)
+		`Flag may be repeated to set multiple attributes (e.g --telemetry-attributes key1=\"value1\" --telemetry-attributes key2=\"value2\" --telemetry-attributes key3=true)`)
 
 	// TLS CA configuration
 	fs.StringVar(&c.CaFile, "ca-cert", "", "Trusted Certificate Authority to verify server certificate")

diff --git a/cmd/telemetrygen/internal/common/config_test.go b/cmd/telemetrygen/internal/common/config_test.go
@@ -17,11 +17,11 @@ func TestKeyValueSet(t *testing.T) {
 	}{
 		{
 			flag:     "key=\"value\"",
-			expected: KeyValue(map[string]string{"key": "value"}),
+			expected: KeyValue(map[string]any{"key": "value"}),
 		},
 		{
 			flag:     "key=\"\"",
-			expected: KeyValue(map[string]string{"key": ""}),
+			expected: KeyValue(map[string]any{"key": ""}),
 		},
 		{
 			flag: "key=\"",
@@ -35,11 +35,19 @@ func TestKeyValueSet(t *testing.T) {
 			flag: "key",
 			err:  errFormatOTLPAttributes,
 		},
+		{
+			flag:     "key=true",
+			expected: KeyValue(map[string]any{"key": true}),
+		},
+		{
+			flag:     "key=false",
+			expected: KeyValue(map[string]any{"key": false}),
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.flag, func(t *testing.T) {
-			kv := KeyValue(make(map[string]string))
+			kv := KeyValue(make(map[string]any))
 			err := kv.Set(tt.flag)
 			if err != nil || tt.err != nil {
 				assert.Equal(t, err, tt.err)

diff --git a/cmd/telemetrygen/internal/logs/exporter.go b/cmd/telemetrygen/internal/logs/exporter.go
@@ -30,7 +30,7 @@ func grpcExporterOptions(cfg *Config) ([]otlploggrpc.Option, error) {
 	}
 
 	if len(cfg.Headers) > 0 {
-		grpcExpOpt = append(grpcExpOpt, otlploggrpc.WithHeaders(cfg.Headers))
+		grpcExpOpt = append(grpcExpOpt, otlploggrpc.WithHeaders(cfg.GetHeaders()))
 	}
 
 	return grpcExpOpt, nil
@@ -55,7 +55,7 @@ func httpExporterOptions(cfg *Config) ([]otlploghttp.Option, error) {
 	}
 
 	if len(cfg.Headers) > 0 {
-		httpExpOpt = append(httpExpOpt, otlploghttp.WithHeaders(cfg.Headers))
+		httpExpOpt = append(httpExpOpt, otlploghttp.WithHeaders(cfg.GetHeaders()))
 	}
 
 	return httpExpOpt, nil

diff --git a/cmd/telemetrygen/internal/metrics/exporter.go b/cmd/telemetrygen/internal/metrics/exporter.go
@@ -30,7 +30,7 @@ func grpcExporterOptions(cfg *Config) ([]otlpmetricgrpc.Option, error) {
 	}
 
 	if len(cfg.Headers) > 0 {
-		grpcExpOpt = append(grpcExpOpt, otlpmetricgrpc.WithHeaders(cfg.Headers))
+		grpcExpOpt = append(grpcExpOpt, otlpmetricgrpc.WithHeaders(cfg.GetHeaders()))
 	}
 
 	return grpcExpOpt, nil
@@ -55,7 +55,7 @@ func httpExporterOptions(cfg *Config) ([]otlpmetrichttp.Option, error) {
 	}
 
 	if len(cfg.Headers) > 0 {
-		httpExpOpt = append(httpExpOpt, otlpmetrichttp.WithHeaders(cfg.Headers))
+		httpExpOpt = append(httpExpOpt, otlpmetrichttp.WithHeaders(cfg.GetHeaders()))
 	}
 
 	return httpExpOpt, nil

diff --git a/cmd/telemetrygen/internal/traces/exporter.go b/cmd/telemetrygen/internal/traces/exporter.go
@@ -30,7 +30,7 @@ func grpcExporterOptions(cfg *Config) ([]otlptracegrpc.Option, error) {
 	}
 
 	if len(cfg.Headers) > 0 {
-		grpcExpOpt = append(grpcExpOpt, otlptracegrpc.WithHeaders(cfg.Headers))
+		grpcExpOpt = append(grpcExpOpt, otlptracegrpc.WithHeaders(cfg.GetHeaders()))
 	}
 
 	return grpcExpOpt, nil
@@ -55,7 +55,7 @@ func httpExporterOptions(cfg *Config) ([]otlptracehttp.Option, error) {
 	}
 
 	if len(cfg.Headers) > 0 {
-		httpExpOpt = append(httpExpOpt, otlptracehttp.WithHeaders(cfg.Headers))
+		httpExpOpt = append(httpExpOpt, otlptracehttp.WithHeaders(cfg.GetHeaders()))
 	}
 
 	return httpExpOpt, nil