diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index aabf69aac888..a175f99a90d8 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -128,6 +128,7 @@ pub async fn exec_from_repl( let mut rl = Editor::new()?; rl.set_helper(Some(CliHelper::new( &ctx.task_ctx().session_config().options().sql_parser.dialect, + print_options.color, ))); rl.load_history(".history").ok(); diff --git a/datafusion-cli/src/helper.rs b/datafusion-cli/src/helper.rs index 0e146d575718..a8e149b4c5c6 100644 --- a/datafusion-cli/src/helper.rs +++ b/datafusion-cli/src/helper.rs @@ -38,20 +38,25 @@ use rustyline::Context; use rustyline::Helper; use rustyline::Result; -use crate::highlighter::SyntaxHighlighter; +use crate::highlighter::{NoSyntaxHighlighter, SyntaxHighlighter}; pub struct CliHelper { completer: FilenameCompleter, dialect: String, - highlighter: SyntaxHighlighter, + highlighter: Box, } impl CliHelper { - pub fn new(dialect: &str) -> Self { + pub fn new(dialect: &str, color: bool) -> Self { + let highlighter: Box = if !color { + Box::new(NoSyntaxHighlighter {}) + } else { + Box::new(SyntaxHighlighter::new(dialect)) + }; Self { completer: FilenameCompleter::new(), dialect: dialect.into(), - highlighter: SyntaxHighlighter::new(dialect), + highlighter, } } @@ -102,7 +107,7 @@ impl CliHelper { impl Default for CliHelper { fn default() -> Self { - Self::new("generic") + Self::new("generic", false) } } diff --git a/datafusion-cli/src/highlighter.rs b/datafusion-cli/src/highlighter.rs index 28732d5b976f..0bb75510b524 100644 --- a/datafusion-cli/src/highlighter.rs +++ b/datafusion-cli/src/highlighter.rs @@ -30,20 +30,22 @@ use datafusion::sql::sqlparser::{ use rustyline::highlight::Highlighter; /// The syntax highlighter. +#[derive(Debug)] pub struct SyntaxHighlighter { dialect: Box, } impl SyntaxHighlighter { pub fn new(dialect: &str) -> Self { - let dialect = match dialect_from_str(dialect) { - Some(dialect) => dialect, - None => Box::new(GenericDialect {}), - }; + let dialect = dialect_from_str(dialect).unwrap_or(Box::new(GenericDialect {})); Self { dialect } } } +pub struct NoSyntaxHighlighter {} + +impl Highlighter for NoSyntaxHighlighter {} + impl Highlighter for SyntaxHighlighter { fn highlight<'l>(&self, line: &'l str, _: usize) -> Cow<'l, str> { let mut out_line = String::new(); diff --git a/datafusion-cli/src/lib.rs b/datafusion-cli/src/lib.rs index 61f9eae7dd53..139a60b8cf16 100644 --- a/datafusion-cli/src/lib.rs +++ b/datafusion-cli/src/lib.rs @@ -23,8 +23,7 @@ pub mod command; pub mod exec; pub mod functions; pub mod helper; +pub mod highlighter; pub mod object_storage; pub mod print_format; pub mod print_options; - -mod highlighter; diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index a9082f2e5351..38537dbd9238 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -136,6 +136,9 @@ struct Args { default_value = "40" )] maxrows: MaxRows, + + #[clap(long, help = "Enables console syntax highlighting")] + color: bool, } #[tokio::main] @@ -169,28 +172,28 @@ async fn main_inner() -> Result<()> { session_config = session_config.with_batch_size(batch_size); }; - let rn_config = RuntimeConfig::new(); - let rn_config = + let rt_config = RuntimeConfig::new(); + let rt_config = // set memory pool size if let Some(memory_limit) = args.memory_limit { let memory_limit = extract_memory_pool_size(&memory_limit).unwrap(); // set memory pool type if let Some(mem_pool_type) = args.mem_pool_type { match mem_pool_type { - PoolType::Greedy => rn_config + PoolType::Greedy => rt_config .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))), - PoolType::Fair => rn_config + PoolType::Fair => rt_config .with_memory_pool(Arc::new(FairSpillPool::new(memory_limit))), } } else { - rn_config + rt_config .with_memory_pool(Arc::new(GreedyMemoryPool::new(memory_limit))) } } else { - rn_config + rt_config }; - let runtime_env = create_runtime_env(rn_config.clone())?; + let runtime_env = create_runtime_env(rt_config.clone())?; let mut ctx = SessionContext::new_with_config_rt(session_config.clone(), Arc::new(runtime_env)); @@ -207,6 +210,7 @@ async fn main_inner() -> Result<()> { format: args.format, quiet: args.quiet, maxrows: args.maxrows, + color: args.color, }; let commands = args.command; diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index b382eb34f62c..f8cd9b3258ef 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -70,6 +70,7 @@ pub struct PrintOptions { pub format: PrintFormat, pub quiet: bool, pub maxrows: MaxRows, + pub color: bool, } fn get_timing_info_str( diff --git a/docs/source/user-guide/cli.md b/docs/source/user-guide/cli.md index 95b3e7125cd4..30ab7d1495a5 100644 --- a/docs/source/user-guide/cli.md +++ b/docs/source/user-guide/cli.md @@ -111,6 +111,9 @@ OPTIONS: -c, --command ... Execute the given command string(s), then exit + --color + Enables console syntax highlighting + -f, --file ... Execute commands from file(s), then exit