feat: add client headers to config (#333)

Related to #332

---------

Co-authored-by: Jules Cheron <jules.cheron@deepki.com>

Author: jules-ch

crates/datafusion-app/src/config.rs

@@ -252,15 +252,22 @@ pub struct FlightSQLConfig {
     pub connection_url: String,
     pub benchmark_iterations: usize,
     pub auth: AuthConfig,
+    pub headers: HashMap<String, String>,
 }
 
 #[cfg(feature = "flightsql")]
 impl FlightSQLConfig {
-    pub fn new(connection_url: String, benchmark_iterations: usize, auth: AuthConfig) -> Self {
+    pub fn new(
+        connection_url: String,
+        benchmark_iterations: usize,
+        auth: AuthConfig,
+        headers: HashMap<String, String>,
+    ) -> Self {
         Self {
             connection_url,
             benchmark_iterations,
             auth,
+            headers,
         }
     }
 }

crates/datafusion-app/src/flightsql.rs

@@ -31,14 +31,14 @@ use datafusion::{
 };
 use log::{debug, error, info, warn};
 
+#[cfg(feature = "flightsql")]
+use crate::config::BasicAuth;
 use color_eyre::eyre::{self, Result};
+use std::collections::HashMap;
 use tokio::sync::Mutex;
 use tokio_stream::StreamExt;
 use tonic::{transport::Channel, IntoRequest};
 
-#[cfg(feature = "flightsql")]
-use crate::config::BasicAuth;
-
 use crate::{
     config::FlightSQLConfig, flightsql_benchmarks::FlightSQLBenchmarkStats, ExecOptions, ExecResult,
 };
@@ -65,7 +65,11 @@ impl FlightSQLContext {
 
     // TODO - Make this part of `new` method
     /// Create FlightSQL client from users FlightSQL config
-    pub async fn create_client(&self, cli_host: Option<String>) -> Result<()> {
+    pub async fn create_client(
+        &self,
+        cli_host: Option<String>,
+        cli_headers: Option<HashMap<String, String>>,
+    ) -> Result<()> {
         let final_url = cli_host.unwrap_or(self.config.connection_url.clone());
         let url = Box::leak(final_url.into_boxed_str());
         info!("Connecting to FlightSQL host: {}", url);
@@ -89,6 +93,14 @@ impl FlightSQLContext {
                         let encoded_basic = STANDARD.encode(format!("{username}:{password}"));
                         client.set_header("Authorization", format!("Basic {encoded_basic}"))
                     }
+
+                    let mut headers = self.config.headers.clone();
+                    if let Some(cli) = cli_headers {
+                        headers.extend(cli);
+                    }
+                    for (name, value) in headers {
+                        client.set_header(name, value);
+                    }
                 }
                 let mut guard = self.client.lock().await;
                 *guard = Some(client);

src/args.rs

@@ -19,6 +19,7 @@
 
 use crate::config::get_data_dir;
 use clap::{Parser, Subcommand};
+use http::{HeaderName, HeaderValue};
 #[cfg(any(feature = "http", feature = "flightsql"))]
 use std::net::SocketAddr;
 use std::path::{Path, PathBuf};
@@ -97,6 +98,14 @@ pub struct DftArgs {
     #[clap(long, help = "Host address to query. Only used for FlightSQL")]
     pub host: Option<String>,
 
+    #[clap(
+        long,
+        help = "Header to add to Flight SQL connection. Only used for FlightSQL",
+        value_parser(parse_header_line),
+        action = clap::ArgAction::Append
+    )]
+    pub header: Option<Vec<(String, String)>>,
+
     #[clap(
         long,
         short,
@@ -218,3 +227,20 @@ fn parse_command(command: &str) -> std::result::Result<String, String> {
         Err("-c flag expects only non empty commands".to_string())
     }
 }
+
+fn parse_header_line(line: &str) -> Result<(String, String), String> {
+    let (name, value) = line
+        .split_once(':')
+        .ok_or_else(|| format!("Invalid header format: '{}'", line))?;
+
+    let name =
+        HeaderName::try_from(name.trim()).map_err(|e| format!("Invalid header name: {}", e))?;
+    let value =
+        HeaderValue::try_from(value.trim()).map_err(|e| format!("Invalid header value: {}", e))?;
+
+    let value_str = value
+        .to_str()
+        .map_err(|e| format!("Header value contains invalid characters: {}", e))?;
+
+    Ok((name.to_string(), value_str.to_string()))
+}

src/cli/mod.rs

@@ -687,9 +687,13 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
                 config.flightsql_client.connection_url,
                 config.flightsql_client.benchmark_iterations,
                 auth,
+                config.flightsql_client.headers,
             );
             let flightsql_ctx = FlightSQLContext::new(flightsql_cfg);
-            flightsql_ctx.create_client(cli.host.clone()).await?;
+            let headers = cli.header.clone().map(|vec| vec.into_iter().collect());
+            flightsql_ctx
+                .create_client(cli.host.clone(), headers)
+                .await?;
             app_execution.with_flightsql_ctx(flightsql_ctx);
         }
     }

src/config.rs

