diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs index 8037532c09ac3..72f5be3d6bb40 100644 --- a/datafusion/sqllogictest/bin/sqllogictests.rs +++ b/datafusion/sqllogictest/bin/sqllogictests.rs @@ -44,6 +44,7 @@ use datafusion::common::runtime::SpawnedTask; use futures::FutureExt; use std::ffi::OsStr; use std::fs; +use std::io::{IsTerminal, stdout}; use std::path::{Path, PathBuf}; #[cfg(feature = "postgres")] @@ -123,6 +124,8 @@ async fn run_tests() -> Result<()> { .unwrap() .progress_chars("##-"); + let colored_output = options.is_colored(); + let start = Instant::now(); let test_files = read_test_files(&options)?; @@ -176,6 +179,7 @@ async fn run_tests() -> Result<()> { m_style_clone, filters.as_ref(), currently_running_sql_tracker_clone, + colored_output, ) .await? } @@ -187,6 +191,7 @@ async fn run_tests() -> Result<()> { m_style_clone, filters.as_ref(), currently_running_sql_tracker_clone, + colored_output, ) .await? } @@ -294,6 +299,7 @@ async fn run_test_file_substrait_round_trip( mp_style: ProgressStyle, filters: &[Filter], currently_executing_sql_tracker: CurrentlyExecutingSqlTracker, + colored_output: bool, ) -> Result<()> { let TestFile { path, @@ -323,7 +329,7 @@ async fn run_test_file_substrait_round_trip( runner.with_column_validator(strict_column_validator); runner.with_normalizer(value_normalizer); runner.with_validator(validator); - let res = run_file_in_runner(path, runner, filters).await; + let res = run_file_in_runner(path, runner, filters, colored_output).await; pb.finish_and_clear(); res } @@ -335,6 +341,7 @@ async fn run_test_file( mp_style: ProgressStyle, filters: &[Filter], currently_executing_sql_tracker: CurrentlyExecutingSqlTracker, + colored_output: bool, ) -> Result<()> { let TestFile { path, @@ -364,7 +371,7 @@ async fn run_test_file( runner.with_column_validator(strict_column_validator); runner.with_normalizer(value_normalizer); runner.with_validator(validator); - let result = run_file_in_runner(path, runner, filters).await; + let result = run_file_in_runner(path, runner, filters, colored_output).await; pb.finish_and_clear(); result } @@ -373,6 +380,7 @@ async fn run_file_in_runner>( path: PathBuf, mut runner: sqllogictest::Runner, filters: &[Filter], + colored_output: bool, ) -> Result<()> { let path = path.canonicalize()?; let records = @@ -386,7 +394,11 @@ async fn run_file_in_runner>( continue; } if let Err(err) = runner.run_async(record).await { - errs.push(format!("{err}")); + if colored_output { + errs.push(format!("{}", err.display(true))); + } else { + errs.push(format!("{err}")); + } } } @@ -479,7 +491,7 @@ async fn run_test_file_with_postgres( runner.with_column_validator(strict_column_validator); runner.with_normalizer(value_normalizer); runner.with_validator(validator); - let result = run_file_in_runner(path, runner, filters).await; + let result = run_file_in_runner(path, runner, filters, false).await; pb.finish_and_clear(); result } @@ -772,6 +784,13 @@ struct Options { default_value_t = get_available_parallelism() )] test_threads: usize, + + #[clap( + long, + value_name = "WHEN", + help = "Enable or disable coloured output: auto (default), always, never" + )] + color: Option, } impl Options { @@ -813,6 +832,29 @@ impl Options { eprintln!("WARNING: Ignoring `--show-output` compatibility option"); } } + + /// Determine if colour output should be enabled, respecting --color, NO_COLOR, CARGO_TERM_COLOR, and terminal detection + fn is_colored(&self) -> bool { + // NO_COLOR takes precedence + if std::env::var_os("NO_COLOR").is_some() { + return false; + } + // CLI flag `--color` takes precedence over `CARGO_TERM_COLOR` + let color_pref: String = match &self.color { + None => std::env::var("CARGO_TERM_COLOR").ok(), + Some(s) => Some(s.clone()), + } + .unwrap_or_else(|| "auto".to_string()); + match color_pref.as_str() { + "always" => true, + "never" => false, + _ => { + // Use colors by default for non-dumb terminals + stdout().is_terminal() + && std::env::var("TERM").unwrap_or_default() != "dumb" + } + } + } } /// Performs scratch file check for all test files.