fix(connections): Add health check to connection flow

pthurlow · web-flow · commit 1c8f8607ce68 · 2026-04-13T12:37:46.000-07:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/api.rs b/src/api.rs
@@ -171,6 +171,24 @@ impl ApiClient {
         }
     }
 
+    /// GET request, exits only on connection error, returns raw (status, body).
+    /// Use for best-effort endpoints (e.g. health checks) where the caller wants
+    /// to handle non-2xx responses gracefully instead of aborting.
+    pub fn get_raw(&self, path: &str) -> (reqwest::StatusCode, String) {
+        let url = format!("{}{path}", self.api_url);
+        self.log_request("GET", &url, None);
+
+        let resp = match self.build_request(reqwest::Method::GET, &url).send() {
+            Ok(r) => r,
+            Err(e) => {
+                eprintln!("error connecting to API: {e}");
+                std::process::exit(1);
+            }
+        };
+
+        util::debug_response(resp)
+    }
+
     /// POST request with JSON body, exits on error, returns raw (status, body).
     pub fn post_raw(&self, path: &str, body: &serde_json::Value) -> (reqwest::StatusCode, String) {
         let url = format!("{}{path}", self.api_url);
diff --git a/src/connections.rs b/src/connections.rs
@@ -1,6 +1,70 @@
 use crate::api::ApiClient;
 use serde::{Deserialize, Serialize};
 
+#[derive(Deserialize, Serialize)]
+struct HealthResponse {
+    #[allow(dead_code)]
+    connection_id: String,
+    healthy: bool,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    latency_ms: Option<u64>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    error: Option<String>,
+}
+
+/// Result of a best-effort health check. Either the endpoint responded with a
+/// parseable body, or it did not — in which case we record why and keep going.
+enum HealthStatus {
+    Available(HealthResponse),
+    Unavailable(String),
+}
+
+impl HealthStatus {
+    fn is_confirmed_unhealthy(&self) -> bool {
+        matches!(self, HealthStatus::Available(h) if !h.healthy)
+    }
+}
+
+impl Serialize for HealthStatus {
+    fn serialize<S: serde::Serializer>(&self, ser: S) -> Result<S::Ok, S::Error> {
+        match self {
+            HealthStatus::Available(h) => h.serialize(ser),
+            HealthStatus::Unavailable(_) => ser.serialize_none(),
+        }
+    }
+}
+
+fn fetch_health(api: &ApiClient, connection_id: &str, show_spinner: bool) -> HealthStatus {
+    let spinner = show_spinner.then(|| crate::util::spinner("Checking connection health..."));
+    let (status, body) = api.get_raw(&format!("/connections/{connection_id}/health"));
+    if let Some(s) = spinner { s.finish_and_clear(); }
+
+    if !status.is_success() {
+        return HealthStatus::Unavailable(crate::util::api_error(body));
+    }
+    match serde_json::from_str::<HealthResponse>(&body) {
+        Ok(h) => HealthStatus::Available(h),
+        Err(e) => HealthStatus::Unavailable(format!("parse error: {e}")),
+    }
+}
+
+fn format_health(health: &HealthStatus) -> String {
+    use crossterm::style::Stylize;
+    match health {
+        HealthStatus::Available(h) if h.healthy => match h.latency_ms {
+            Some(ms) => format!("{} {}", "healthy".green(), format!("({ms}ms)").dark_grey()),
+            None => "healthy".green().to_string(),
+        },
+        HealthStatus::Available(h) => {
+            let err = h.error.as_deref().unwrap_or("unknown error");
+            format!("{} — {}", "unhealthy".red(), err)
+        }
+        HealthStatus::Unavailable(err) => {
+            format!("{} — {}", "unavailable".yellow(), err)
+        }
+    }
+}
+
 #[derive(Deserialize, Serialize)]
 struct ConnectionType {
     name: String,
@@ -88,23 +152,60 @@ struct ListResponse {
 
 pub fn get(workspace_id: &str, connection_id: &str, format: &str) {
     let api = ApiClient::new(Some(workspace_id));
+    let is_table = format == "table";
+
+    let spinner = is_table.then(|| crate::util::spinner("Fetching connection..."));
     let detail: ConnectionDetail = api.get(&format!("/connections/{connection_id}"));
+    if let Some(s) = spinner { s.finish_and_clear(); }
+
+    let health = fetch_health(&api, connection_id, is_table);
 
     match format {
-        "json" => println!("{}", serde_json::to_string_pretty(&detail).unwrap()),
-        "yaml" => print!("{}", serde_yaml::to_string(&detail).unwrap()),
+        "json" => {
+            let combined = serde_json::json!({
+                "id": detail.id,
+                "name": detail.name,
+                "source_type": detail.source_type,
+                "table_count": detail.table_count,
+                "synced_table_count": detail.synced_table_count,
+                "health": &health,
+            });
+            println!("{}", serde_json::to_string_pretty(&combined).unwrap());
+        }
+        "yaml" => {
+            let combined = serde_json::json!({
+                "id": detail.id,
+                "name": detail.name,
+                "source_type": detail.source_type,
+                "table_count": detail.table_count,
+                "synced_table_count": detail.synced_table_count,
+                "health": &health,
+            });
+            print!("{}", serde_yaml::to_string(&combined).unwrap());
+        }
         "table" => {
             use crossterm::style::Stylize;
             let label = |l: &str| format!("{:<16}", l).dark_grey().to_string();
             println!("{}{}", label("id:"), detail.id.dark_cyan());
             println!("{}{}", label("name:"), detail.name.white());
             println!("{}{}", label("source_type:"), detail.source_type.green());
             println!("{}{}", label("tables:"), format!("{} synced / {} total", detail.synced_table_count.to_string().cyan(), detail.table_count.to_string().cyan()));
+            println!("{}{}", label("health:"), format_health(&health));
         }
         _ => unreachable!(),
     }
 }
 
+#[derive(Deserialize, Serialize)]
+struct CreateResponse {
+    id: String,
+    name: String,
+    source_type: String,
+    tables_discovered: u64,
+    discovery_status: String,
+    discovery_error: Option<String>,
+}
+
 pub fn create(
     workspace_id: &str,
     name: &str,
@@ -127,22 +228,53 @@ pub fn create(
     });
 
     let api = ApiClient::new(Some(workspace_id));
+    let is_table = format == "table";
+
+    let spinner = is_table.then(|| crate::util::spinner("Creating connection..."));
+    let (status, resp_body) = api.post_raw("/connections", &body);
+    if let Some(s) = &spinner { s.finish_and_clear(); }
 
-    #[derive(Deserialize, Serialize)]
-    struct CreateResponse {
-        id: String,
-        name: String,
-        source_type: String,
-        tables_discovered: u64,
-        discovery_status: String,
-        discovery_error: Option<String>,
+    if !status.is_success() {
+        use crossterm::style::Stylize;
+        eprintln!("{}", crate::util::api_error(resp_body).red());
+        std::process::exit(1);
     }
 
-    let result: CreateResponse = api.post("/connections", &body);
+    let result: CreateResponse = match serde_json::from_str(&resp_body) {
+        Ok(v) => v,
+        Err(e) => {
+            eprintln!("error parsing response: {e}");
+            std::process::exit(1);
+        }
+    };
+
+    let health = fetch_health(&api, &result.id, is_table);
 
     match format {
-        "json" => println!("{}", serde_json::to_string_pretty(&result).unwrap()),
-        "yaml" => print!("{}", serde_yaml::to_string(&result).unwrap()),
+        "json" => {
+            let combined = serde_json::json!({
+                "id": result.id,
+                "name": result.name,
+                "source_type": result.source_type,
+                "tables_discovered": result.tables_discovered,
+                "discovery_status": result.discovery_status,
+                "discovery_error": result.discovery_error,
+                "health": &health,
+            });
+            println!("{}", serde_json::to_string_pretty(&combined).unwrap());
+        }
+        "yaml" => {
+            let combined = serde_json::json!({
+                "id": result.id,
+                "name": result.name,
+                "source_type": result.source_type,
+                "tables_discovered": result.tables_discovered,
+                "discovery_status": result.discovery_status,
+                "discovery_error": result.discovery_error,
+                "health": &health,
+            });
+            print!("{}", serde_yaml::to_string(&combined).unwrap());
+        }
         "table" => {
             use crossterm::style::Stylize;
             println!("{}", "Connection created".green());
@@ -156,9 +288,14 @@ pub fn create(
                 _         => result.discovery_status.yellow().to_string(),
             };
             println!("discovery_status:  {status_colored}");
+            println!("health:            {}", format_health(&health));
         }
         _ => unreachable!(),
     }
+
+    if health.is_confirmed_unhealthy() {
+        std::process::exit(1);
+    }
 }
 
 pub fn list(workspace_id: &str, format: &str) {
diff --git a/src/connections_new.rs b/src/connections_new.rs
@@ -268,9 +268,51 @@ pub fn run(workspace_id: &str) {
         discovery_error: Option<String>,
     }
 
-    let result: CreateResponse = api.post("/connections", &body);
+    #[derive(serde::Deserialize)]
+    struct HealthResponse {
+        healthy: bool,
+        #[serde(default)]
+        latency_ms: Option<u64>,
+        #[serde(default)]
+        error: Option<String>,
+    }
+
+    enum HealthStatus {
+        Available(HealthResponse),
+        Unavailable(String),
+    }
+
+    let create_spinner = crate::util::spinner("Creating connection...");
+    let (status_code, resp_body) = api.post_raw("/connections", &body);
+    create_spinner.finish_and_clear();
 
     use crossterm::style::Stylize;
+    if !status_code.is_success() {
+        eprintln!("{}", crate::util::api_error(resp_body).red());
+        std::process::exit(1);
+    }
+
+    let result: CreateResponse = match serde_json::from_str(&resp_body) {
+        Ok(v) => v,
+        Err(e) => {
+            eprintln!("error parsing response: {e}");
+            std::process::exit(1);
+        }
+    };
+
+    let health_spinner = crate::util::spinner("Checking connection health...");
+    let (hstatus, hbody) = api.get_raw(&format!("/connections/{}/health", result.id));
+    health_spinner.finish_and_clear();
+
+    let health = if !hstatus.is_success() {
+        HealthStatus::Unavailable(crate::util::api_error(hbody))
+    } else {
+        match serde_json::from_str::<HealthResponse>(&hbody) {
+            Ok(h) => HealthStatus::Available(h),
+            Err(e) => HealthStatus::Unavailable(format!("parse error: {e}")),
+        }
+    };
+
     println!("{}", "Connection created".green());
     println!("id:                {}", result.id);
     println!("name:              {}", result.name);
@@ -282,4 +324,22 @@ pub fn run(workspace_id: &str) {
         _         => result.discovery_status.yellow().to_string(),
     };
     println!("discovery_status:  {status}");
+    let health_str = match &health {
+        HealthStatus::Available(h) if h.healthy => match h.latency_ms {
+            Some(ms) => format!("{} {}", "healthy".green(), format!("({ms}ms)").dark_grey()),
+            None => "healthy".green().to_string(),
+        },
+        HealthStatus::Available(h) => {
+            let err = h.error.as_deref().unwrap_or("unknown error");
+            format!("{} — {}", "unhealthy".red(), err)
+        }
+        HealthStatus::Unavailable(err) => {
+            format!("{} — {}", "unavailable".yellow(), err)
+        }
+    };
+    println!("health:            {health_str}");
+
+    if matches!(&health, HealthStatus::Available(h) if !h.healthy) {
+        std::process::exit(1);
+    }
 }
diff --git a/src/query.rs b/src/query.rs
@@ -57,13 +57,7 @@ pub fn execute(sql: &str, workspace_id: &str, connection: Option<&str>, format:
         body["connection_id"] = Value::String(conn.to_string());
     }
 
-    let spinner = indicatif::ProgressBar::new_spinner();
-    spinner.set_style(
-        indicatif::ProgressStyle::with_template("{spinner:.cyan} {msg}")
-            .unwrap(),
-    );
-    spinner.set_message("running query...");
-    spinner.enable_steady_tick(std::time::Duration::from_millis(80));
+    let spinner = crate::util::spinner("running query...");
 
     let (status, resp_body) = api.post_raw("/query", &body);
     spinner.finish_and_clear();
diff --git a/src/util.rs b/src/util.rs
