feat: support sqllogictest output coloring

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<D: AsyncDB, M: MakeConnection<Conn = D>>(
     path: PathBuf,
     mut runner: sqllogictest::Runner<D, M>,
     filters: &[Filter],
+    colored_output: bool,
 ) -> Result<()> {
     let path = path.canonicalize()?;
     let records =
@@ -386,7 +394,11 @@ async fn run_file_in_runner<D: AsyncDB, M: MakeConnection<Conn = D>>(
             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<String>,
 }
 
 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.