@@ -27,6 +27,7 @@ use directories::{ProjectDirs, UserDirs};
 use lazy_static::lazy_static;
 use log::{debug, error};
 use serde::Deserialize;
+use std::collections::HashMap;
 
 #[cfg(any(feature = "flightsql", feature = "http"))]
 use datafusion_app::config::AuthConfig;
@@ -115,6 +116,8 @@ pub struct FlightSQLClientConfig {
     pub benchmark_iterations: usize,
     #[serde(default = "default_auth_config")]
     pub auth: AuthConfig,
+    #[serde(default = "default_headers")]
+    pub headers: HashMap<String, String>,
 }
 
 #[cfg(feature = "flightsql")]
@@ -124,6 +127,7 @@ impl Default for FlightSQLClientConfig {
             connection_url: default_connection_url(),
             benchmark_iterations: default_benchmark_iterations(),
             auth: default_auth_config(),
+            headers: default_headers(),
         }
     }
 }
@@ -266,6 +270,11 @@ pub fn default_connection_url() -> String {
     "http://localhost:50051".to_string()
 }
 
+#[cfg(feature = "flightsql")]
+pub fn default_headers() -> HashMap<String, String> {
+    HashMap::new()
+}
+
 #[cfg(any(feature = "flightsql", feature = "http"))]
 fn default_server_metrics_addr() -> SocketAddr {
     SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9000)

src/server/http/mod.rs

@@ -142,12 +142,16 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
             config.flightsql_client.connection_url.clone(),
             config.flightsql_client.benchmark_iterations,
             auth,
+            config.flightsql_client.headers.clone(),
         );
 
         let flightsql_context = FlightSQLContext::new(flightsql_cfg.clone());
         // TODO - Consider adding flag to allow startup even if FlightSQL initiation fails
         if let Err(e) = flightsql_context
-            .create_client(Some(flightsql_cfg.connection_url))
+            .create_client(
+                Some(flightsql_cfg.connection_url),
+                Some(flightsql_cfg.headers),
+            )
             .await
         {
             error!("{}", e.to_string())

src/server/http/router.rs

@@ -463,7 +463,7 @@ mod flightsql_test {
         };
         let flightsql_ctx = FlightSQLContext::new(flightsql_cfg);
         flightsql_ctx
-            .create_client(Some("http://localhost:50051".to_string()))
+            .create_client(Some("http://localhost:50051".to_string()), None)
             .await
             .unwrap();
         execution.with_flightsql_ctx(flightsql_ctx);

src/tui/execution.rs

@@ -28,6 +28,8 @@ use datafusion::execution::SendableRecordBatchStream;
 use datafusion::physical_plan::execute_stream;
 use futures::StreamExt;
 use log::{error, info};
+#[cfg(feature = "flightsql")]
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::mpsc::UnboundedSender;
@@ -404,8 +406,15 @@ impl TuiExecution {
     // TODO: Maybe just expose `inner` and use that rather than re-implementing the same
     // functions here.
     #[cfg(feature = "flightsql")]
-    pub async fn create_flightsql_client(&self, cli_host: Option<String>) -> Result<()> {
-        self.inner.flightsql_ctx().create_client(cli_host).await
+    pub async fn create_flightsql_client(
+        &self,
+        cli_host: Option<String>,
+        cli_headers: Option<HashMap<String, String>>,
+    ) -> Result<()> {
+        self.inner
+            .flightsql_ctx()
+            .create_client(cli_host, cli_headers)
+            .await
     }
 
     #[cfg(feature = "flightsql")]

src/tui/handlers/mod.rs

@@ -246,8 +246,9 @@ pub fn app_event_handler(app: &mut App, event: AppEvent) -> Result<()> {
             let execution = Arc::clone(&app.execution);
             let _event_tx = app.event_tx.clone();
             let host = app.args.host.clone();
+            let headers = app.args.header.clone().map(|vec| vec.into_iter().collect());
             tokio::spawn(async move {
-                if let Err(e) = execution.create_flightsql_client(host).await {
+                if let Err(e) = execution.create_flightsql_client(host, headers).await {
                     error!("Error creating FlightSQL client: {:?}", e);
                     if let Err(e) = _event_tx.send(AppEvent::FlightSQLFailedToConnect) {
                         error!("Error sending FlightSQLFailedToConnect message: {e}");

tests/config.rs

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#[cfg(feature = "flightsql")]
+use std::collections::HashMap;
 use std::{io::Write, path::PathBuf};
 
 use tempfile::{tempdir, TempDir};
@@ -198,6 +200,19 @@ impl TestConfigBuilder {
         self
     }
 
+    #[cfg(feature = "flightsql")]
+    pub fn with_client_headers(&mut self, headers: Option<HashMap<String, String>>) -> &mut Self {
+        self.config_text.push_str("[flightsql_client.headers]\n");
+
+        if let Some(headers) = &headers {
+            for (name, value) in headers {
+                self.config_text.push_str(&format!("{name} = {value}\n"));
+            }
+        }
+
+        self
+    }
+
     // TODO: Update this to work with HTTP server
     #[allow(dead_code)]
     #[cfg(feature = "flightsql")]