From ea0cd6b85f5c325de60bce56a63ccde51f27f01c Mon Sep 17 00:00:00 2001 From: Zac Farrell Date: Tue, 7 Jul 2026 09:35:21 -0700 Subject: [PATCH 1/2] feat(queries,results): scope to database for SDK 0.8.0 --- Cargo.lock | 4 +-- Cargo.toml | 2 +- src/cli.rs | 8 ++++++ src/client/sdk.rs | 54 +++++++++++++++++++++++++++++++++++------ src/commands/queries.rs | 24 ++++++++++++++---- src/commands/query.rs | 48 ++++++++++++++++++++++++------------ src/commands/results.rs | 29 +++++++++++++++++++--- src/main.rs | 13 +++++++--- 8 files changed, 142 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a6a0a56..4145627 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1165,9 +1165,9 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" [[package]] name = "hotdata" -version = "0.7.0" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d820a420f191d02017b94c3814e8305cab0442f47f785732ae4df0a240c1754" +checksum = "2e63f030edd9865312928059d492b3caf3c64f44b4a26abbfc92779517ec9026" dependencies = [ "arrow-array", "arrow-ipc", diff --git a/Cargo.toml b/Cargo.toml index ec5d668..4f74a6f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,7 @@ path = "src/main.rs" # behind a shared multi-thread tokio runtime. The `arrow` feature backs result # decode; `upload_file` runs the presigned direct-to-storage upload flow # (see src/sdk.rs::Api::upload). -hotdata = { version = "0.7.0", features = ["arrow"] } +hotdata = { version = "0.8.0", features = ["arrow"] } # Shared multi-thread runtime for the sync wrapper; block_on is called # concurrently from rayon worker threads (see src/indexes.rs). The presigned # upload (src/sdk.rs::Api::upload) drives the SDK's async `upload_file` on this diff --git a/src/cli.rs b/src/cli.rs index 41827ef..896b4a5 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -103,6 +103,10 @@ pub enum Commands { #[arg(long, short = 'w', global = true)] workspace_id: Option, + /// Managed database to scope to (defaults to the current database set via `databases set`) + #[arg(long, short = 'd', global = true)] + database: Option, + /// Output format #[arg(long = "output", short = 'o', default_value = "table", value_parser = ["table", "json", "csv"])] output: String, @@ -199,6 +203,10 @@ pub enum Commands { /// Query run ID to show details id: Option, + /// Managed database to scope to (defaults to the current database set via `databases set`) + #[arg(long, short = 'd', global = true)] + database: Option, + /// Output format (used with query run ID) #[arg(long = "output", short = 'o', default_value = "table", value_parser = ["table", "json", "yaml"])] output: String, diff --git a/src/client/sdk.rs b/src/client/sdk.rs index 69e88a1..d8273be 100644 --- a/src/client/sdk.rs +++ b/src/client/sdk.rs @@ -609,6 +609,39 @@ impl Api { self.database_id.as_deref() } + /// Return a clone of this `Api` scoped to `database_id`, overriding any + /// database resolved at construction (`HOTDATA_DATABASE` / current + /// database). Used by commands whose `--database` flag scopes a single + /// invocation without switching the persisted active database. + pub fn scoped_to_database(&self, database_id: String) -> Self { + Self { + database_id: Some(database_id), + ..self.clone() + } + } + + /// The active database scope, or exit with actionable guidance when none is + /// set. + /// + /// Results and query runs are scoped to a managed database (the required + /// `X-Database-Id` header). Commands that hit those endpoints call this to + /// resolve the scope — set via `--database`, `HOTDATA_DATABASE`, or + /// `databases set` — and to fail with a hint rather than surfacing the raw + /// server "X-Database-Id header is required" error. + pub fn require_database(&self) -> &str { + self.database_id.as_deref().unwrap_or_else(|| { + use crossterm::style::Stylize; + eprintln!("{}", "error: no active database.".red()); + eprintln!( + "{}", + "Results and query runs are scoped to a managed database. Set one with \ + `hotdata databases set `, or pass `--database `." + .dark_grey() + ); + std::process::exit(1); + }) + } + /// Best-effort probe of the scoped workspace's runtimedb scale state, via /// the control-plane `GET /v1/workspaces/{id}/runtime/status` endpoint. /// @@ -809,12 +842,16 @@ impl Api { /// `get_result_arrow`, returning the fully-buffered [`hotdata::ArrowResult`]. /// /// The SDK owns transport (same reqwest client, bearer via the - /// `token_provider`, `X-Workspace-Id`) and decode. Its - /// `ArrowError` (the Arrow-path error type, which is not an `Error`) is - /// mapped to [`ApiError`] via [`from_arrow`](ApiError::from_arrow) so callers - /// keep the same `.exit()` handling. + /// `token_provider`, `X-Workspace-Id`) and decode. Results are + /// database-scoped, so the active database is forwarded as the required + /// `X-Database-Id`; [`require_database`](Self::require_database) exits with a + /// hint when none is set. Its `ArrowError` (the Arrow-path error type, which + /// is not an `Error`) is mapped to [`ApiError`] via + /// [`from_arrow`](ApiError::from_arrow) so callers keep the same `.exit()` + /// handling. pub fn get_result_arrow(&self, id: &str) -> Result { - rt().block_on(self.client.get_result_arrow(id, None, None)) + let database_id = self.require_database(); + rt().block_on(self.client.get_result_arrow(id, database_id, None, None)) .map_err(ApiError::from_arrow) } @@ -1303,13 +1340,14 @@ mod tests { )) .match_header("Authorization", "Bearer test-jwt") .match_header("X-Workspace-Id", "ws-1") + .match_header("X-Database-Id", "db-1") .match_header("Accept", "application/vnd.apache.arrow.stream") .with_status(200) .with_header("content-type", "application/vnd.apache.arrow.stream") .with_body(ipc) .create(); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); let result = api .get_result_arrow("res_1") .expect("get_result_arrow should succeed"); @@ -1332,7 +1370,7 @@ mod tests { .with_body("not found") .create(); - let api = Api::test_new(&server.url(), "test-jwt", None); + let api = Api::test_new_scoped(&server.url(), "test-jwt", None, Some("db-1")); match api.get_result_arrow("missing").unwrap_err() { ApiError::Status { status, .. } => { assert_eq!(status, reqwest::StatusCode::NOT_FOUND); @@ -1355,7 +1393,7 @@ mod tests { .with_body("forbidden") .create(); - let api = Api::test_new(&server.url(), "test-jwt", None); + let api = Api::test_new_scoped(&server.url(), "test-jwt", None, Some("db-1")); match api.get_result_arrow("forbidden").unwrap_err() { ApiError::Status { status, body } => { assert_eq!(status, reqwest::StatusCode::FORBIDDEN); diff --git a/src/commands/queries.rs b/src/commands/queries.rs index 29ac339..5780fff 100644 --- a/src/commands/queries.rs +++ b/src/commands/queries.rs @@ -194,19 +194,21 @@ fn truncate_sql(sql: &str, max: usize) -> String { pub fn list( workspace_id: &str, + database: Option<&str>, limit: Option, cursor: Option<&str>, status: Option<&str>, format: &str, ) { - let api = Api::new(Some(workspace_id)); + let api = scoped_api(workspace_id, database); + let database_id = api.require_database(); let resp = crate::client::sdk::block_with_wakeup( &api, "Loading query runs…", api.client() .query_runs() - .list(limit.map(|l| l as i32), cursor, status, None), + .list(database_id, limit.map(|l| l as i32), cursor, status, None), ) .unwrap_or_else(|e| e.exit()); @@ -257,18 +259,30 @@ pub fn list( } } -pub fn get(query_run_id: &str, workspace_id: &str, format: &str) { - let api = Api::new(Some(workspace_id)); +pub fn get(query_run_id: &str, workspace_id: &str, database: Option<&str>, format: &str) { + let api = scoped_api(workspace_id, database); + let database_id = api.require_database(); let run: QueryRun = crate::client::sdk::block_with_wakeup( &api, "Loading query run…", - api.client().query_runs().get(query_run_id), + api.client().query_runs().get(query_run_id, database_id), ) .unwrap_or_else(|e| e.exit()) .into(); print_detail(&run, format); } +/// Build an [`Api`] scoped to `database` when the caller passed `--database`, +/// otherwise the database resolved at construction (`HOTDATA_DATABASE` / +/// current database). +fn scoped_api(workspace_id: &str, database: Option<&str>) -> Api { + let api = Api::new(Some(workspace_id)); + match database { + Some(db) => api.scoped_to_database(db.to_string()), + None => api, + } +} + fn print_detail(r: &QueryRun, format: &str) { match format { "json" => println!("{}", serde_json::to_string_pretty(r).unwrap()), diff --git a/src/commands/query.rs b/src/commands/query.rs index 3a97b48..357d339 100644 --- a/src/commands/query.rs +++ b/src/commands/query.rs @@ -426,12 +426,16 @@ fn fail_run(error_msg: &str) -> ! { } pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &str) { - let api = Api::new(Some(workspace_id)); - // Scope to the explicit --database flag, else the active database resolved - // at construction (HOTDATA_DATABASE / current database). submit_query sends - // it as the X-Database-Id header. - let database = database.or(api.database_id()); + // at construction (HOTDATA_DATABASE / current database). The scoped `Api` + // carries the database into submit_query's `X-Database-Id` header and into + // the database-scoped follow-up fetches (query-run poll, Arrow result). + let api = Api::new(Some(workspace_id)); + let api = match database { + Some(db) => api.scoped_to_database(db.to_string()), + None => api, + }; + let database = api.database_id(); let mut request = hotdata::models::QueryRequest::new(sql.to_string()); request.r#async = Some(true); @@ -471,8 +475,12 @@ pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &s loop { // Drive the poll loop ourselves to preserve the 5-minute deadline and // 500ms cadence (NOT the SDK's PollConfig defaults). - let run = crate::client::sdk::block(api.client().query_runs().get(run_id)) - .unwrap_or_else(|e| e.exit()); + let run = crate::client::sdk::block( + api.client() + .query_runs() + .get(run_id, api.require_database()), + ) + .unwrap_or_else(|e| e.exit()); match run.status.as_str() { "succeeded" => { spinner.finish_and_clear(); @@ -525,11 +533,19 @@ pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &s } /// Poll a query run by ID. If succeeded and has a result_id, fetch and display the result. -pub fn poll(query_run_id: &str, workspace_id: &str, format: &str) { +pub fn poll(query_run_id: &str, workspace_id: &str, database: Option<&str>, format: &str) { let api = Api::new(Some(workspace_id)); + let api = match database { + Some(db) => api.scoped_to_database(db.to_string()), + None => api, + }; - let run = crate::client::sdk::block(api.client().query_runs().get(query_run_id)) - .unwrap_or_else(|e| e.exit()); + let run = crate::client::sdk::block( + api.client() + .query_runs() + .get(query_run_id, api.require_database()), + ) + .unwrap_or_else(|e| e.exit()); match run.status.as_str() { "succeeded" => { @@ -779,7 +795,7 @@ mod tests { let mut resp = truncated_preview(Some("res_1")); resp.warning = Some(Some("approximate aggregate".to_string())); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); let resolved = resolve_inline(&api, resp); // Followed the truncated preview to the full 3-row result. @@ -813,7 +829,7 @@ mod tests { .with_body("boom") .create(); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); let resolved = resolve_inline(&api, truncated_preview(Some("res_1"))); // Preview kept, flagged incomplete so print_result fails closed. @@ -836,7 +852,7 @@ mod tests { // truncated=false short-circuits before any network call; point the Api // at a server with no mocks so an erroneous fetch would fail loudly. let server = mockito::Server::new(); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); let mut resp = hotdata::models::QueryResponse::new( vec!["x".to_string()], @@ -867,7 +883,7 @@ mod tests { // Truncated but persistence never started (result_id is null): the full // result is unfetchable, so keep the preview and surface a warning. let server = mockito::Server::new(); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); // The server reports the grand total even though it couldn't persist; // it must survive onto the preview so structured output exposes it. @@ -998,7 +1014,7 @@ mod tests { // warning explaining why persistence didn't start. The truncation note // is appended to it, not allowed to clobber it. let server = mockito::Server::new(); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); let mut resp = truncated_preview(None); resp.warning = Some(Some( @@ -1061,7 +1077,7 @@ mod tests { .with_body(ipc) .create(); - let api = Api::test_new(&server.url(), "test-jwt", Some("ws-1")); + let api = Api::test_new_scoped(&server.url(), "test-jwt", Some("ws-1"), Some("db-1")); // The async poll response reported the run took 4200ms. let result = fetch_arrow_result_with_timing(&api, "res_1", Some(Some(4200))); diff --git a/src/commands/results.rs b/src/commands/results.rs index b3cb48a..c969c5b 100644 --- a/src/commands/results.rs +++ b/src/commands/results.rs @@ -41,8 +41,18 @@ struct ListResponse { has_more: bool, } -pub fn list(workspace_id: &str, limit: Option, offset: Option, format: &str) { - let api = Api::new(Some(workspace_id)); +pub fn list( + workspace_id: &str, + database: Option<&str>, + limit: Option, + offset: Option, + format: &str, +) { + let api = scoped_api(workspace_id, database); + // Results are database-scoped (the required `X-Database-Id` header the seam + // sends from the active database). Fail early with a hint when none is set, + // rather than surfacing the raw server error. + api.require_database(); let mut params: Vec<(&str, String)> = Vec::new(); if let Some(l) = limit { @@ -130,8 +140,19 @@ pub fn list(workspace_id: &str, limit: Option, offset: Option, format: } } -pub fn get(result_id: &str, workspace_id: &str, format: &str) { - let api = Api::new(Some(workspace_id)); +pub fn get(result_id: &str, workspace_id: &str, database: Option<&str>, format: &str) { + let api = scoped_api(workspace_id, database); let result = crate::commands::query::fetch_arrow_result(&api, result_id); crate::commands::query::print_result(&result, format); } + +/// Build an [`Api`] scoped to `database` when the caller passed `--database`, +/// otherwise the database resolved at construction (`HOTDATA_DATABASE` / +/// current database). +fn scoped_api(workspace_id: &str, database: Option<&str>) -> Api { + let api = Api::new(Some(workspace_id)); + match database { + Some(db) => api.scoped_to_database(db.to_string()), + None => api, + } +} diff --git a/src/main.rs b/src/main.rs index 2bc9260..bd3b325 100644 --- a/src/main.rs +++ b/src/main.rs @@ -208,7 +208,9 @@ fn main() { } => { let workspace_id = resolve_workspace(workspace_id); match command { - Some(QueryCommands::Status { id }) => query::poll(&id, &workspace_id, &output), + Some(QueryCommands::Status { id }) => { + query::poll(&id, &workspace_id, database.as_deref(), &output) + } None => match sql { Some(sql) => { query::execute(&sql, &workspace_id, database.as_deref(), &output) @@ -517,6 +519,7 @@ fn main() { Commands::Results { result_id, workspace_id, + database, output, command, } => { @@ -526,9 +529,9 @@ fn main() { limit, offset, output, - }) => results::list(&workspace_id, limit, offset, &output), + }) => results::list(&workspace_id, database.as_deref(), limit, offset, &output), None => match result_id { - Some(id) => results::get(&id, &workspace_id, &output), + Some(id) => results::get(&id, &workspace_id, database.as_deref(), &output), None => { use clap::CommandFactory; let mut cmd = Cli::command(); @@ -814,12 +817,13 @@ fn main() { } Commands::Queries { id, + database, output, command, } => { let workspace_id = resolve_workspace(None); if let Some(id) = id { - queries::get(&id, &workspace_id, &output) + queries::get(&id, &workspace_id, database.as_deref(), &output) } else { match command { Some(QueriesCommands::List { @@ -829,6 +833,7 @@ fn main() { output, }) => queries::list( &workspace_id, + database.as_deref(), Some(limit), cursor.as_deref(), status.as_deref(), From a0e1050f5c7f1e64cb47014761992f936f4f6325 Mon Sep 17 00:00:00 2001 From: Zac Farrell Date: Tue, 7 Jul 2026 09:54:27 -0700 Subject: [PATCH 2/2] refactor(seam): consolidate database scoping into scoped_to_database_opt --- src/client/sdk.rs | 17 +++++++++-------- src/commands/queries.rs | 15 ++------------- src/commands/query.rs | 12 ++---------- src/commands/results.rs | 15 ++------------- 4 files changed, 15 insertions(+), 44 deletions(-) diff --git a/src/client/sdk.rs b/src/client/sdk.rs index d8273be..9cf931b 100644 --- a/src/client/sdk.rs +++ b/src/client/sdk.rs @@ -609,15 +609,16 @@ impl Api { self.database_id.as_deref() } - /// Return a clone of this `Api` scoped to `database_id`, overriding any - /// database resolved at construction (`HOTDATA_DATABASE` / current - /// database). Used by commands whose `--database` flag scopes a single - /// invocation without switching the persisted active database. - pub fn scoped_to_database(&self, database_id: String) -> Self { - Self { - database_id: Some(database_id), - ..self.clone() + /// Scope this `Api` to `database` when `Some`, overriding any database + /// resolved at construction (`HOTDATA_DATABASE` / current database); a + /// `None` leaves the resolved scope untouched. Used by commands whose + /// `--database` flag scopes a single invocation without switching the + /// persisted active database — `Api::new(..).scoped_to_database_opt(flag)`. + pub fn scoped_to_database_opt(mut self, database: Option<&str>) -> Self { + if let Some(db) = database { + self.database_id = Some(db.to_string()); } + self } /// The active database scope, or exit with actionable guidance when none is diff --git a/src/commands/queries.rs b/src/commands/queries.rs index 5780fff..ea0b9d5 100644 --- a/src/commands/queries.rs +++ b/src/commands/queries.rs @@ -200,7 +200,7 @@ pub fn list( status: Option<&str>, format: &str, ) { - let api = scoped_api(workspace_id, database); + let api = Api::new(Some(workspace_id)).scoped_to_database_opt(database); let database_id = api.require_database(); let resp = crate::client::sdk::block_with_wakeup( @@ -260,7 +260,7 @@ pub fn list( } pub fn get(query_run_id: &str, workspace_id: &str, database: Option<&str>, format: &str) { - let api = scoped_api(workspace_id, database); + let api = Api::new(Some(workspace_id)).scoped_to_database_opt(database); let database_id = api.require_database(); let run: QueryRun = crate::client::sdk::block_with_wakeup( &api, @@ -272,17 +272,6 @@ pub fn get(query_run_id: &str, workspace_id: &str, database: Option<&str>, forma print_detail(&run, format); } -/// Build an [`Api`] scoped to `database` when the caller passed `--database`, -/// otherwise the database resolved at construction (`HOTDATA_DATABASE` / -/// current database). -fn scoped_api(workspace_id: &str, database: Option<&str>) -> Api { - let api = Api::new(Some(workspace_id)); - match database { - Some(db) => api.scoped_to_database(db.to_string()), - None => api, - } -} - fn print_detail(r: &QueryRun, format: &str) { match format { "json" => println!("{}", serde_json::to_string_pretty(r).unwrap()), diff --git a/src/commands/query.rs b/src/commands/query.rs index 357d339..c56f99d 100644 --- a/src/commands/query.rs +++ b/src/commands/query.rs @@ -430,11 +430,7 @@ pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &s // at construction (HOTDATA_DATABASE / current database). The scoped `Api` // carries the database into submit_query's `X-Database-Id` header and into // the database-scoped follow-up fetches (query-run poll, Arrow result). - let api = Api::new(Some(workspace_id)); - let api = match database { - Some(db) => api.scoped_to_database(db.to_string()), - None => api, - }; + let api = Api::new(Some(workspace_id)).scoped_to_database_opt(database); let database = api.database_id(); let mut request = hotdata::models::QueryRequest::new(sql.to_string()); @@ -534,11 +530,7 @@ pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &s /// Poll a query run by ID. If succeeded and has a result_id, fetch and display the result. pub fn poll(query_run_id: &str, workspace_id: &str, database: Option<&str>, format: &str) { - let api = Api::new(Some(workspace_id)); - let api = match database { - Some(db) => api.scoped_to_database(db.to_string()), - None => api, - }; + let api = Api::new(Some(workspace_id)).scoped_to_database_opt(database); let run = crate::client::sdk::block( api.client() diff --git a/src/commands/results.rs b/src/commands/results.rs index c969c5b..2904944 100644 --- a/src/commands/results.rs +++ b/src/commands/results.rs @@ -48,7 +48,7 @@ pub fn list( offset: Option, format: &str, ) { - let api = scoped_api(workspace_id, database); + let api = Api::new(Some(workspace_id)).scoped_to_database_opt(database); // Results are database-scoped (the required `X-Database-Id` header the seam // sends from the active database). Fail early with a hint when none is set, // rather than surfacing the raw server error. @@ -141,18 +141,7 @@ pub fn list( } pub fn get(result_id: &str, workspace_id: &str, database: Option<&str>, format: &str) { - let api = scoped_api(workspace_id, database); + let api = Api::new(Some(workspace_id)).scoped_to_database_opt(database); let result = crate::commands::query::fetch_arrow_result(&api, result_id); crate::commands::query::print_result(&result, format); } - -/// Build an [`Api`] scoped to `database` when the caller passed `--database`, -/// otherwise the database resolved at construction (`HOTDATA_DATABASE` / -/// current database). -fn scoped_api(workspace_id: &str, database: Option<&str>) -> Api { - let api = Api::new(Some(workspace_id)); - match database { - Some(db) => api.scoped_to_database(db.to_string()), - None => api, - } -}