Skip to content

Commit 841e282

Browse files
matthewmturnermatthewmturner
committed
Headers file
1 parent a8eedd3 commit 841e282

5 files changed

Lines changed: 509 additions & 5 deletions

File tree

‎src/args.rs‎

Lines changed: 217 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,12 @@ pub struct DftArgs {
109109
)]
110110
pub header: Option<Vec<(String, String)>>,
111111

112+
#[clap(
113+
long,
114+
help = "Path to file containing Flight SQL headers. Supports simple format ('Name: Value') and curl config format ('header = Name: Value' or '-H \"Name: Value\"'). Only used for FlightSQL"
115+
)]
116+
pub headers_file: Option<PathBuf>,
117+
112118
#[clap(
113119
long,
114120
short,
@@ -257,7 +263,7 @@ fn parse_command(command: &str) -> std::result::Result<String, String> {
257263
fn parse_header_line(line: &str) -> Result<(String, String), String> {
258264
let (name, value) = line
259265
.split_once(':')
260-
.ok_or_else(|| format!("Invalid header format: '{}'", line))?;
266+
.ok_or_else(|| format!("Invalid header format: '{}'\n Expected format: 'Header-Name: Header-Value', 'header = Name: Value', or '-H \"Name: Value\"'", line))?;
261267

262268
let name =
263269
HeaderName::try_from(name.trim()).map_err(|e| format!("Invalid header name: {}", e))?;
@@ -270,3 +276,213 @@ fn parse_header_line(line: &str) -> Result<(String, String), String> {
270276

271277
Ok((name.to_string(), value_str.to_string()))
272278
}
279+
280+
/// Parse headers from a file supporting both simple and curl config formats
281+
///
282+
/// Supported formats:
283+
/// - Simple: `Header-Name: Header-Value`
284+
/// - Curl config: `header = Name: Value` or `-H "Name: Value"`
285+
/// - Comments: Lines starting with `#`
286+
/// - Blank lines are ignored
287+
///
288+
/// Both formats can be mixed in the same file.
289+
pub fn parse_headers_file(path: &Path) -> Result<Vec<(String, String)>, String> {
290+
let content = std::fs::read_to_string(path)
291+
.map_err(|e| format!("Failed to read headers file '{}': {}", path.display(), e))?;
292+
293+
let mut headers = Vec::new();
294+
for (line_num, line) in content.lines().enumerate() {
295+
let line = line.trim();
296+
297+
// Skip comments and blank lines
298+
if line.is_empty() || line.starts_with('#') {
299+
continue;
300+
}
301+
302+
// Detect and parse format
303+
let header_value = if let Some(stripped) = line.strip_prefix("header") {
304+
// Curl config format: "header = Name: Value" or "header=Name: Value"
305+
let stripped = stripped.trim_start();
306+
if let Some(value) = stripped.strip_prefix('=') {
307+
value.trim()
308+
} else {
309+
line // Not curl format, try simple format
310+
}
311+
} else if let Some(stripped) = line.strip_prefix("-H") {
312+
// Curl config format: -H "Name: Value" or -H Name: Value
313+
let stripped = stripped.trim();
314+
// Remove surrounding quotes if present
315+
stripped.trim_matches(|c| c == '"' || c == '\'')
316+
} else {
317+
// Simple format: "Name: Value"
318+
line
319+
};
320+
321+
// Parse header line
322+
match parse_header_line(header_value) {
323+
Ok(header) => headers.push(header),
324+
Err(e) => {
325+
return Err(format!(
326+
"Invalid header format at line {} in '{}': '{}'\n{}",
327+
line_num + 1,
328+
path.display(),
329+
line,
330+
e
331+
));
332+
}
333+
}
334+
}
335+
336+
Ok(headers)
337+
}
338+
339+
#[cfg(test)]
340+
mod tests {
341+
use super::*;
342+
use std::io::Write;
343+
use tempfile::NamedTempFile;
344+
345+
#[test]
346+
fn test_parse_headers_file_simple_format() {
347+
let mut file = NamedTempFile::new().unwrap();
348+
writeln!(file, "x-api-key: secret123").unwrap();
349+
writeln!(file, "database: production").unwrap();
350+
file.flush().unwrap();
351+
352+
let headers = parse_headers_file(file.path()).unwrap();
353+
assert_eq!(headers.len(), 2);
354+
assert_eq!(
355+
headers[0],
356+
("x-api-key".to_string(), "secret123".to_string())
357+
);
358+
assert_eq!(
359+
headers[1],
360+
("database".to_string(), "production".to_string())
361+
);
362+
}
363+
364+
#[test]
365+
fn test_parse_headers_file_curl_format() {
366+
let mut file = NamedTempFile::new().unwrap();
367+
writeln!(file, "header = x-api-key: secret123").unwrap();
368+
writeln!(file, "-H \"database: production\"").unwrap();
369+
file.flush().unwrap();
370+
371+
let headers = parse_headers_file(file.path()).unwrap();
372+
assert_eq!(headers.len(), 2);
373+
assert_eq!(
374+
headers[0],
375+
("x-api-key".to_string(), "secret123".to_string())
376+
);
377+
assert_eq!(
378+
headers[1],
379+
("database".to_string(), "production".to_string())
380+
);
381+
}
382+
383+
#[test]
384+
fn test_parse_headers_file_mixed_format() {
385+
let mut file = NamedTempFile::new().unwrap();
386+
writeln!(file, "# Simple format").unwrap();
387+
writeln!(file, "x-test: value1").unwrap();
388+
writeln!(file, "").unwrap();
389+
writeln!(file, "# Curl config format").unwrap();
390+
writeln!(file, "header = x-api-key: secret123").unwrap();
391+
writeln!(file, "-H \"database: production\"").unwrap();
392+
file.flush().unwrap();
393+
394+
let headers = parse_headers_file(file.path()).unwrap();
395+
assert_eq!(headers.len(), 3);
396+
assert_eq!(headers[0], ("x-test".to_string(), "value1".to_string()));
397+
assert_eq!(
398+
headers[1],
399+
("x-api-key".to_string(), "secret123".to_string())
400+
);
401+
assert_eq!(
402+
headers[2],
403+
("database".to_string(), "production".to_string())
404+
);
405+
}
406+
407+
#[test]
408+
fn test_parse_headers_file_with_comments() {
409+
let mut file = NamedTempFile::new().unwrap();
410+
writeln!(file, "# This is a comment").unwrap();
411+
writeln!(file, "x-api-key: secret123").unwrap();
412+
writeln!(file, "# Another comment").unwrap();
413+
writeln!(file, "database: production").unwrap();
414+
file.flush().unwrap();
415+
416+
let headers = parse_headers_file(file.path()).unwrap();
417+
assert_eq!(headers.len(), 2);
418+
assert_eq!(
419+
headers[0],
420+
("x-api-key".to_string(), "secret123".to_string())
421+
);
422+
assert_eq!(
423+
headers[1],
424+
("database".to_string(), "production".to_string())
425+
);
426+
}
427+
428+
#[test]
429+
fn test_parse_headers_file_blank_lines() {
430+
let mut file = NamedTempFile::new().unwrap();
431+
writeln!(file, "x-api-key: secret123").unwrap();
432+
writeln!(file, "").unwrap();
433+
writeln!(file, " ").unwrap();
434+
writeln!(file, "database: production").unwrap();
435+
file.flush().unwrap();
436+
437+
let headers = parse_headers_file(file.path()).unwrap();
438+
assert_eq!(headers.len(), 2);
439+
assert_eq!(
440+
headers[0],
441+
("x-api-key".to_string(), "secret123".to_string())
442+
);
443+
assert_eq!(
444+
headers[1],
445+
("database".to_string(), "production".to_string())
446+
);
447+
}
448+
449+
#[test]
450+
fn test_parse_headers_file_curl_with_quotes() {
451+
let mut file = NamedTempFile::new().unwrap();
452+
writeln!(file, "-H \"x-api-key: secret123\"").unwrap();
453+
writeln!(file, "-H 'database: production'").unwrap();
454+
file.flush().unwrap();
455+
456+
let headers = parse_headers_file(file.path()).unwrap();
457+
assert_eq!(headers.len(), 2);
458+
assert_eq!(
459+
headers[0],
460+
("x-api-key".to_string(), "secret123".to_string())
461+
);
462+
assert_eq!(
463+
headers[1],
464+
("database".to_string(), "production".to_string())
465+
);
466+
}
467+
468+
#[test]
469+
fn test_parse_headers_file_invalid_format() {
470+
let mut file = NamedTempFile::new().unwrap();
471+
writeln!(file, "x-api-key: secret123").unwrap();
472+
writeln!(file, "invalid-line-without-colon").unwrap();
473+
file.flush().unwrap();
474+
475+
let result = parse_headers_file(file.path());
476+
assert!(result.is_err());
477+
assert!(result
478+
.unwrap_err()
479+
.contains("Invalid header format at line 2"));
480+
}
481+
482+
#[test]
483+
fn test_parse_headers_file_not_found() {
484+
let result = parse_headers_file(Path::new("/nonexistent/file.txt"));
485+
assert!(result.is_err());
486+
assert!(result.unwrap_err().contains("Failed to read headers file"));
487+
}
488+
}

‎src/cli/mod.rs‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ use std::io::Write;
4141
use std::path::{Path, PathBuf};
4242
#[cfg(feature = "flightsql")]
4343
use {
44-
crate::args::{Command, FlightSqlCommand},
44+
crate::args::{parse_headers_file, Command, FlightSqlCommand},
4545
datafusion_app::{
4646
config::{AuthConfig, FlightSQLConfig},
4747
flightsql::FlightSQLContext,
@@ -870,10 +870,41 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
870870
config.flightsql_client.connection_url,
871871
config.flightsql_client.benchmark_iterations,
872872
auth,
873-
config.flightsql_client.headers,
873+
config.flightsql_client.headers.clone(),
874874
);
875875
let flightsql_ctx = FlightSQLContext::new(flightsql_cfg);
876-
let headers = cli.header.clone().map(|vec| vec.into_iter().collect());
876+
877+
// Three-way header merge: config < file < CLI
878+
let mut all_headers = config.flightsql_client.headers.clone();
879+
880+
// Merge headers from file (if specified in config or CLI)
881+
let headers_file = cli
882+
.headers_file
883+
.as_ref()
884+
.or(config.flightsql_client.headers_file.as_ref());
885+
886+
if let Some(file_path) = headers_file {
887+
match parse_headers_file(file_path) {
888+
Ok(file_headers) => {
889+
all_headers.extend(file_headers);
890+
}
891+
Err(e) => {
892+
return Err(eyre!("Error reading headers file: {}", e));
893+
}
894+
}
895+
}
896+
897+
// Merge CLI headers (highest precedence)
898+
if let Some(cli_headers) = &cli.header {
899+
all_headers.extend(cli_headers.iter().cloned());
900+
}
901+
902+
let headers = if all_headers.is_empty() {
903+
None
904+
} else {
905+
Some(all_headers)
906+
};
907+
877908
flightsql_ctx
878909
.create_client(cli.host.clone(), headers)
879910
.await?;

‎src/config.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,8 @@ pub struct FlightSQLClientConfig {
120120
pub auth: AuthConfig,
121121
#[serde(default = "default_headers")]
122122
pub headers: HashMap<String, String>,
123+
#[serde(default)]
124+
pub headers_file: Option<PathBuf>,
123125
}
124126

125127
#[cfg(feature = "flightsql")]
@@ -130,6 +132,7 @@ impl Default for FlightSQLClientConfig {
130132
benchmark_iterations: default_benchmark_iterations(),
131133
auth: default_auth_config(),
132134
headers: default_headers(),
135+
headers_file: None,
133136
}
134137
}
135138
}

‎src/tui/mod.rs‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -389,14 +389,36 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
389389

390390
#[cfg(feature = "flightsql")]
391391
{
392+
use crate::args::parse_headers_file;
392393
use datafusion_app::config::FlightSQLConfig;
393394
use datafusion_app::flightsql::FlightSQLContext;
394395

396+
// Merge headers: config < file (CLI headers are merged later in the handler)
397+
let mut all_headers = config.flightsql_client.headers.clone();
398+
399+
// Load headers from file if specified in config or CLI args
400+
let headers_file = cli
401+
.headers_file
402+
.as_ref()
403+
.or(config.flightsql_client.headers_file.as_ref());
404+
405+
if let Some(file_path) = headers_file {
406+
match parse_headers_file(file_path) {
407+
Ok(file_headers) => {
408+
all_headers.extend(file_headers);
409+
}
410+
Err(e) => {
411+
// TUI silently logs file errors to avoid disrupting UI
412+
error!("Error reading headers file: {}", e);
413+
}
414+
}
415+
}
416+
395417
let flightsql_config = FlightSQLConfig::new(
396418
config.flightsql_client.connection_url.clone(),
397419
config.flightsql_client.benchmark_iterations,
398420
config.flightsql_client.auth.clone(),
399-
config.flightsql_client.headers.clone(),
421+
all_headers,
400422
);
401423
app_execution.with_flightsql_ctx(FlightSQLContext::new(flightsql_config));
402424
}

0 commit comments

Comments
 (0)
⚡