Allow specifying OTLP HTTP headers from env variable (#1290)

lalitb · web-flow · commit 5eaaf530538f · 2023-10-08T15:20:48.000-04:00
diff --git a/opentelemetry-otlp/CHANGELOG.md b/opentelemetry-otlp/CHANGELOG.md
@@ -6,6 +6,7 @@
 
 - Add `build_{signal}_exporter` methods to client builders (#1187)
 - Add `grpcio` metrics exporter (#1202)
+- Allow specifying OTLP HTTP headers from env variable (#1290)
 
 ### Changed
 
diff --git a/opentelemetry-otlp/src/exporter/http/mod.rs b/opentelemetry-otlp/src/exporter/http/mod.rs
@@ -1,4 +1,7 @@
-use crate::{ExportConfig, Protocol, OTEL_EXPORTER_OTLP_ENDPOINT, OTEL_EXPORTER_OTLP_TIMEOUT};
+use crate::{
+    ExportConfig, Protocol, OTEL_EXPORTER_OTLP_ENDPOINT, OTEL_EXPORTER_OTLP_HEADERS,
+    OTEL_EXPORTER_OTLP_TIMEOUT,
+};
 use http::{HeaderName, HeaderValue, Uri};
 use opentelemetry_http::HttpClient;
 use std::collections::HashMap;
@@ -143,6 +146,7 @@ impl HttpExporterBuilder {
         signal_endpoint_var: &str,
         signal_endpoint_path: &str,
         signal_timeout_var: &str,
+        signal_http_headers_var: &str,
     ) -> Result<OtlpHttpClient, crate::Error> {
         let endpoint = resolve_endpoint(
             signal_endpoint_var,
@@ -168,7 +172,7 @@ impl HttpExporterBuilder {
             .ok_or(crate::Error::NoHttpClient)?;
 
         #[allow(clippy::mutable_key_type)] // http headers are not mutated
-        let headers = self
+        let mut headers: HashMap<HeaderName, HeaderValue> = self
             .http_config
             .headers
             .take()
@@ -182,6 +186,13 @@ impl HttpExporterBuilder {
             })
             .collect();
 
+        // read headers from env var - signal specific env var is preferred over general
+        if let Ok(input) =
+            env::var(signal_http_headers_var).or_else(|_| env::var(OTEL_EXPORTER_OTLP_HEADERS))
+        {
+            add_header_from_string(&input, &mut headers);
+        }
+
         Ok(OtlpHttpClient::new(http_client, endpoint, headers, timeout))
     }
 
@@ -190,12 +201,16 @@ impl HttpExporterBuilder {
     pub fn build_span_exporter(
         mut self,
     ) -> Result<crate::SpanExporter, opentelemetry::trace::TraceError> {
-        use crate::{OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, OTEL_EXPORTER_OTLP_TRACES_TIMEOUT};
+        use crate::{
+            OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, OTEL_EXPORTER_OTLP_TRACES_HEADERS,
+            OTEL_EXPORTER_OTLP_TRACES_TIMEOUT,
+        };
 
         let client = self.build_client(
             OTEL_EXPORTER_OTLP_TRACES_ENDPOINT,
             "/v1/traces",
             OTEL_EXPORTER_OTLP_TRACES_TIMEOUT,
+            OTEL_EXPORTER_OTLP_TRACES_HEADERS,
         )?;
 
         Ok(crate::SpanExporter::new(client))
@@ -204,12 +219,16 @@ impl HttpExporterBuilder {
     /// Create a log exporter with the current configuration
     #[cfg(feature = "logs")]
     pub fn build_log_exporter(mut self) -> opentelemetry::logs::LogResult<crate::LogExporter> {
-        use crate::{OTEL_EXPORTER_OTLP_LOGS_ENDPOINT, OTEL_EXPORTER_OTLP_LOGS_TIMEOUT};
+        use crate::{
+            OTEL_EXPORTER_OTLP_LOGS_ENDPOINT, OTEL_EXPORTER_OTLP_LOGS_HEADERS,
+            OTEL_EXPORTER_OTLP_LOGS_TIMEOUT,
+        };
 
         let client = self.build_client(
             OTEL_EXPORTER_OTLP_LOGS_ENDPOINT,
             "/v1/logs",
             OTEL_EXPORTER_OTLP_LOGS_TIMEOUT,
+            OTEL_EXPORTER_OTLP_LOGS_HEADERS,
         )?;
 
         Ok(crate::LogExporter::new(client))
@@ -222,12 +241,16 @@ impl HttpExporterBuilder {
         aggregation_selector: Box<dyn opentelemetry_sdk::metrics::reader::AggregationSelector>,
         temporality_selector: Box<dyn opentelemetry_sdk::metrics::reader::TemporalitySelector>,
     ) -> opentelemetry::metrics::Result<crate::MetricsExporter> {
-        use crate::{OTEL_EXPORTER_OTLP_METRICS_ENDPOINT, OTEL_EXPORTER_OTLP_METRICS_TIMEOUT};
+        use crate::{
+            OTEL_EXPORTER_OTLP_METRICS_ENDPOINT, OTEL_EXPORTER_OTLP_METRICS_HEADERS,
+            OTEL_EXPORTER_OTLP_METRICS_TIMEOUT,
+        };
 
         let client = self.build_client(
             OTEL_EXPORTER_OTLP_METRICS_ENDPOINT,
             "/v1/metrics",
             OTEL_EXPORTER_OTLP_METRICS_TIMEOUT,
+            OTEL_EXPORTER_OTLP_METRICS_HEADERS,
         )?;
 
         Ok(crate::MetricsExporter::new(
@@ -291,6 +314,25 @@ fn resolve_endpoint(
         .map_err(From::from)
 }
 
+#[allow(clippy::mutable_key_type)] // http headers are not mutated
+fn add_header_from_string(input: &str, headers: &mut HashMap<HeaderName, HeaderValue>) {
+    for pair in input.split_terminator(',') {
+        if pair.trim().is_empty() {
+            continue;
+        }
+        if let Some((k, v)) = pair.trim().split_once('=') {
+            if !k.trim().is_empty() && !v.trim().is_empty() {
+                if let (Ok(key), Ok(value)) = (
+                    HeaderName::from_str(k.trim()),
+                    HeaderValue::from_str(v.trim()),
+                ) {
+                    headers.insert(key, value);
+                }
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{OTEL_EXPORTER_OTLP_ENDPOINT, OTEL_EXPORTER_OTLP_TRACES_ENDPOINT};
@@ -410,4 +452,90 @@ mod tests {
             // You may also want to assert on the specific error type if applicable
         });
     }
+
+    #[test]
+    fn test_add_header_from_string() {
+        use http::{HeaderName, HeaderValue};
+        use std::collections::HashMap;
+        let test_cases = vec![
+            // Format: (input_str, expected_headers)
+            ("k1=v1", vec![("k1", "v1")]),
+            ("k1=v1,k2=v2", vec![("k1", "v1"), ("k2", "v2")]),
+            ("k1=v1=10,k2,k3", vec![("k1", "v1=10")]),
+            ("k1=v1,,,k2,k3=10", vec![("k1", "v1"), ("k3", "10")]),
+        ];
+
+        for (input_str, expected_headers) in test_cases {
+            #[allow(clippy::mutable_key_type)] // http headers are not mutated
+            let mut headers: HashMap<HeaderName, HeaderValue> = HashMap::new();
+            super::add_header_from_string(input_str, &mut headers);
+
+            assert_eq!(
+                headers.len(),
+                expected_headers.len(),
+                "Failed on input: {}",
+                input_str
+            );
+
+            for (expected_key, expected_value) in expected_headers {
+                assert_eq!(
+                    headers.get(&HeaderName::from_static(expected_key)),
+                    Some(&HeaderValue::from_static(expected_value)),
+                    "Failed on key: {} with input: {}",
+                    expected_key,
+                    input_str
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_merge_header_from_string() {
+        use http::{HeaderName, HeaderValue};
+        use std::collections::HashMap;
+        #[allow(clippy::mutable_key_type)] // http headers are not mutated
+        let mut headers: HashMap<HeaderName, HeaderValue> = std::collections::HashMap::new();
+        headers.insert(
+            HeaderName::from_static("k1"),
+            HeaderValue::from_static("v1"),
+        );
+        headers.insert(
+            HeaderName::from_static("k2"),
+            HeaderValue::from_static("v2"),
+        );
+        let test_cases = vec![
+            // Format: (input_str, expected_headers)
+            ("k1=v1_new", vec![("k1", "v1_new"), ("k2", "v2")]),
+            (
+                "k3=val=10,22,34,k4=,k5=10",
+                vec![
+                    ("k1", "v1_new"),
+                    ("k2", "v2"),
+                    ("k3", "val=10"),
+                    ("k5", "10"),
+                ],
+            ),
+        ];
+
+        for (input_str, expected_headers) in test_cases {
+            super::add_header_from_string(input_str, &mut headers);
+
+            assert_eq!(
+                headers.len(),
+                expected_headers.len(),
+                "Failed on input: {}",
+                input_str
+            );
+
+            for (expected_key, expected_value) in expected_headers {
+                assert_eq!(
+                    headers.get(&HeaderName::from_static(expected_key)),
+                    Some(&HeaderValue::from_static(expected_value)),
+                    "Failed on key: {} with input: {}",
+                    expected_key,
+                    input_str
+                );
+            }
+        }
+    }
 }
diff --git a/opentelemetry-otlp/src/exporter/mod.rs b/opentelemetry-otlp/src/exporter/mod.rs
@@ -21,6 +21,10 @@ use std::time::Duration;
 pub const OTEL_EXPORTER_OTLP_ENDPOINT: &str = "OTEL_EXPORTER_OTLP_ENDPOINT";
 /// Default target to which the exporter is going to send signals.
 pub const OTEL_EXPORTER_OTLP_ENDPOINT_DEFAULT: &str = OTEL_EXPORTER_OTLP_HTTP_ENDPOINT_DEFAULT;
+/// Key-value pairs to be used as headers associated with gRPC or HTTP requests
+/// Example: `k1=v1,k2=v2`
+/// Note: as of now, this is only supported for HTTP requests.
+pub const OTEL_EXPORTER_OTLP_HEADERS: &str = "OTEL_EXPORTER_OTLP_HEADERS";
 /// Protocol the exporter will use. Either `http/protobuf` or `grpc`.
 pub const OTEL_EXPORTER_OTLP_PROTOCOL: &str = "OTEL_EXPORTER_OTLP_PROTOCOL";
 /// Compression algorithm to use, defaults to none.
diff --git a/opentelemetry-otlp/src/exporter/tonic/mod.rs b/opentelemetry-otlp/src/exporter/tonic/mod.rs
@@ -215,7 +215,7 @@ impl TonicExporterBuilder {
         signal_compression_var: &str,
     ) -> Result<(Channel, BoxInterceptor, Option<CompressionEncoding>), crate::Error> {
         let config = &mut self.exporter_config;
-        let tonic_config = &mut self.tonic_config;
+        let tonic_config: &mut TonicConfig = &mut self.tonic_config;
 
         let endpoint = match env::var(signal_endpoint_var)
             .ok()
diff --git a/opentelemetry-otlp/src/lib.rs b/opentelemetry-otlp/src/lib.rs
@@ -211,25 +211,27 @@ pub use crate::exporter::ExportConfig;
 #[cfg(feature = "trace")]
 pub use crate::span::{
     OtlpTracePipeline, SpanExporter, SpanExporterBuilder, OTEL_EXPORTER_OTLP_TRACES_COMPRESSION,
-    OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, OTEL_EXPORTER_OTLP_TRACES_TIMEOUT,
+    OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, OTEL_EXPORTER_OTLP_TRACES_HEADERS,
+    OTEL_EXPORTER_OTLP_TRACES_TIMEOUT,
 };
 
 #[cfg(feature = "metrics")]
 pub use crate::metric::{
     MetricsExporter, MetricsExporterBuilder, OtlpMetricPipeline,
     OTEL_EXPORTER_OTLP_METRICS_COMPRESSION, OTEL_EXPORTER_OTLP_METRICS_ENDPOINT,
-    OTEL_EXPORTER_OTLP_METRICS_TIMEOUT,
+    OTEL_EXPORTER_OTLP_METRICS_HEADERS, OTEL_EXPORTER_OTLP_METRICS_TIMEOUT,
 };
 
 #[cfg(feature = "logs")]
 pub use crate::logs::{
     LogExporter, LogExporterBuilder, OtlpLogPipeline, OTEL_EXPORTER_OTLP_LOGS_COMPRESSION,
-    OTEL_EXPORTER_OTLP_LOGS_ENDPOINT, OTEL_EXPORTER_OTLP_LOGS_TIMEOUT,
+    OTEL_EXPORTER_OTLP_LOGS_ENDPOINT, OTEL_EXPORTER_OTLP_LOGS_HEADERS,
+    OTEL_EXPORTER_OTLP_LOGS_TIMEOUT,
 };
 
 pub use crate::exporter::{
     HasExportConfig, WithExportConfig, OTEL_EXPORTER_OTLP_COMPRESSION, OTEL_EXPORTER_OTLP_ENDPOINT,
-    OTEL_EXPORTER_OTLP_ENDPOINT_DEFAULT, OTEL_EXPORTER_OTLP_PROTOCOL,
+    OTEL_EXPORTER_OTLP_ENDPOINT_DEFAULT, OTEL_EXPORTER_OTLP_HEADERS, OTEL_EXPORTER_OTLP_PROTOCOL,
     OTEL_EXPORTER_OTLP_PROTOCOL_DEFAULT, OTEL_EXPORTER_OTLP_TIMEOUT,
     OTEL_EXPORTER_OTLP_TIMEOUT_DEFAULT,
 };
diff --git a/opentelemetry-otlp/src/logs.rs b/opentelemetry-otlp/src/logs.rs
@@ -30,6 +30,12 @@ pub const OTEL_EXPORTER_OTLP_LOGS_ENDPOINT: &str = "OTEL_EXPORTER_OTLP_LOGS_ENDP
 /// Maximum time the OTLP exporter will wait for each batch logs export.
 pub const OTEL_EXPORTER_OTLP_LOGS_TIMEOUT: &str = "OTEL_EXPORTER_OTLP_LOGS_TIMEOUT";
 
+/// Key-value pairs to be used as headers associated with gRPC or HTTP requests
+/// for sending logs.
+/// Example: `k1=v1,k2=v2`
+/// Note: this is only supported for HTTP.
+pub const OTEL_EXPORTER_OTLP_LOGS_HEADERS: &str = "OTEL_EXPORTER_OTLP_LOGS_HEADERS";
+
 impl OtlpPipeline {
     /// Create a OTLP logging pipeline.
     pub fn logging(self) -> OtlpLogPipeline<NoExporterConfig> {
diff --git a/opentelemetry-otlp/src/metric.rs b/opentelemetry-otlp/src/metric.rs
@@ -40,6 +40,11 @@ pub const OTEL_EXPORTER_OTLP_METRICS_ENDPOINT: &str = "OTEL_EXPORTER_OTLP_METRIC
 pub const OTEL_EXPORTER_OTLP_METRICS_TIMEOUT: &str = "OTEL_EXPORTER_OTLP_METRICS_TIMEOUT";
 /// Compression algorithm to use, defaults to none.
 pub const OTEL_EXPORTER_OTLP_METRICS_COMPRESSION: &str = "OTEL_EXPORTER_OTLP_METRICS_COMPRESSION";
+/// Key-value pairs to be used as headers associated with gRPC or HTTP requests
+/// for sending metrics.
+/// Example: `k1=v1,k2=v2`
+/// Note: this is only supported for HTTP.
+pub const OTEL_EXPORTER_OTLP_METRICS_HEADERS: &str = "OTEL_EXPORTER_OTLP_METRICS_HEADERS";
 impl OtlpPipeline {
     /// Create a OTLP metrics pipeline.
     pub fn metrics<RT>(self, rt: RT) -> OtlpMetricPipeline<RT, NoExporterConfig>
diff --git a/opentelemetry-otlp/src/span.rs b/opentelemetry-otlp/src/span.rs
@@ -36,6 +36,11 @@ pub const OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: &str = "OTEL_EXPORTER_OTLP_TRACES_
 pub const OTEL_EXPORTER_OTLP_TRACES_TIMEOUT: &str = "OTEL_EXPORTER_OTLP_TRACES_TIMEOUT";
 /// Compression algorithm to use, defaults to none.
 pub const OTEL_EXPORTER_OTLP_TRACES_COMPRESSION: &str = "OTEL_EXPORTER_OTLP_TRACES_COMPRESSION";
+/// Key-value pairs to be used as headers associated with gRPC or HTTP requests
+/// for sending spans.
+/// Example: `k1=v1,k2=v2`
+/// Note: this is only supported for HTTP.
+pub const OTEL_EXPORTER_OTLP_TRACES_HEADERS: &str = "OTEL_EXPORTER_OTLP_TRACES_HEADERS";
 
 impl OtlpPipeline {
     /// Create a OTLP tracing pipeline.
