diff --git a/README.md b/README.md index 6bfa99d..dc8290d 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,8 @@ hotdata query status [-o table|json|csv] - Long-running queries automatically fall back to async execution and return a `query_run_id`. - Use `hotdata query status ` to poll for results. - Exit codes for `query status`: `0` = succeeded, `1` = failed, `2` = still running (poll again). +- `json` output carries `truncated`, `preview_row_count`, and `total_row_count` so a consumer can detect a partial result from the body alone. +- If the server returns only a bounded preview that can't be completed (truncated and unfetchable), the CLI prints the preview, warns on stderr, and exits `3` — so pipelines break rather than silently ingest partial data. ## Query Run History diff --git a/src/query.rs b/src/query.rs index 72fe43e..ebdf2cb 100644 --- a/src/query.rs +++ b/src/query.rs @@ -1,4 +1,4 @@ -use crate::sdk::Api; +use crate::sdk::{Api, ApiError}; use serde::Deserialize; use serde_json::Value; @@ -7,11 +7,29 @@ pub struct QueryResponse { pub result_id: Option, pub columns: Vec, pub rows: Vec>, + /// Rows actually carried here (`rows.len()`). For a complete result this is + /// the whole result; for an incomplete preview (`truncated`) it's a bounded + /// subset. pub row_count: u64, + /// Grand total rows in the full result when the server reported it. `None` + /// when unknown — e.g. a truncated result whose persistence never started. + pub total_row_count: Option, + /// True when `rows` is an *incomplete* subset the CLI could not complete: + /// the server truncated the result and it was either never persisted (no + /// `result_id`) or the follow-up fetch failed. A truncated result the CLI + /// successfully follows to the full set is `false` — the rows held are the + /// whole result. Drives the fail-closed exit in [`print_result`]. + pub truncated: bool, pub execution_time_ms: Option, pub warning: Option, } +/// Exit code emitted when the CLI prints an incomplete preview (a truncated +/// result it could not complete). Distinct from the generic `1`/`2` so pipelines +/// can tell "partial data" apart from a hard failure and break rather than +/// silently ingest a subset. +pub const EXIT_INCOMPLETE_RESULT: i32 = 3; + /// Convert the SDK's inline `QueryResponse` (200 path) into the CLI's display /// model. The async path decodes Arrow instead (see `fetch_arrow_result`). /// @@ -22,11 +40,32 @@ pub struct QueryResponse { /// warning). A truncated result that *can* be fetched never reaches here — it's /// followed to the full set via Arrow in `resolve_inline`. So counting the held /// rows can never overstate or understate what the user sees. +/// +/// `truncated` and `total_row_count` are carried through verbatim so structured +/// output (`--output json`) exposes whether the rows are a partial preview and, +/// when known, the grand total. For a complete (non-truncated) result the server +/// reports `total_row_count == preview`; if it's absent we backfill it from the +/// held rows, which for a complete result is the whole truth. fn query_response_from_sdk(resp: hotdata::models::QueryResponse) -> QueryResponse { + let row_count = resp.rows.len() as u64; + // A negative total is impossible per the API contract; treat one as absent + // rather than clamping to 0, which would emit a contradictory + // `preview_row_count > total_row_count`. + let total_row_count = resp + .total_row_count + .flatten() + .and_then(|t| u64::try_from(t).ok()) + .or(if resp.truncated { + None + } else { + Some(row_count) + }); QueryResponse { result_id: resp.result_id.flatten(), columns: resp.columns, - row_count: resp.rows.len() as u64, + row_count, + total_row_count, + truncated: resp.truncated, rows: resp.rows, execution_time_ms: Some(resp.execution_time_ms.max(0) as u64), warning: resp.warning.flatten(), @@ -193,24 +232,46 @@ fn arrow_result_to_query_response( } let row_count = rows.len() as u64; + // The fetched Arrow result is the full persisted result, so it's complete: + // `truncated` is false and the total is the authoritative `X-Total-Row-Count` + // the SDK parsed (falling back to the rows we hold). + let total_row_count = result + .total_row_count + .and_then(|t| u64::try_from(t).ok()) + .or(Some(row_count)); QueryResponse { result_id: Some(result_id), columns, rows, row_count, + total_row_count, + truncated: false, execution_time_ms: None, warning: None, } } -/// Fetch `/results/{result_id}` as Arrow and return a `QueryResponse`. +/// Fetch `/results/{result_id}` as Arrow and return a `QueryResponse`, returning +/// the error instead of exiting. /// /// Both transport and decode are owned by the SDK's `get_result_arrow` (via the /// [`Api::get_result_arrow`] seam), so the CLI shares one `arrow` major version /// with the SDK. +/// +/// The fallible form lets callers that hold a fallback (e.g. an inline preview) +/// degrade gracefully rather than terminate the process; [`fetch_arrow_result`] +/// is the exiting wrapper for callers with nothing else to show. +pub(crate) fn try_fetch_arrow_result( + api: &Api, + result_id: &str, +) -> Result { + let result = api.get_result_arrow(result_id)?; + Ok(arrow_result_to_query_response(result, result_id.to_owned())) +} + +/// Fetch `/results/{result_id}` as Arrow, exiting the process on failure. pub(crate) fn fetch_arrow_result(api: &Api, result_id: &str) -> QueryResponse { - let result = api.get_result_arrow(result_id).unwrap_or_else(|e| e.exit()); - arrow_result_to_query_response(result, result_id.to_owned()) + try_fetch_arrow_result(api, result_id).unwrap_or_else(|e| e.exit()) } /// Resolve an inline (HTTP 200) query response for display. @@ -224,34 +285,50 @@ pub(crate) fn fetch_arrow_result(api: &Api, result_id: &str) -> QueryResponse { /// print only the preview rows. /// /// If a truncated response has no `result_id` (persistence could not be -/// initiated — see the SDK's `warning` field), the full result is unfetchable, -/// so fall back to the preview and surface a warning rather than failing. +/// initiated — see the SDK's `warning` field), or the follow-up Arrow fetch +/// fails, the full result is unreachable. Rather than exiting and discarding the +/// preview the inline body already carries, fall back to that preview, mark it +/// incomplete (`truncated`), and surface a warning. `print_result` then fails +/// closed (non-zero exit) so the partial data can't be silently consumed. fn resolve_inline(api: &Api, resp: hotdata::models::QueryResponse) -> QueryResponse { if !resp.truncated { return query_response_from_sdk(resp); } match resp.result_id.clone().flatten() { - Some(result_id) => { + Some(result_id) => match try_fetch_arrow_result(api, &result_id) { // The Arrow fetch returns only schema + rows; carry the query-level // warning and execution time the inline response reported, which // `arrow_result_to_query_response` otherwise hardcodes to None. - let mut full = fetch_arrow_result(api, &result_id); - full.warning = resp.warning.flatten(); - full.execution_time_ms = Some(resp.execution_time_ms.max(0) as u64); - full - } - None => { - let mut preview = query_response_from_sdk(resp); - let note = "result truncated to a preview; full result unavailable (persistence not initiated)"; - preview.warning = Some(match preview.warning { - Some(w) => format!("{w}; {note}"), - None => note.to_string(), - }); - preview - } + Ok(mut full) => { + full.warning = resp.warning.flatten(); + full.execution_time_ms = Some(resp.execution_time_ms.max(0) as u64); + full + } + // The full result is persisted but the follow-up fetch failed (e.g. + // transport error, persistence still in progress). Degrade to the + // preview instead of hard-exiting and losing the rows in hand. + Err(e) => incomplete_preview( + resp, + &format!("could not fetch full result ({})", e.message()), + ), + }, + None => incomplete_preview(resp, "full result unavailable (persistence not initiated)"), } } +/// Keep a truncated inline response's preview rows, mark it incomplete, and fold +/// `note` into the warning (preserving any SDK-provided warning that explains +/// *why* the full result is unreachable). +fn incomplete_preview(resp: hotdata::models::QueryResponse, note: &str) -> QueryResponse { + let mut preview = query_response_from_sdk(resp); + let note = format!("result truncated to a preview; {note}"); + preview.warning = Some(match preview.warning { + Some(w) => format!("{w}; {note}"), + None => note, + }); + preview +} + pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &str) { let api = Api::new(Some(workspace_id)); @@ -386,6 +463,71 @@ pub fn poll(query_run_id: &str, workspace_id: &str, format: &str) { } } +/// Build the `--output json` body for a result. +/// +/// `row_count` (rows in this body) is kept for backward compatibility, but the +/// truncation truth is carried explicitly: `truncated`, `preview_row_count` +/// (rows held), and `total_row_count` (grand total, `null` when unknown). A +/// JSON consumer can detect a partial result from these rather than having to +/// notice a stderr-only warning. +fn result_json(result: &QueryResponse) -> Value { + serde_json::json!({ + "result_id": result.result_id, + "columns": result.columns, + "rows": result.rows, + "row_count": result.row_count, + "preview_row_count": result.row_count, + "total_row_count": result.total_row_count, + "truncated": result.truncated, + "execution_time_ms": result.execution_time_ms, + }) +} + +/// Process exit code after rendering a result: [`EXIT_INCOMPLETE_RESULT`] when +/// the rows are an incomplete preview (fail closed so pipelines break), else `0`. +fn result_exit_code(result: &QueryResponse) -> i32 { + if result.truncated { + EXIT_INCOMPLETE_RESULT + } else { + 0 + } +} + +/// The unstyled summary line printed under a `table` result. +/// +/// A complete result reads `N rows (time) [result-id]`. An incomplete preview is +/// loud — `N of TOTAL rows — INCOMPLETE PREVIEW (...)` — with `?` standing in for +/// a total the server didn't report. The caller colours it (red vs grey). +fn table_footer(result: &QueryResponse) -> String { + let id_part = result + .result_id + .as_deref() + .map(|id| format!(" [result-id: {id}]")) + .unwrap_or_default(); + let time_part = match result.execution_time_ms { + Some(ms) => format!("{ms} ms"), + None => "\u{2014}".to_string(), // em dash + }; + if result.truncated { + let total = result + .total_row_count + .map(|t| t.to_string()) + .unwrap_or_else(|| "?".to_string()); + format!( + "{} of {} rows — INCOMPLETE PREVIEW ({}){}", + result.row_count, total, time_part, id_part + ) + } else { + format!( + "{} row{} ({}){}", + result.row_count, + if result.row_count == 1 { "" } else { "s" }, + time_part, + id_part + ) + } +} + pub fn print_result(result: &QueryResponse, format: &str) { if let Some(ref warning) = result.warning { eprintln!("warning: {warning}"); @@ -393,14 +535,10 @@ pub fn print_result(result: &QueryResponse, format: &str) { match format { "json" => { - let out = serde_json::json!({ - "result_id": result.result_id, - "columns": result.columns, - "rows": result.rows, - "row_count": result.row_count, - "execution_time_ms": result.execution_time_ms, - }); - println!("{}", serde_json::to_string_pretty(&out).unwrap()); + println!( + "{}", + serde_json::to_string_pretty(&result_json(result)).unwrap() + ); } "csv" => { println!("{}", result.columns.join(",")); @@ -422,29 +560,25 @@ pub fn print_result(result: &QueryResponse, format: &str) { "table" => { crate::table::print_json(&result.columns, &result.rows); use crossterm::style::Stylize; - let id_part = result - .result_id - .as_deref() - .map(|id| format!(" [result-id: {id}]")) - .unwrap_or_default(); - let time_part = match result.execution_time_ms { - Some(ms) => format!("{ms} ms"), - None => "\u{2014}".to_string(), // em dash - }; - eprintln!( - "{}", - format!( - "\n{} row{} ({}){}", - result.row_count, - if result.row_count == 1 { "" } else { "s" }, - time_part, - id_part - ) - .dark_grey() - ); + let footer = table_footer(result); + // Loud (red) when the preview is incomplete so it can't be mistaken + // for a complete result; quiet (grey) otherwise. + if result.truncated { + eprintln!("\n{}", footer.red()); + } else { + eprintln!("\n{}", footer.dark_grey()); + } } _ => unreachable!(), } + + // Fail closed: an incomplete preview was just printed (with a stderr + // warning). Exit non-zero so a pipeline consuming the output breaks rather + // than silently ingesting a subset of the result as if it were complete. + let code = result_exit_code(result); + if code != 0 { + std::process::exit(code); + } } #[cfg(test)] @@ -516,12 +650,51 @@ mod tests { assert_eq!(resolved.row_count, 3); assert_eq!(resolved.rows.len(), 3); assert_eq!(resolved.result_id.as_deref(), Some("res_1")); + // The held rows are now the whole result: complete, not a preview. + assert!(!resolved.truncated); + assert_eq!(resolved.total_row_count, Some(3)); // Inline warning + timing carried through, not dropped by the fetch. assert_eq!(resolved.warning.as_deref(), Some("approximate aggregate")); assert_eq!(resolved.execution_time_ms, Some(5)); m.assert(); } + #[test] + fn resolve_inline_falls_back_to_preview_when_follow_fetch_fails() { + // Truncated with a result_id, but the follow-up Arrow fetch fails (500). + // The CLI must NOT hard-exit (which would also discard the preview it + // already holds) — it degrades to the preview, marks it incomplete, and + // explains. A returning call here is itself the assertion that no exit + // happened. + let mut server = mockito::Server::new(); + let m = server + .mock("GET", "/v1/results/res_1") + .match_query(mockito::Matcher::UrlEncoded( + "format".into(), + "arrow".into(), + )) + .with_status(500) + .with_body("boom") + .create(); + + let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let resolved = resolve_inline(&api, truncated_preview(Some("res_1"))); + + // Preview kept, flagged incomplete so print_result fails closed. + assert!(resolved.truncated); + assert_eq!(resolved.row_count, 1); + assert_eq!(resolved.rows.len(), 1); + // The fixture reports no total, so the table renders "1 of ? rows". + assert_eq!(resolved.total_row_count, None); + let warning = resolved.warning.as_deref().unwrap_or(""); + assert!(warning.contains("truncated"), "warning: {warning:?}"); + assert!( + warning.contains("could not fetch full result"), + "warning: {warning:?}" + ); + m.assert(); + } + #[test] fn resolve_inline_returns_untruncated_preview_without_fetching() { // truncated=false short-circuits before any network call; point the Api @@ -548,6 +721,9 @@ mod tests { vec![vec![serde_json::json!(1)], vec![serde_json::json!(2)]] ); assert_eq!(resolved.result_id.as_deref(), Some("res_2")); + // Complete result: not a preview, total backfilled from held rows. + assert!(!resolved.truncated); + assert_eq!(resolved.total_row_count, Some(2)); } #[test] @@ -557,9 +733,16 @@ mod tests { let server = mockito::Server::new(); let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); - let resolved = resolve_inline(&api, truncated_preview(None)); + // The server reports the grand total even though it couldn't persist; + // it must survive onto the preview so structured output exposes it. + let mut resp = truncated_preview(None); + resp.total_row_count = Some(Some(100)); + + let resolved = resolve_inline(&api, resp); + assert!(resolved.truncated); assert_eq!(resolved.row_count, 1); assert_eq!(resolved.rows.len(), 1); + assert_eq!(resolved.total_row_count, Some(100)); assert!( resolved .warning @@ -569,6 +752,110 @@ mod tests { ); } + #[test] + fn result_json_exposes_truncation_for_incomplete_preview() { + // A JSON consumer must be able to detect a partial result from the body + // alone — not only from a stderr warning. truncated=true, preview < + // total, and both counts present. + let result = QueryResponse { + result_id: None, + columns: vec!["id".to_string()], + rows: vec![vec![serde_json::json!(1)]], + row_count: 1, + total_row_count: Some(100), + truncated: true, + execution_time_ms: Some(5), + warning: Some("result truncated to a preview".to_string()), + }; + let json = result_json(&result); + assert_eq!(json["truncated"], serde_json::json!(true)); + assert_eq!(json["preview_row_count"], serde_json::json!(1)); + assert_eq!(json["row_count"], serde_json::json!(1)); + assert_eq!(json["total_row_count"], serde_json::json!(100)); + } + + #[test] + fn result_json_marks_complete_result_not_truncated() { + let result = QueryResponse { + result_id: Some("res_9".to_string()), + columns: vec!["id".to_string()], + rows: vec![vec![serde_json::json!(1)], vec![serde_json::json!(2)]], + row_count: 2, + total_row_count: Some(2), + truncated: false, + execution_time_ms: Some(5), + warning: None, + }; + let json = result_json(&result); + assert_eq!(json["truncated"], serde_json::json!(false)); + assert_eq!(json["preview_row_count"], serde_json::json!(2)); + assert_eq!(json["total_row_count"], serde_json::json!(2)); + // A complete result exits 0 — pipelines proceed. + assert_eq!(result_exit_code(&result), 0); + } + + #[test] + fn incomplete_preview_fails_closed_with_distinct_exit_code() { + // The fail-closed contract: an incomplete preview maps to a non-zero, + // non-generic exit code so a pipeline breaks instead of ingesting a + // partial result that exited 0. + let result = QueryResponse { + result_id: None, + columns: vec!["id".to_string()], + rows: vec![vec![serde_json::json!(1)]], + row_count: 1, + total_row_count: Some(100), + truncated: true, + execution_time_ms: Some(5), + warning: Some("result truncated to a preview".to_string()), + }; + assert_eq!(result_exit_code(&result), EXIT_INCOMPLETE_RESULT); + assert_ne!(EXIT_INCOMPLETE_RESULT, 0); + // Distinct from the generic failure codes used elsewhere in this module. + assert_ne!(EXIT_INCOMPLETE_RESULT, 1); + assert_ne!(EXIT_INCOMPLETE_RESULT, 2); + } + + /// Build a minimal display `QueryResponse` for footer rendering tests. + fn display_result( + row_count: u64, + total_row_count: Option, + truncated: bool, + ) -> QueryResponse { + QueryResponse { + result_id: None, + columns: vec!["id".to_string()], + rows: Vec::new(), + row_count, + total_row_count, + truncated, + execution_time_ms: Some(5), + warning: None, + } + } + + #[test] + fn table_footer_marks_incomplete_preview_loudly() { + let footer = table_footer(&display_result(1, Some(100), true)); + assert!(footer.contains("INCOMPLETE PREVIEW"), "footer: {footer}"); + assert!(footer.contains("1 of 100 rows"), "footer: {footer}"); + } + + #[test] + fn table_footer_renders_question_mark_for_unknown_total() { + // Truncated with no server-reported total → "1 of ? rows". + let footer = table_footer(&display_result(1, None, true)); + assert!(footer.contains("1 of ? rows"), "footer: {footer}"); + assert!(footer.contains("INCOMPLETE PREVIEW"), "footer: {footer}"); + } + + #[test] + fn table_footer_for_complete_result_is_plain() { + let footer = table_footer(&display_result(2, Some(2), false)); + assert!(!footer.contains("INCOMPLETE"), "footer: {footer}"); + assert!(footer.starts_with("2 rows"), "footer: {footer}"); + } + #[test] fn resolve_inline_preserves_existing_warning_when_following_fails() { // A truncated response with no result_id often arrives with an SDK diff --git a/src/sdk.rs b/src/sdk.rs index a2404d0..43b10c3 100644 --- a/src/sdk.rs +++ b/src/sdk.rs @@ -210,6 +210,17 @@ impl ApiError { } } + /// A printable, single-line description of the failure. + /// + /// Used where the error is surfaced inline (e.g. folded into a query + /// `warning`) rather than printed-and-exited via [`exit`](Self::exit). + pub fn message(&self) -> String { + match self { + ApiError::Status { status, body } => format!("{status}: {body}"), + ApiError::Transport(msg) => msg.clone(), + } + } + /// Print the standard error and exit, reproducing `ApiClient::fail_response`. /// /// On a 4xx, re-probe the auth status so a masked 404/403 is upgraded into @@ -673,6 +684,20 @@ mod tests { use super::*; use auth::AuthStatus; + #[test] + fn api_error_message_formats_status_and_transport() { + let status = ApiError::Status { + status: reqwest::StatusCode::INTERNAL_SERVER_ERROR, + body: "boom".to_string(), + }; + let m = status.message(); + assert!(m.contains("500"), "{m}"); + assert!(m.contains("boom"), "{m}"); + + let transport = ApiError::Transport("error connecting to API: refused".to_string()); + assert_eq!(transport.message(), "error connecting to API: refused"); + } + // --- format_fail_message: ported verbatim from api.rs (9 cases) ---------- #[test]