From e5cce59c23076ee3fe216e70b953df5a8404f4db Mon Sep 17 00:00:00 2001 From: Mosha Pasumansky Date: Thu, 21 Aug 2025 16:33:26 -0700 Subject: [PATCH 1/2] Implement \set PROMPT1, PROMPT2, and PROMPT3 commands for Firebolt CLI --- src/context.rs | 17 ++- src/main.rs | 32 +++- src/meta_commands.rs | 340 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 385 insertions(+), 4 deletions(-) create mode 100644 src/meta_commands.rs diff --git a/src/context.rs b/src/context.rs index a37fd5a..b9b868c 100644 --- a/src/context.rs +++ b/src/context.rs @@ -13,17 +13,32 @@ pub struct Context { pub args: Args, pub url: String, pub sa_token: Option, + pub prompt1: Option, + pub prompt2: Option, + pub prompt3: Option, } impl Context { pub fn new(args: Args) -> Self { let url = get_url(&args); - Self { args, url, sa_token: None } + Self { args, url, sa_token: None, prompt1: None, prompt2: None, prompt3: None } } pub fn update_url(&mut self) { self.url = get_url(&self.args); } + + pub fn set_prompt1(&mut self, prompt: String) { + self.prompt1 = Some(prompt); + } + + pub fn set_prompt2(&mut self, prompt: String) { + self.prompt2 = Some(prompt); + } + + pub fn set_prompt3(&mut self, prompt: String) { + self.prompt3 = Some(prompt); + } } #[cfg(test)] diff --git a/src/main.rs b/src/main.rs index f0ee579..70edc9b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,12 +3,14 @@ use rustyline::{config::Configurer, error::ReadlineError, Cmd, DefaultEditor, Ev mod args; mod auth; mod context; +mod meta_commands; mod query; mod utils; use args::get_args; use auth::maybe_authenticate; use context::Context; +use meta_commands::handle_meta_command; use query::{query, try_split_queries}; use utils::history_path; @@ -57,11 +59,26 @@ async fn main() -> Result<(), Box> { let mut buffer: String = String::new(); loop { let prompt = if !buffer.trim_start().is_empty() { - "~> " + // Continuation prompt (PROMPT2) + if let Some(custom_prompt) = &context.prompt2 { + custom_prompt.as_str() + } else { + "~> " + } } else if context.args.extra.iter().any(|arg| arg.starts_with("transaction_id=")) { - "*> " + // Transaction prompt (PROMPT3) + if let Some(custom_prompt) = &context.prompt3 { + custom_prompt.as_str() + } else { + "*> " + } } else { - "=> " + // Normal prompt (PROMPT1) + if let Some(custom_prompt) = &context.prompt1 { + custom_prompt.as_str() + } else { + "=> " + } }; let readline = rl.readline(prompt); @@ -75,6 +92,15 @@ async fn main() -> Result<(), Box> { buffer += "\n"; if !line.is_empty() { + // Check if this is a meta-command (backslash command) + if line.trim().starts_with('\\') { + if let Err(e) = handle_meta_command(&mut context, line.trim()) { + eprintln!("Error processing meta-command: {}", e); + } + buffer.clear(); + continue; + } + let queries = try_split_queries(&buffer).unwrap_or_default(); if !queries.is_empty() { diff --git a/src/meta_commands.rs b/src/meta_commands.rs new file mode 100644 index 0000000..8c9a931 --- /dev/null +++ b/src/meta_commands.rs @@ -0,0 +1,340 @@ +use crate::context::Context; +use regex::Regex; +use once_cell::sync::Lazy; + +// Handle meta-commands (backslash commands) +pub fn handle_meta_command(context: &mut Context, command: &str) -> Result> { + // Handle \set PROMPT1 command + if let Some(prompt) = parse_set_prompt1(command) { + context.set_prompt1(prompt); + return Ok(true); + } + + // Handle \set PROMPT2 command + if let Some(prompt) = parse_set_prompt2(command) { + context.set_prompt2(prompt); + return Ok(true); + } + + // Handle \set PROMPT3 command + if let Some(prompt) = parse_set_prompt3(command) { + context.set_prompt3(prompt); + return Ok(true); + } + + // Handle \unset PROMPT1 command + if parse_unset_prompt1(command) { + context.prompt1 = None; + return Ok(true); + } + + // Handle \unset PROMPT2 command + if parse_unset_prompt2(command) { + context.prompt2 = None; + return Ok(true); + } + + // Handle \unset PROMPT3 command + if parse_unset_prompt3(command) { + context.prompt3 = None; + return Ok(true); + } + + Ok(false) +} + +// Parse \set PROMPT1 'value' command +fn parse_set_prompt1(command: &str) -> Option { + static SET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\set\s+PROMPT1\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() + }); + + if let Some(captures) = SET_PROMPT_RE.captures(command) { + // Check which capture group matched + if let Some(prompt) = captures.get(1) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(2) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(3) { + return Some(prompt.as_str().to_string()); + } + } + + None +} + +// Parse \set PROMPT2 'value' command +fn parse_set_prompt2(command: &str) -> Option { + static SET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\set\s+PROMPT2\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() + }); + + if let Some(captures) = SET_PROMPT_RE.captures(command) { + // Check which capture group matched + if let Some(prompt) = captures.get(1) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(2) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(3) { + return Some(prompt.as_str().to_string()); + } + } + + None +} + +// Parse \set PROMPT3 'value' command +fn parse_set_prompt3(command: &str) -> Option { + static SET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\set\s+PROMPT3\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() + }); + + if let Some(captures) = SET_PROMPT_RE.captures(command) { + // Check which capture group matched + if let Some(prompt) = captures.get(1) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(2) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(3) { + return Some(prompt.as_str().to_string()); + } + } + + None +} + +// Parse \unset PROMPT1 command +fn parse_unset_prompt1(command: &str) -> bool { + static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\unset\s+PROMPT1\s*$"#).unwrap() + }); + + UNSET_PROMPT_RE.is_match(command) +} + +// Parse \unset PROMPT2 command +fn parse_unset_prompt2(command: &str) -> bool { + static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\unset\s+PROMPT2\s*$"#).unwrap() + }); + + UNSET_PROMPT_RE.is_match(command) +} + +// Parse \unset PROMPT3 command +fn parse_unset_prompt3(command: &str) -> bool { + static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\unset\s+PROMPT3\s*$"#).unwrap() + }); + + UNSET_PROMPT_RE.is_match(command) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::args::get_args; + + #[test] + fn test_set_prompt1_single_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT1 'custom_prompt> '"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt1, Some("custom_prompt> ".to_string())); + } + + #[test] + fn test_set_prompt1_double_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT1 "custom_prompt> ""#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt1, Some("custom_prompt> ".to_string())); + } + + #[test] + fn test_set_prompt1_no_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT1 custom_prompt>"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt1, Some("custom_prompt>".to_string())); + } + + #[test] + fn test_set_prompt2_single_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT2 'custom_prompt> '"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt2, Some("custom_prompt> ".to_string())); + } + + #[test] + fn test_set_prompt2_double_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT2 "custom_prompt> ""#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt2, Some("custom_prompt> ".to_string())); + } + + #[test] + fn test_set_prompt2_no_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT2 custom_prompt>"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt2, Some("custom_prompt>".to_string())); + } + + #[test] + fn test_set_prompt3_single_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT3 'custom_prompt> '"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt3, Some("custom_prompt> ".to_string())); + } + + #[test] + fn test_set_prompt3_double_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT3 "custom_prompt> ""#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt3, Some("custom_prompt> ".to_string())); + } + + #[test] + fn test_set_prompt3_no_quotes() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + let command = r#"\set PROMPT3 custom_prompt>"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt3, Some("custom_prompt>".to_string())); + } + + #[test] + fn test_unset_prompt1() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // First set a prompt + context.set_prompt1("test> ".to_string()); + assert_eq!(context.prompt1, Some("test> ".to_string())); + + // Then unset it + let command = r#"\unset PROMPT1"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt1, None); + } + + #[test] + fn test_unset_prompt2() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // First set a prompt + context.set_prompt2("test> ".to_string()); + assert_eq!(context.prompt2, Some("test> ".to_string())); + + // Then unset it + let command = r#"\unset PROMPT2"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt2, None); + } + + #[test] + fn test_unset_prompt3() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // First set a prompt + context.set_prompt3("test> ".to_string()); + assert_eq!(context.prompt3, Some("test> ".to_string())); + + // Then unset it + let command = r#"\unset PROMPT3"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt3, None); + } + + #[test] + fn test_invalid_commands() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // Invalid commands should return false + let command = r#"\invalid command"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(!result); + + let command = r#"\set INVALID value"#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(!result); + } + + #[test] + fn test_whitespace_handling() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // Test with various whitespace + let command = r#" \set PROMPT1 'test>' "#; + let result = handle_meta_command(&mut context, command).unwrap(); + assert!(result); + assert_eq!(context.prompt1, Some("test>".to_string())); + } + + #[test] + fn test_prompt_independence() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // Set all three prompts to different values + let command1 = r#"\set PROMPT1 'prompt1> '"#; + let command2 = r#"\set PROMPT2 'prompt2> '"#; + let command3 = r#"\set PROMPT3 'prompt3> '"#; + + handle_meta_command(&mut context, command1).unwrap(); + handle_meta_command(&mut context, command2).unwrap(); + handle_meta_command(&mut context, command3).unwrap(); + + // Verify all prompts are set independently + assert_eq!(context.prompt1, Some("prompt1> ".to_string())); + assert_eq!(context.prompt2, Some("prompt2> ".to_string())); + assert_eq!(context.prompt3, Some("prompt3> ".to_string())); + + // Unset only PROMPT2 + let unset_command = r#"\unset PROMPT2"#; + handle_meta_command(&mut context, unset_command).unwrap(); + + // Verify only PROMPT2 was unset + assert_eq!(context.prompt1, Some("prompt1> ".to_string())); + assert_eq!(context.prompt2, None); + assert_eq!(context.prompt3, Some("prompt3> ".to_string())); + } +} From 5f5af62f115d57cf23269e677bab05f696426014 Mon Sep 17 00:00:00 2001 From: Mosha Pasumansky Date: Tue, 26 Aug 2025 17:10:11 -0700 Subject: [PATCH 2/2] Refactor parse prompt functions --- src/meta_commands.rs | 125 +++++++++++++++++-------------------------- 1 file changed, 49 insertions(+), 76 deletions(-) diff --git a/src/meta_commands.rs b/src/meta_commands.rs index 8c9a931..d4f604c 100644 --- a/src/meta_commands.rs +++ b/src/meta_commands.rs @@ -5,37 +5,37 @@ use once_cell::sync::Lazy; // Handle meta-commands (backslash commands) pub fn handle_meta_command(context: &mut Context, command: &str) -> Result> { // Handle \set PROMPT1 command - if let Some(prompt) = parse_set_prompt1(command) { + if let Some(prompt) = parse_set_prompt(command, "PROMPT1") { context.set_prompt1(prompt); return Ok(true); } // Handle \set PROMPT2 command - if let Some(prompt) = parse_set_prompt2(command) { + if let Some(prompt) = parse_set_prompt(command, "PROMPT2") { context.set_prompt2(prompt); return Ok(true); } // Handle \set PROMPT3 command - if let Some(prompt) = parse_set_prompt3(command) { + if let Some(prompt) = parse_set_prompt(command, "PROMPT3") { context.set_prompt3(prompt); return Ok(true); } // Handle \unset PROMPT1 command - if parse_unset_prompt1(command) { + if parse_unset_prompt(command, "PROMPT1") { context.prompt1 = None; return Ok(true); } // Handle \unset PROMPT2 command - if parse_unset_prompt2(command) { + if parse_unset_prompt(command, "PROMPT2") { context.prompt2 = None; return Ok(true); } // Handle \unset PROMPT3 command - if parse_unset_prompt3(command) { + if parse_unset_prompt(command, "PROMPT3") { context.prompt3 = None; return Ok(true); } @@ -43,91 +43,44 @@ pub fn handle_meta_command(context: &mut Context, command: &str) -> Result Option { +// Generic function to parse \set PROMPT command +fn parse_set_prompt(command: &str, prompt_type: &str) -> Option { static SET_PROMPT_RE: Lazy = Lazy::new(|| { - Regex::new(r#"(?i)^\s*\\set\s+PROMPT1\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() + Regex::new(r#"(?i)^\s*\\set\s+(\w+)\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() }); if let Some(captures) = SET_PROMPT_RE.captures(command) { - // Check which capture group matched - if let Some(prompt) = captures.get(1) { - return Some(prompt.as_str().to_string()); - } else if let Some(prompt) = captures.get(2) { - return Some(prompt.as_str().to_string()); - } else if let Some(prompt) = captures.get(3) { - return Some(prompt.as_str().to_string()); + // Check if the prompt type matches + if let Some(cmd_prompt_type) = captures.get(1) { + if cmd_prompt_type.as_str().eq_ignore_ascii_case(prompt_type) { + // Check which capture group matched for the value + if let Some(prompt) = captures.get(2) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(3) { + return Some(prompt.as_str().to_string()); + } else if let Some(prompt) = captures.get(4) { + return Some(prompt.as_str().to_string()); + } + } } } None } -// Parse \set PROMPT2 'value' command -fn parse_set_prompt2(command: &str) -> Option { - static SET_PROMPT_RE: Lazy = Lazy::new(|| { - Regex::new(r#"(?i)^\s*\\set\s+PROMPT2\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() - }); - - if let Some(captures) = SET_PROMPT_RE.captures(command) { - // Check which capture group matched - if let Some(prompt) = captures.get(1) { - return Some(prompt.as_str().to_string()); - } else if let Some(prompt) = captures.get(2) { - return Some(prompt.as_str().to_string()); - } else if let Some(prompt) = captures.get(3) { - return Some(prompt.as_str().to_string()); - } - } - - None -} - -// Parse \set PROMPT3 'value' command -fn parse_set_prompt3(command: &str) -> Option { - static SET_PROMPT_RE: Lazy = Lazy::new(|| { - Regex::new(r#"(?i)^\s*\\set\s+PROMPT3\s+(?:'([^']*)'|"([^"]*)"|(\S+))\s*$"#).unwrap() +// Generic function to parse \unset PROMPT command +fn parse_unset_prompt(command: &str, prompt_type: &str) -> bool { + static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { + Regex::new(r#"(?i)^\s*\\unset\s+(\w+)\s*$"#).unwrap() }); - if let Some(captures) = SET_PROMPT_RE.captures(command) { - // Check which capture group matched - if let Some(prompt) = captures.get(1) { - return Some(prompt.as_str().to_string()); - } else if let Some(prompt) = captures.get(2) { - return Some(prompt.as_str().to_string()); - } else if let Some(prompt) = captures.get(3) { - return Some(prompt.as_str().to_string()); + if let Some(captures) = UNSET_PROMPT_RE.captures(command) { + if let Some(cmd_prompt_type) = captures.get(1) { + return cmd_prompt_type.as_str().eq_ignore_ascii_case(prompt_type); } } - None -} - -// Parse \unset PROMPT1 command -fn parse_unset_prompt1(command: &str) -> bool { - static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { - Regex::new(r#"(?i)^\s*\\unset\s+PROMPT1\s*$"#).unwrap() - }); - - UNSET_PROMPT_RE.is_match(command) -} - -// Parse \unset PROMPT2 command -fn parse_unset_prompt2(command: &str) -> bool { - static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { - Regex::new(r#"(?i)^\s*\\unset\s+PROMPT2\s*$"#).unwrap() - }); - - UNSET_PROMPT_RE.is_match(command) -} - -// Parse \unset PROMPT3 command -fn parse_unset_prompt3(command: &str) -> bool { - static UNSET_PROMPT_RE: Lazy = Lazy::new(|| { - Regex::new(r#"(?i)^\s*\\unset\s+PROMPT3\s*$"#).unwrap() - }); - - UNSET_PROMPT_RE.is_match(command) + false } #[cfg(test)] @@ -337,4 +290,24 @@ mod tests { assert_eq!(context.prompt2, None); assert_eq!(context.prompt3, Some("prompt3> ".to_string())); } + + #[test] + fn test_case_insensitive_prompt_types() { + let args = get_args().unwrap(); + let mut context = Context::new(args); + + // Test case insensitive prompt type matching + let command1 = r#"\set prompt1 'test1> '"#; + let command2 = r#"\set Prompt2 'test2> '"#; + let command3 = r#"\set PROMPT3 'test3> '"#; + + handle_meta_command(&mut context, command1).unwrap(); + handle_meta_command(&mut context, command2).unwrap(); + handle_meta_command(&mut context, command3).unwrap(); + + // Verify all prompts are set correctly regardless of case + assert_eq!(context.prompt1, Some("test1> ".to_string())); + assert_eq!(context.prompt2, Some("test2> ".to_string())); + assert_eq!(context.prompt3, Some("test3> ".to_string())); + } }