From 6c8f8c2fe3ecda51fc7f805967dbfd6c1b70d979 Mon Sep 17 00:00:00 2001 From: arsenii Date: Tue, 1 Jul 2025 16:59:44 -0700 Subject: [PATCH] report transaction status; exit via quit/exit --- src/main.rs | 13 ++++++++- src/query.rs | 82 ++++++++++++++++++++++++++++------------------------ tests/cli.rs | 30 +++++++++++++++++++ 3 files changed, 87 insertions(+), 38 deletions(-) diff --git a/src/main.rs b/src/main.rs index 8ac1e54..f0ee579 100644 --- a/src/main.rs +++ b/src/main.rs @@ -56,12 +56,23 @@ async fn main() -> Result<(), Box> { let mut buffer: String = String::new(); loop { - let prompt = if !buffer.trim_start().is_empty() { "~> " } else { "=> " }; + let prompt = if !buffer.trim_start().is_empty() { + "~> " + } else if context.args.extra.iter().any(|arg| arg.starts_with("transaction_id=")) { + "*> " + } else { + "=> " + }; let readline = rl.readline(prompt); match readline { Ok(line) => { buffer += line.as_str(); + + if buffer.trim() == "quit" || buffer.trim() == "exit" { + break; + } + buffer += "\n"; if !line.is_empty() { let queries = try_split_queries(&buffer).unwrap_or_default(); diff --git a/src/query.rs b/src/query.rs index b91928a..19691d9 100644 --- a/src/query.rs +++ b/src/query.rs @@ -6,12 +6,12 @@ use std::time::Instant; use tokio::{select, signal, task}; use tokio_util::sync::CancellationToken; -use crate::USER_AGENT; -use crate::FIREBOLT_PROTOCOL_VERSION; use crate::args::normalize_extras; use crate::auth::authenticate_service_account; use crate::context::Context; use crate::utils::spin; +use crate::FIREBOLT_PROTOCOL_VERSION; +use crate::USER_AGENT; // Set parameters via query pub fn set_args(context: &mut Context, query: &str) -> Result> { @@ -44,10 +44,6 @@ pub fn set_args(context: &mut Context, query: &str) -> Result Result Result Result<(), Box> { // Handle set/unset commands if set_args(context, &query_text)? { + if !context.args.concise && !context.args.hide_pii { + eprintln!("URL: {}", context.url); + } + return Ok(()); } if unset_args(context, &query_text)? { + if !context.args.concise && !context.args.hide_pii { + eprintln!("URL: {}", context.url); + } + return Ok(()); } @@ -150,37 +150,45 @@ pub async fn query(context: &mut Context, query_text: String) -> Result<(), Box< let mut maybe_request_id: Option = None; match response { Ok(resp) => { - if let Some(header) = resp.headers().get("X-REQUEST-ID") { - maybe_request_id = header.to_str().map_or(None, |l| Some(String::from(l))); - } - if let Some(header) = resp.headers().get("firebolt-update-parameters") { - set_args(context, format!("set {}", header.to_str().unwrap()).as_str())?; - } - if let Some(header) = resp.headers().get("firebolt-remove-parameters") { - unset_args(context, format!("unset {}", header.to_str().unwrap()).as_str())?; - } - if let Some(header) = resp.headers().get("firebolt-update-endpoint") { - let header_str = header.to_str().unwrap(); - // Split the header at the '?' character - if let Some(pos) = header_str.find('?') { - // Extract base URL and query part - let base_url = &header_str[..pos]; - let query_part = &header_str[pos+1..]; - - // Update the context URL with just the base part - context.args.host = base_url.to_string(); - - // Process each query parameter - for param in query_part.split('&') { - if !param.is_empty() { - set_args(context, format!("set {};", param).as_str())?; + let mut updated_url = false; + for (header, value) in resp.headers() { + if header == "firebolt-remove-parameters" { + unset_args(context, format!("unset {}", value.to_str()?).as_str())?; + updated_url = true; + } else if header == "firebolt-update-parameters" { + set_args(context, format!("set {}", value.to_str()?).as_str())?; + updated_url = true; + } else if header == "X-REQUEST-ID" { + maybe_request_id = value.to_str().map_or(None, |l| Some(String::from(l))); + updated_url = true; + } else if header == "firebolt-update-endpoint" { + let header_str = value.to_str()?; + // Split the header at the '?' character + if let Some(pos) = header_str.find('?') { + // Extract base URL and query part + let base_url = &header_str[..pos]; + let query_part = &header_str[pos+1..]; + + // Update the context URL with just the base part + context.args.host = base_url.to_string(); + + // Process each query parameter + for param in query_part.split('&') { + if !param.is_empty() { + set_args(context, format!("set {};", param).as_str())?; + } } + } else { + // No query parameters, just set the URL + context.args.host = header_str.to_string(); } - } else { - // No query parameters, just set the URL - context.args.host = header_str.to_string(); + updated_url = true; } } + if updated_url && !context.args.concise && !context.args.hide_pii { + eprintln!("URL: {}", context.url); + } + // on stdout, on purpose println!("{}", resp.text().await?); } diff --git a/tests/cli.rs b/tests/cli.rs index f93a1f9..07a366f 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -175,3 +175,33 @@ fn test_command_parsing() { assert!(stdout.contains("1339")); } + +#[test] +fn test_exiting() { + let mut child = Command::new(env!("CARGO_BIN_EXE_fb")) + .args(&[ + "--core", + "--concise", + "-f", + "TabSeparatedWithNamesAndTypes", + ]) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .spawn() + .unwrap(); + + let mut stdin = child.stdin.take().unwrap(); + writeln!(stdin, "SELECT 42;").unwrap(); + writeln!(stdin, "quit").unwrap(); + drop(stdin); // Close stdin to end interactive mode + + let output = child.wait_with_output().unwrap(); + let stdout = String::from_utf8(output.stdout).unwrap(); + + assert!(output.status.success()); + let mut lines = stdout.lines(); + assert_eq!(lines.next().unwrap(), "?column?"); + lines.next(); + assert_eq!(lines.next().unwrap(), "42"); + lines.next(); +}