From 4bbe756b521d7e3cfc4ea5a64c99bad63f84fb47 Mon Sep 17 00:00:00 2001
From: Marco Neumann <marco@crepererum.net>
Date: Wed, 1 Dec 2021 15:02:15 +0100
Subject: [PATCH] feat: make jaeger-debug-id configurable

---
 influxdb2_client/src/lib.rs                   | 20 +++++++++----------
 .../src/bin/iox_data_generator.rs             |  5 +++--
 iox_data_generator/src/write.rs               |  6 +++---
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/influxdb2_client/src/lib.rs b/influxdb2_client/src/lib.rs
index fd75e6584b..6e36bb3fe5 100644
--- a/influxdb2_client/src/lib.rs
+++ b/influxdb2_client/src/lib.rs
@@ -100,10 +100,13 @@ pub struct Client {
     pub url: String,
     auth_header: Option<String>,
     reqwest: reqwest::Client,
-    jaeger_debug: bool,
+    jaeger_debug_header: Option<String>,
 }
 
 impl Client {
+    /// Default [jaeger debug header](Self::with_jaeger_debug) that should work in many environments.
+    pub const DEFAULT_JAEGER_DEBUG_HEADER: &'static str = "jaeger-debug-id";
+
     /// Create a new client pointing to the URL specified in
     /// `protocol://server:port` format and using the specified token for
     /// authorization.
@@ -125,14 +128,14 @@ impl Client {
             url: url.into(),
             auth_header,
             reqwest: reqwest::Client::new(),
-            jaeger_debug: false,
+            jaeger_debug_header: None,
         }
     }
 
-    /// Enable generation of `jaeger-debug-id` headers.
-    pub fn with_jaeger_debug(self) -> Self {
+    /// Enable generation of jaeger debug headers with the given header name.
+    pub fn with_jaeger_debug(self, header: String) -> Self {
         Self {
-            jaeger_debug: true,
+            jaeger_debug_header: Some(header),
             ..self
         }
     }
@@ -144,11 +147,8 @@ impl Client {
         if let Some(auth) = &self.auth_header {
             req = req.header("Authorization", auth);
         }
-        if self.jaeger_debug {
-            req = req.header(
-                "jaeger-debug-id",
-                format!("influxdb_client-{}", uuid::Uuid::new_v4()),
-            );
+        if let Some(header) = &self.jaeger_debug_header {
+            req = req.header(header, format!("influxdb_client-{}", uuid::Uuid::new_v4()));
         }
 
         req
diff --git a/iox_data_generator/src/bin/iox_data_generator.rs b/iox_data_generator/src/bin/iox_data_generator.rs
index e18e957ae2..4f6bc9c5a1 100644
--- a/iox_data_generator/src/bin/iox_data_generator.rs
+++ b/iox_data_generator/src/bin/iox_data_generator.rs
@@ -142,7 +142,8 @@ Logging:
         .arg(
             Arg::with_name("jaeger_debug")
                 .long("jaeger_debug")
-                .help("Generate `jaeger-debug-id` headers during write")
+                .help("Generate jaeger debug header with given key during write")
+                .takes_value(true)
         )
         .get_matches();
 
@@ -203,7 +204,7 @@ Logging:
             token,
             create_bucket,
             org_id,
-            matches.is_present("jaeger_debug"),
+            matches.value_of("jaeger_debug"),
         )
         .await?
     } else if matches.is_present("PRINT") {
diff --git a/iox_data_generator/src/write.rs b/iox_data_generator/src/write.rs
index f3675e7f01..7b9f79fb04 100644
--- a/iox_data_generator/src/write.rs
+++ b/iox_data_generator/src/write.rs
@@ -118,7 +118,7 @@ impl PointsWriterBuilder {
         token: impl Into<String> + Send,
         create_bucket: bool,
         org_id: Option<&str>,
-        jaeger_debug: bool,
+        jaeger_debug: Option<&str>,
     ) -> Result<Self> {
         let host = host.into();
 
@@ -132,8 +132,8 @@ impl PointsWriterBuilder {
         };
 
         let mut client = influxdb2_client::Client::new(host, token.into());
-        if jaeger_debug {
-            client = client.with_jaeger_debug();
+        if let Some(header) = jaeger_debug {
+            client = client.with_jaeger_debug(header.to_string());
         }
         let org = org.into();
         let bucket = bucket.into();