From d431aa220b152d38117ee03559ac654483878025 Mon Sep 17 00:00:00 2001 From: Eddie A Tejeda <669988+eddietejeda@users.noreply.github.com> Date: Fri, 19 Jun 2026 11:06:10 -0700 Subject: [PATCH] feat(cli): show "waking up worker" hint on KEDA cold starts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap worker-bound commands in block_with_wakeup: show a spinner and, if the request hasn't returned within ~1.5s, probe the control plane's /v1/workspaces//runtime/status and upgrade the message when the worker is cold. Warm requests pay nothing — the probe only fires after the delay and is dropped the instant the real response lands. The probe omits X-Workspace-Id so it reaches the always-warm control plane, not the KEDA interceptor. Wired into list/get/create/update/delete/refresh across databases, connections, datasets, tables, context, embedding providers, queries, jobs, indexes, and query execute. Polling loops, pagination, rayon scans, and internal helpers stay on plain block(). --- src/connections.rs | 56 +++++++---- src/connections_new.rs | 10 +- src/context.rs | 15 ++- src/databases.rs | 33 ++++--- src/datasets.rs | 29 ++++-- src/embedding_providers.rs | 49 ++++++--- src/indexes.rs | 10 +- src/jobs.rs | 7 +- src/queries.rs | 23 +++-- src/query.rs | 10 +- src/sdk.rs | 198 +++++++++++++++++++++++++++++++++++++ src/tables.rs | 20 ++-- 12 files changed, 374 insertions(+), 86 deletions(-) diff --git a/src/connections.rs b/src/connections.rs index 5b989bc..4b58dc2 100644 --- a/src/connections.rs +++ b/src/connections.rs @@ -1,4 +1,4 @@ -use crate::sdk::{Api, ApiError, block, none_if_404}; +use crate::sdk::{Api, ApiError, block, block_with_wakeup, none_if_404}; use serde::{Deserialize, Serialize}; #[derive(Deserialize, Serialize)] @@ -99,7 +99,12 @@ struct ConnectionTypeDetail { pub fn types_list(workspace_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let resp = block(api.client().connection_types().list()).unwrap_or_else(|e| e.exit()); + let resp = block_with_wakeup( + &api, + "Loading connection types…", + api.client().connection_types().list(), + ) + .unwrap_or_else(|e| e.exit()); let body = ListConnectionTypesResponse { connection_types: resp .connection_types @@ -136,7 +141,12 @@ pub fn types_list(workspace_id: &str, format: &str) { pub fn types_get(workspace_id: &str, name: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let resp = block(api.client().connection_types().get(name)).unwrap_or_else(|e| e.exit()); + let resp = block_with_wakeup( + &api, + "Loading connection type…", + api.client().connection_types().get(name), + ) + .unwrap_or_else(|e| e.exit()); // The SDK models nullable fields as `Option>`; flatten and // drop an explicit JSON `null` to match the old behavior (the old struct // deserialized a missing/`null` field to `None`). @@ -243,11 +253,18 @@ pub fn get(workspace_id: &str, connection_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); let is_table = format == "table"; - let spinner = is_table.then(|| crate::util::spinner("Fetching connection...")); - let resp = block(api.client().connections().get(connection_id)).unwrap_or_else(|e| e.exit()); - if let Some(s) = spinner { - s.finish_and_clear(); + // Keep the spinner table-only (scripting output stays clean), but route the + // interactive path through the wakeup hint so a cold KEDA start is explained. + let resp = if is_table { + block_with_wakeup( + &api, + "Fetching connection...", + api.client().connections().get(connection_id), + ) + } else { + block(api.client().connections().get(connection_id)) } + .unwrap_or_else(|e| e.exit()); let detail = ConnectionDetail { id: resp.id, name: resp.name, @@ -327,16 +344,16 @@ pub fn create(workspace_id: &str, name: &str, source_type: &str, config: &str, f source_type.to_string(), ); - let spinner = is_table.then(|| crate::util::spinner("Creating connection...")); - let resp = block(api.client().connections().create(request)).unwrap_or_else(|e| { - if let Some(s) = &spinner { - s.finish_and_clear(); - } - e.exit() - }); - if let Some(s) = &spinner { - s.finish_and_clear(); + let resp = if is_table { + block_with_wakeup( + &api, + "Creating connection...", + api.client().connections().create(request), + ) + } else { + block(api.client().connections().create(request)) } + .unwrap_or_else(|e| e.exit()); let result = CreateResponse { id: resp.id, @@ -404,7 +421,12 @@ pub fn create(workspace_id: &str, name: &str, source_type: &str, config: &str, f pub fn list(workspace_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let resp = block(api.client().connections().list()).unwrap_or_else(|e| e.exit()); + let resp = block_with_wakeup( + &api, + "Loading connections…", + api.client().connections().list(), + ) + .unwrap_or_else(|e| e.exit()); let body = ListResponse { connections: resp .connections diff --git a/src/connections_new.rs b/src/connections_new.rs index be9465b..feb1d05 100644 --- a/src/connections_new.rs +++ b/src/connections_new.rs @@ -2,7 +2,7 @@ use inquire::validator::Validation; use inquire::{Confirm, Password, Select, Text}; use serde_json::{Map, Number, Value}; -use crate::sdk::{Api, ApiError, block}; +use crate::sdk::{Api, ApiError, block, block_with_wakeup}; // ── SDK helpers ───────────────────────────────────────────────────────────── @@ -326,9 +326,11 @@ pub fn run(workspace_id: &str) { } } - let create_spinner = crate::util::spinner("Creating connection..."); - let result = block(api.client().connections().create(request)); - create_spinner.finish_and_clear(); + let result = block_with_wakeup( + &api, + "Creating connection...", + api.client().connections().create(request), + ); use crossterm::style::Stylize; let result = match result { diff --git a/src/context.rs b/src/context.rs index bfd1b31..768b8ff 100644 --- a/src/context.rs +++ b/src/context.rs @@ -131,8 +131,12 @@ fn fetch_context(api: &Api, database_id: &str, name: &str) -> Option, format: &str) { let api = Api::new(Some(workspace_id)); - let body = crate::sdk::block(api.client().database_context().list(database_id)) - .unwrap_or_else(|e| e.exit()); + let body = crate::sdk::block_with_wakeup( + &api, + "Loading context…", + api.client().database_context().list(database_id), + ) + .unwrap_or_else(|e| e.exit()); let mut rows: Vec = body.contexts.into_iter().map(ContextRow::from).collect(); if let Some(p) = prefix { @@ -295,8 +299,11 @@ pub fn push(workspace_id: &str, database_id: &str, name: &str, dry_run: bool) { let api = Api::new(Some(workspace_id)); let request = UpsertDatabaseContextRequest::new(content, name.clone()); - let resp = match crate::sdk::block(api.client().database_context().upsert(database_id, request)) - { + let resp = match crate::sdk::block_with_wakeup( + &api, + "Pushing context…", + api.client().database_context().upsert(database_id, request), + ) { Ok(resp) => resp, Err(ApiError::Status { status: _, body }) => { let msg = crate::util::api_error(body); diff --git a/src/databases.rs b/src/databases.rs index b4f7de1..52c6b33 100644 --- a/src/databases.rs +++ b/src/databases.rs @@ -1,4 +1,4 @@ -use crate::sdk::{Api, ApiError, block, none_if_404}; +use crate::sdk::{Api, ApiError, block, block_with_wakeup, none_if_404}; use indicatif::{ProgressBar, ProgressStyle}; use serde::{Deserialize, Serialize}; use std::path::Path; @@ -139,8 +139,12 @@ pub(crate) fn get_database(api: &Api, id: &str) -> Result { /// List databases through the SDK's typed `databases().list` handle, mapped /// into the CLI's summary rows. +/// +/// Routed through [`block_with_wakeup`] so a cold KEDA start surfaces a "waking +/// up worker" hint instead of an unexplained pause — `databases list` is a +/// common first command against an idle workspace. fn list_database_summaries(api: &Api) -> Result, ApiError> { - block(api.client().databases().list()) + block_with_wakeup(api, "Loading databases…", api.client().databases().list()) .map(|r| r.databases.into_iter().map(DatabaseSummary::from).collect()) } @@ -678,16 +682,16 @@ pub fn create( let request = create_database_typed_request(name, catalog, schema, tables, expires_at); let api = Api::new(Some(workspace_id)); - let spinner = (format == "table").then(|| crate::util::spinner("Creating database...")); - let resp = block(api.client().databases().create(request)).unwrap_or_else(|e| { - if let Some(s) = &spinner { - s.finish_and_clear(); - } - e.exit() - }); - if let Some(s) = &spinner { - s.finish_and_clear(); + let resp = if format == "table" { + block_with_wakeup( + &api, + "Creating database...", + api.client().databases().create(request), + ) + } else { + block(api.client().databases().create(request)) } + .unwrap_or_else(|e| e.exit()); let result = CreateDatabaseResponse { id: resp.id, @@ -794,7 +798,12 @@ pub fn delete(workspace_id: &str, id_or_name: &str) { let api = Api::new(Some(workspace_id)); let db = resolve_database(&api, id_or_name); - block(api.client().databases().delete(&db.id)).unwrap_or_else(|e| e.exit()); + block_with_wakeup( + &api, + "Deleting database…", + api.client().databases().delete(&db.id), + ) + .unwrap_or_else(|e| e.exit()); // If the deleted database was the current one, clear it so subsequent // commands don't silently send a stale X-Database-Id header. diff --git a/src/datasets.rs b/src/datasets.rs index ab88ed7..ff009e5 100644 --- a/src/datasets.rs +++ b/src/datasets.rs @@ -158,7 +158,9 @@ pub fn create_from_saved_query( pub fn list(workspace_id: &str, limit: Option, offset: Option, format: &str) { let api = Api::new(Some(workspace_id)); - let body = crate::sdk::block( + let body = crate::sdk::block_with_wakeup( + &api, + "Loading datasets…", api.client() .datasets() .list(limit.map(|l| l as i32), offset.map(|o| o as i32)), @@ -219,8 +221,12 @@ pub fn list(workspace_id: &str, limit: Option, offset: Option, format: pub fn get(dataset_id: &str, workspace_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let resp: GetDatasetResponse = - crate::sdk::block(api.client().datasets().get(dataset_id)).unwrap_or_else(|e| e.exit()); + let resp: GetDatasetResponse = crate::sdk::block_with_wakeup( + &api, + "Loading dataset…", + api.client().datasets().get(dataset_id), + ) + .unwrap_or_else(|e| e.exit()); let d = DatasetDetail { id: resp.id, @@ -295,9 +301,12 @@ pub fn update( request.table_name = Some(Some(n.to_string())); } - let resp: UpdateDatasetResponse = - crate::sdk::block(api.client().datasets().update(dataset_id, request)) - .unwrap_or_else(|e| e.exit()); + let resp: UpdateDatasetResponse = crate::sdk::block_with_wakeup( + &api, + "Updating dataset…", + api.client().datasets().update(dataset_id, request), + ) + .unwrap_or_else(|e| e.exit()); let d = UpdateView::from(resp); use crossterm::style::Stylize; @@ -341,8 +350,12 @@ pub fn refresh(workspace_id: &str, dataset_id: &str, async_mode: bool) { request.r#async = Some(true); } - let resp = - crate::sdk::block(api.client().refresh().refresh(request)).unwrap_or_else(|e| e.exit()); + let resp = crate::sdk::block_with_wakeup( + &api, + "Refreshing dataset…", + api.client().refresh().refresh(request), + ) + .unwrap_or_else(|e| e.exit()); if async_mode { let job_id = match &resp { diff --git a/src/embedding_providers.rs b/src/embedding_providers.rs index 2e99525..c5525d7 100644 --- a/src/embedding_providers.rs +++ b/src/embedding_providers.rs @@ -44,12 +44,16 @@ fn parse_config(raw: Option<&str>) -> Option { pub fn list(workspace_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let providers: Vec = crate::sdk::block(api.client().embedding_providers().list()) - .unwrap_or_else(|e| e.exit()) - .embedding_providers - .into_iter() - .map(Provider::from) - .collect(); + let providers: Vec = crate::sdk::block_with_wakeup( + &api, + "Loading embedding providers…", + api.client().embedding_providers().list(), + ) + .unwrap_or_else(|e| e.exit()) + .embedding_providers + .into_iter() + .map(Provider::from) + .collect(); use crossterm::style::Stylize; match format { @@ -80,9 +84,13 @@ pub fn list(workspace_id: &str, format: &str) { pub fn get(workspace_id: &str, id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let p: Provider = crate::sdk::block(api.client().embedding_providers().get(id)) - .unwrap_or_else(|e| e.exit()) - .into(); + let p: Provider = crate::sdk::block_with_wakeup( + &api, + "Loading embedding provider…", + api.client().embedding_providers().get(id), + ) + .unwrap_or_else(|e| e.exit()) + .into(); match format { "json" => println!("{}", serde_json::to_string_pretty(&p).unwrap()), @@ -128,8 +136,12 @@ pub fn create( req.secret_name = Some(Some(s.to_string())); } - let resp = crate::sdk::block(api.client().embedding_providers().create(req)) - .unwrap_or_else(|e| e.exit()); + let resp = crate::sdk::block_with_wakeup( + &api, + "Creating embedding provider…", + api.client().embedding_providers().create(req), + ) + .unwrap_or_else(|e| e.exit()); let parsed = serde_json::to_value(&resp).unwrap_or_default(); eprintln!("{}", "Embedding provider created.".green()); @@ -180,8 +192,12 @@ pub fn update( req.secret_name = Some(Some(s.to_string())); } - let resp = crate::sdk::block(api.client().embedding_providers().update(id, req)) - .unwrap_or_else(|e| e.exit()); + let resp = crate::sdk::block_with_wakeup( + &api, + "Updating embedding provider…", + api.client().embedding_providers().update(id, req), + ) + .unwrap_or_else(|e| e.exit()); let resp = serde_json::to_value(&resp).unwrap_or_default(); eprintln!("{}", "Embedding provider updated.".green()); @@ -202,7 +218,12 @@ pub fn update( pub fn delete(workspace_id: &str, id: &str) { use crossterm::style::Stylize; let api = Api::new(Some(workspace_id)); - crate::sdk::block(api.client().embedding_providers().delete(id)).unwrap_or_else(|e| e.exit()); + crate::sdk::block_with_wakeup( + &api, + "Deleting embedding provider…", + api.client().embedding_providers().delete(id), + ) + .unwrap_or_else(|e| e.exit()); println!("{}", format!("Embedding provider '{id}' deleted.").green()); } diff --git a/src/indexes.rs b/src/indexes.rs index ffdac4a..dd71e0a 100644 --- a/src/indexes.rs +++ b/src/indexes.rs @@ -1,4 +1,4 @@ -use crate::sdk::{Api, block, none_if_404}; +use crate::sdk::{Api, block, block_with_wakeup, none_if_404}; use rayon::prelude::*; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -548,12 +548,16 @@ pub fn delete(workspace_id: &str, scope: IndexScope<'_>, index_name: &str) { connection_id, schema, table, - } => block( + } => block_with_wakeup( + &api, + "Deleting index…", api.client() .indexes() .delete_index(connection_id, schema, table, index_name), ), - IndexScope::Dataset { dataset_id } => block( + IndexScope::Dataset { dataset_id } => block_with_wakeup( + &api, + "Deleting index…", api.client() .indexes() .delete_dataset_index(dataset_id, index_name), diff --git a/src/jobs.rs b/src/jobs.rs index ce1ff11..1f9c634 100644 --- a/src/jobs.rs +++ b/src/jobs.rs @@ -49,9 +49,10 @@ fn parse_job_type(s: &str) -> Option { pub fn get(job_id: &str, workspace_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let job: Job = crate::sdk::block(api.client().jobs().get(job_id)) - .unwrap_or_else(|e| e.exit()) - .into(); + let job: Job = + crate::sdk::block_with_wakeup(&api, "Loading job…", api.client().jobs().get(job_id)) + .unwrap_or_else(|e| e.exit()) + .into(); match format { "json" => println!("{}", serde_json::to_string_pretty(&job).unwrap()), diff --git a/src/queries.rs b/src/queries.rs index 42c05b6..cdbde3e 100644 --- a/src/queries.rs +++ b/src/queries.rs @@ -178,12 +178,13 @@ pub fn list( ) { let api = Api::new(Some(workspace_id)); - let resp = crate::sdk::block(api.client().query_runs().list( - limit.map(|l| l as i32), - cursor, - status, - None, - )) + let resp = crate::sdk::block_with_wakeup( + &api, + "Loading query runs…", + api.client() + .query_runs() + .list(limit.map(|l| l as i32), cursor, status, None), + ) .unwrap_or_else(|e| e.exit()); let query_runs: Vec = resp.query_runs.into_iter().map(QueryRun::from).collect(); @@ -235,9 +236,13 @@ pub fn list( pub fn get(query_run_id: &str, workspace_id: &str, format: &str) { let api = Api::new(Some(workspace_id)); - let run: QueryRun = crate::sdk::block(api.client().query_runs().get(query_run_id)) - .unwrap_or_else(|e| e.exit()) - .into(); + let run: QueryRun = crate::sdk::block_with_wakeup( + &api, + "Loading query run…", + api.client().query_runs().get(query_run_id), + ) + .unwrap_or_else(|e| e.exit()) + .into(); print_detail(&run, format); } diff --git a/src/query.rs b/src/query.rs index 72fe43e..e4ed03e 100644 --- a/src/query.rs +++ b/src/query.rs @@ -264,10 +264,12 @@ pub fn execute(sql: &str, workspace_id: &str, database: Option<&str>, format: &s request.r#async = Some(true); request.async_after_ms = Some(Some(1000)); - let spinner = crate::util::spinner("running query..."); - let outcome = crate::sdk::block(api.client().submit_query(request, database)) - .unwrap_or_else(|e| e.exit()); - spinner.finish_and_clear(); + let outcome = crate::sdk::block_with_wakeup( + &api, + "running query...", + api.client().submit_query(request, database), + ) + .unwrap_or_else(|e| e.exit()); let async_resp = match outcome { // Completed within async_after_ms — inline results. A large result can diff --git a/src/sdk.rs b/src/sdk.rs index a2404d0..dcd23a7 100644 --- a/src/sdk.rs +++ b/src/sdk.rs @@ -248,6 +248,91 @@ where rt().block_on(fut).map_err(ApiError::from_sdk) } +/// KEDA scale state of a workspace's runtimedb worker, as reported by the +/// always-warm control plane. Used only to upgrade a spinner message on a +/// cold start — it never affects control flow. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RuntimeState { + /// At least one ready replica; the request is being served warm. + Ready, + /// Scaling up (desired >= 1) but no ready replica yet. + Waking, + /// Scaled to zero; the request triggered a cold start. + Asleep, + /// Couldn't determine (no control plane, error, non-2xx, unscoped). + Unknown, +} + +impl RuntimeState { + fn from_state_str(s: &str) -> Self { + match s { + "ready" => Self::Ready, + "waking" => Self::Waking, + "asleep" => Self::Asleep, + _ => Self::Unknown, + } + } + + /// True when the worker is cold — i.e. the in-flight request is (or will be) + /// blocked waiting for KEDA to bring a replica up. + fn is_cold(self) -> bool { + matches!(self, Self::Waking | Self::Asleep) + } +} + +/// How long a request may run before we suspect a cold start and probe the +/// control plane. Short enough that a real wake-up is flagged promptly, long +/// enough that a warm-but-slow request never triggers the probe. +const WAKE_PROBE_DELAY: std::time::Duration = std::time::Duration::from_millis(1500); +/// Upper bound on the status probe itself, so a slow/stuck probe never delays +/// the real response (the probe runs in a select arm that's dropped the moment +/// the real request completes, but this caps its own footprint regardless). +const WAKE_PROBE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3); +/// Spinner message shown once a cold start is confirmed. +const WAKE_MESSAGE: &str = "waking up worker after inactivity (this can take ~20s)…"; + +/// Like [`block`], but shows a spinner with `msg` and, if the request hasn't +/// returned within [`WAKE_PROBE_DELAY`], probes the control plane for the +/// worker's scale state. On a confirmed cold start the spinner message is +/// upgraded to explain the wait. Warm requests pay nothing: the probe only +/// fires after the delay, and returns the real result the instant it lands. +pub fn block_with_wakeup(api: &Api, msg: &str, fut: F) -> Result +where + F: std::future::Future>>, + E: std::fmt::Debug, +{ + let pb = util::spinner(msg); + let hint_pb = pb.clone(); + let result = rt().block_on(async { + tokio::pin!(fut); + // After the delay, probe once and (if cold) upgrade the message, then + // idle forever so the select keeps driving the real request to + // completion. Dropped — cancelling any in-flight probe — as soon as + // `fut` wins. + let hint = async { + tokio::time::sleep(WAKE_PROBE_DELAY).await; + let state = tokio::time::timeout(WAKE_PROBE_TIMEOUT, api.probe_runtime_status()) + .await + .unwrap_or(RuntimeState::Unknown); + if state.is_cold() { + hint_pb.set_message(WAKE_MESSAGE); + } + std::future::pending::<()>().await; + }; + tokio::pin!(hint); + loop { + tokio::select! { + r = &mut fut => break r, + // `hint` ends in `pending()`, so this arm never completes; it + // exists only to drive the probe alongside the real request. + _ = &mut hint => {} + } + } + }); + pb.finish_and_clear(); + result.map_err(ApiError::from_sdk) +} + /// Map a result, returning `Ok(None)` on HTTP 404 instead of an error. /// /// Reproduces `ApiClient::get_none_if_not_found` / the context-404 / indexes-404 @@ -479,6 +564,53 @@ impl Api { self.database_id.as_deref() } + /// Best-effort probe of the scoped workspace's runtimedb scale state, via + /// the control-plane `GET /v1/workspaces/{id}/runtime/status` endpoint. + /// + /// Returns [`RuntimeState::Unknown`] on any failure (no workspace scope, + /// transport error, non-2xx, or unparseable body) so callers degrade to + /// their latency heuristic rather than surfacing an error. + /// + /// Deliberately omits the `X-Workspace-Id` header: that header is what the + /// gateway matches to route `/v1` traffic to the KEDA interceptor, and + /// routing this probe there would wake the very worker we're asking about. + /// Without it the request lands on the always-warm control plane, which + /// answers from Kubernetes state. + async fn probe_runtime_status(&self) -> RuntimeState { + let Some(ws) = self.workspace_id.as_deref() else { + return RuntimeState::Unknown; + }; + let cfg = self.client.configuration(); + // `base_path` has no `/v1` suffix (see `sdk_base_path`); add the one + // the control-plane route expects. + let url = format!( + "{}/v1/workspaces/{}/runtime/status", + cfg.base_path.trim_end_matches('/'), + ws + ); + let mut req = cfg.client.get(&url); + if let Some(ref user_agent) = cfg.user_agent { + req = req.header(reqwest::header::USER_AGENT, user_agent.clone()); + } + if let Some(token) = cfg.resolve_bearer_token().await { + req = req.bearer_auth(token); + } + let Ok(resp) = req.send().await else { + return RuntimeState::Unknown; + }; + if !resp.status().is_success() { + return RuntimeState::Unknown; + } + match resp.json::().await { + Ok(body) => body + .get("state") + .and_then(|s| s.as_str()) + .map(RuntimeState::from_state_str) + .unwrap_or(RuntimeState::Unknown), + Err(_) => RuntimeState::Unknown, + } + } + /// Borrow the underlying SDK client (for command modules calling resource /// handles directly through [`block`]). pub fn client(&self) -> &Client { @@ -1294,4 +1426,70 @@ mod tests { other => panic!("expected Status error, got {other:?}"), } } + + #[test] + fn runtime_state_parsing_and_coldness() { + assert_eq!(RuntimeState::from_state_str("ready"), RuntimeState::Ready); + assert_eq!(RuntimeState::from_state_str("waking"), RuntimeState::Waking); + assert_eq!(RuntimeState::from_state_str("asleep"), RuntimeState::Asleep); + assert_eq!( + RuntimeState::from_state_str("garbage"), + RuntimeState::Unknown + ); + // Only the not-yet-serving states count as cold. + assert!(RuntimeState::Asleep.is_cold()); + assert!(RuntimeState::Waking.is_cold()); + assert!(!RuntimeState::Ready.is_cold()); + assert!(!RuntimeState::Unknown.is_cold()); + } + + #[test] + fn probe_runtime_status_reads_state_without_workspace_header() { + let mut server = mockito::Server::new(); + // The probe must hit the control-plane route *without* X-Workspace-Id, + // or the gateway would send it to the KEDA interceptor and wake the + // worker we're only trying to inspect. + let m = server + .mock("GET", "/v1/workspaces/work-1/runtime/status") + .match_header("x-workspace-id", mockito::Matcher::Missing) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"ok":true,"state":"asleep","estimated_wake_seconds":20}"#) + .create(); + + let api = Api::test_new(&server.url(), "test-jwt", Some("work-1")); + let state = rt().block_on(api.probe_runtime_status()); + assert_eq!(state, RuntimeState::Asleep); + m.assert(); + } + + #[test] + fn probe_runtime_status_non_2xx_is_unknown() { + let mut server = mockito::Server::new(); + let _m = server + .mock("GET", "/v1/workspaces/work-1/runtime/status") + .with_status(500) + .with_body("boom") + .create(); + + let api = Api::test_new(&server.url(), "test-jwt", Some("work-1")); + assert_eq!( + rt().block_on(api.probe_runtime_status()), + RuntimeState::Unknown + ); + } + + #[test] + fn probe_runtime_status_unscoped_is_unknown_without_request() { + // No workspace scope -> nothing to probe; must not make a request. + let mut server = mockito::Server::new(); + let m = server.mock("GET", mockito::Matcher::Any).expect(0).create(); + + let api = Api::test_new(&server.url(), "test-jwt", None); + assert_eq!( + rt().block_on(api.probe_runtime_status()), + RuntimeState::Unknown + ); + m.assert(); + } } diff --git a/src/tables.rs b/src/tables.rs index 7015e09..525d1d0 100644 --- a/src/tables.rs +++ b/src/tables.rs @@ -42,14 +42,18 @@ pub fn list( // the old behavior (include_columns=true iff connection_id is set). let include_columns = connection_id.map(|_| true); - let body = crate::sdk::block(api.client().information_schema().get( - connection_id, - schema, - table_filter, - include_columns, - limit.map(|l| l as i32), - cursor, - )) + let body = crate::sdk::block_with_wakeup( + &api, + "Loading tables…", + api.client().information_schema().get( + connection_id, + schema, + table_filter, + include_columns, + limit.map(|l| l as i32), + cursor, + ), + ) .unwrap_or_else(|e| e.exit()); let has_more = body.has_more;