Skip to content

Commit 57a8862

Browse files
arsenii-fireboltarsenii-firebolt
authored
report transaction status; exit via quit/exit (#2)
1 parent f7c2d85 commit 57a8862

File tree

3 files changed

+87
-38
lines changed

3 files changed

+87
-38
lines changed

‎src/main.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,12 +56,23 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
5656

5757
let mut buffer: String = String::new();
5858
loop {
59-
let prompt = if !buffer.trim_start().is_empty() { "~> " } else { "=> " };
59+
let prompt = if !buffer.trim_start().is_empty() {
60+
"~> "
61+
} else if context.args.extra.iter().any(|arg| arg.starts_with("transaction_id=")) {
62+
"*> "
63+
} else {
64+
"=> "
65+
};
6066
let readline = rl.readline(prompt);
6167

6268
match readline {
6369
Ok(line) => {
6470
buffer += line.as_str();
71+
72+
if buffer.trim() == "quit" || buffer.trim() == "exit" {
73+
break;
74+
}
75+
6576
buffer += "\n";
6677
if !line.is_empty() {
6778
let queries = try_split_queries(&buffer).unwrap_or_default();

‎src/query.rs‎

Lines changed: 45 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@ use std::time::Instant;
66
use tokio::{select, signal, task};
77
use tokio_util::sync::CancellationToken;
88

9-
use crate::USER_AGENT;
10-
use crate::FIREBOLT_PROTOCOL_VERSION;
119
use crate::args::normalize_extras;
1210
use crate::auth::authenticate_service_account;
1311
use crate::context::Context;
1412
use crate::utils::spin;
13+
use crate::FIREBOLT_PROTOCOL_VERSION;
14+
use crate::USER_AGENT;
1515

1616
// Set parameters via query
1717
pub fn set_args(context: &mut Context, query: &str) -> Result<bool, Box<dyn std::error::Error>> {
@@ -44,10 +44,6 @@ pub fn set_args(context: &mut Context, query: &str) -> Result<bool, Box<dyn std:
4444

4545
context.update_url();
4646

47-
if !context.args.concise && !context.args.hide_pii {
48-
eprintln!("URL: {}", context.url);
49-
}
50-
5147
return Ok(true);
5248
}
5349

@@ -67,10 +63,6 @@ pub fn unset_args(context: &mut Context, query: &str) -> Result<bool, Box<dyn st
6763

6864
context.update_url();
6965

70-
if !context.args.concise && !context.args.hide_pii {
71-
eprintln!("URL: {}", context.url);
72-
}
73-
7466
return Ok(true);
7567
}
7668

@@ -81,10 +73,18 @@ pub fn unset_args(context: &mut Context, query: &str) -> Result<bool, Box<dyn st
8173
pub async fn query(context: &mut Context, query_text: String) -> Result<(), Box<dyn std::error::Error>> {
8274
// Handle set/unset commands
8375
if set_args(context, &query_text)? {
76+
if !context.args.concise && !context.args.hide_pii {
77+
eprintln!("URL: {}", context.url);
78+
}
79+
8480
return Ok(());
8581
}
8682

8783
if unset_args(context, &query_text)? {
84+
if !context.args.concise && !context.args.hide_pii {
85+
eprintln!("URL: {}", context.url);
86+
}
87+
8888
return Ok(());
8989
}
9090

@@ -150,37 +150,45 @@ pub async fn query(context: &mut Context, query_text: String) -> Result<(), Box<
150150
let mut maybe_request_id: Option<String> = None;
151151
match response {
152152
Ok(resp) => {
153-
if let Some(header) = resp.headers().get("X-REQUEST-ID") {
154-
maybe_request_id = header.to_str().map_or(None, |l| Some(String::from(l)));
155-
}
156-
if let Some(header) = resp.headers().get("firebolt-update-parameters") {
157-
set_args(context, format!("set {}", header.to_str().unwrap()).as_str())?;
158-
}
159-
if let Some(header) = resp.headers().get("firebolt-remove-parameters") {
160-
unset_args(context, format!("unset {}", header.to_str().unwrap()).as_str())?;
161-
}
162-
if let Some(header) = resp.headers().get("firebolt-update-endpoint") {
163-
let header_str = header.to_str().unwrap();
164-
// Split the header at the '?' character
165-
if let Some(pos) = header_str.find('?') {
166-
// Extract base URL and query part
167-
let base_url = &header_str[..pos];
168-
let query_part = &header_str[pos+1..];
169-
170-
// Update the context URL with just the base part
171-
context.args.host = base_url.to_string();
172-
173-
// Process each query parameter
174-
for param in query_part.split('&') {
175-
if !param.is_empty() {
176-
set_args(context, format!("set {};", param).as_str())?;
153+
let mut updated_url = false;
154+
for (header, value) in resp.headers() {
155+
if header == "firebolt-remove-parameters" {
156+
unset_args(context, format!("unset {}", value.to_str()?).as_str())?;
157+
updated_url = true;
158+
} else if header == "firebolt-update-parameters" {
159+
set_args(context, format!("set {}", value.to_str()?).as_str())?;
160+
updated_url = true;
161+
} else if header == "X-REQUEST-ID" {
162+
maybe_request_id = value.to_str().map_or(None, |l| Some(String::from(l)));
163+
updated_url = true;
164+
} else if header == "firebolt-update-endpoint" {
165+
let header_str = value.to_str()?;
166+
// Split the header at the '?' character
167+
if let Some(pos) = header_str.find('?') {
168+
// Extract base URL and query part
169+
let base_url = &header_str[..pos];
170+
let query_part = &header_str[pos+1..];
171+
172+
// Update the context URL with just the base part
173+
context.args.host = base_url.to_string();
174+
175+
// Process each query parameter
176+
for param in query_part.split('&') {
177+
if !param.is_empty() {
178+
set_args(context, format!("set {};", param).as_str())?;
179+
}
177180
}
181+
} else {
182+
// No query parameters, just set the URL
183+
context.args.host = header_str.to_string();
178184
}
179-
} else {
180-
// No query parameters, just set the URL
181-
context.args.host = header_str.to_string();
185+
updated_url = true;
182186
}
183187
}
188+
if updated_url && !context.args.concise && !context.args.hide_pii {
189+
eprintln!("URL: {}", context.url);
190+
}
191+
184192
// on stdout, on purpose
185193
println!("{}", resp.text().await?);
186194
}

‎tests/cli.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,3 +175,33 @@ fn test_command_parsing() {
175175

176176
assert!(stdout.contains("1339"));
177177
}
178+
179+
#[test]
180+
fn test_exiting() {
181+
let mut child = Command::new(env!("CARGO_BIN_EXE_fb"))
182+
.args(&[
183+
"--core",
184+
"--concise",
185+
"-f",
186+
"TabSeparatedWithNamesAndTypes",
187+
])
188+
.stdin(std::process::Stdio::piped())
189+
.stdout(std::process::Stdio::piped())
190+
.spawn()
191+
.unwrap();
192+
193+
let mut stdin = child.stdin.take().unwrap();
194+
writeln!(stdin, "SELECT 42;").unwrap();
195+
writeln!(stdin, "quit").unwrap();
196+
drop(stdin); // Close stdin to end interactive mode
197+
198+
let output = child.wait_with_output().unwrap();
199+
let stdout = String::from_utf8(output.stdout).unwrap();
200+
201+
assert!(output.status.success());
202+
let mut lines = stdout.lines();
203+
assert_eq!(lines.next().unwrap(), "?column?");
204+
lines.next();
205+
assert_eq!(lines.next().unwrap(), "42");
206+
lines.next();
207+
}

0 commit comments

Comments
 (0)