feat(api-client): always catch 4xx and 5xx errors regardless of body type (#2052)

* feat(api-client): always catch 4xx and 5xx errors regardless of body type

* clip it

Author: jonaro00

Cargo.lock

@@ -28,6 +28,7 @@ assert_cmd = "2.0.6"
 async-trait = "0.1.58"
 axum = { version = "0.8.1", default-features = false }
 bollard = { version = "0.18.1", features = ["ssl_providerless"] }
+bytes = "1"
 cargo_metadata = "0.19.1"
 chrono = { version = "0.4.34", default-features = false }
 clap = { version = "4.2.7", features = ["derive"] }

Cargo.toml

@@ -1,6 +1,6 @@
 use anyhow::Result;
 use serde_json::{json, Value};
-use shuttle_api_client::ShuttleApiClient;
+use shuttle_api_client::{util::ToBodyContent, ShuttleApiClient};
 use shuttle_common::models::{
     project::{ProjectResponse, ProjectUpdateRequest},
     team::AddTeamMemberRequest,
@@ -91,44 +91,39 @@ impl Client {
             .await
     }
 
-    pub async fn add_team_member(&self, team_user_id: &str, user_id: String) -> Result<()> {
+    pub async fn add_team_member(&self, team_user_id: &str, user_id: String) -> Result<String> {
         self.inner
-            .post(
+            .post_json(
                 format!("/teams/{team_user_id}/members"),
                 Some(AddTeamMemberRequest {
                     user_id: Some(user_id),
                     email: None,
                     role: None,
                 }),
             )
-            .await?;
-
-        Ok(())
+            .await
     }
 
     pub async fn feature_flag(&self, entity: &str, flag: &str, set: bool) -> Result<()> {
-        let resp = if set {
+        if set {
             self.inner
                 .put(
                     format!("/admin/feature-flag/{entity}/{flag}"),
                     Option::<()>::None,
                 )
                 .await?
+                .to_empty()
+                .await
         } else {
             self.inner
                 .delete(
                     format!("/admin/feature-flag/{entity}/{flag}"),
                     Option::<()>::None,
                 )
                 .await?
-        };
-
-        if !resp.status().is_success() {
-            dbg!(resp);
-            panic!("request failed");
+                .to_empty()
+                .await
         }
-
-        Ok(())
     }
 
     pub async fn gc_free_tier(&self, days: u32) -> Result<Vec<String>> {
@@ -163,9 +158,9 @@ impl Client {
                 format!("/admin/users/{user_id}/account_tier"),
                 Some(UpdateAccountTierRequest { account_tier }),
             )
-            .await?;
-
-        Ok(())
+            .await?
+            .to_empty()
+            .await
     }
 
     pub async fn get_expired_protrials(&self) -> Result<Vec<String>> {

admin/src/client.rs

@@ -12,6 +12,7 @@ shuttle-common = { workspace = true, features = ["models", "unknown-variants"] }
 
 anyhow = { workspace = true }
 async-trait = { workspace = true }
+bytes = { workspace = true }
 headers = { workspace = true }
 http = { workspace = true }
 percent-encoding = { workspace = true }

api-client/Cargo.toml

@@ -32,8 +32,8 @@ use crate::middleware::LoggingMiddleware;
 #[cfg(feature = "tracing")]
 use tracing::{debug, error};
 
-mod util;
-use util::ToJson;
+pub mod util;
+use util::ToBodyContent;
 
 #[derive(Clone)]
 pub struct ShuttleApiClient {
@@ -155,19 +155,20 @@ impl ShuttleApiClient {
         let r#type = resource_type.to_string();
         let r#type = utf8_percent_encode(&r#type, percent_encoding::NON_ALPHANUMERIC).to_owned();
 
-        let res = self
+        let bytes = self
             .get(
                 format!(
                     "/projects/{project}/services/{project}/resources/{}/dump",
                     r#type
                 ),
                 Option::<()>::None,
             )
-            .await?;
+            .await?
+            .to_bytes()
+            .await?
+            .to_vec();
 
-        let bytes = res.bytes().await?;
-
-        Ok(bytes.to_vec())
+        Ok(bytes)
     }
 
     pub async fn delete_service_resource(
@@ -301,15 +302,22 @@ impl ShuttleApiClient {
         self.get_json(path).await
     }
 
-    pub async fn reset_api_key(&self) -> Result<Response> {
-        self.put("/users/reset-api-key", Option::<()>::None).await
+    pub async fn reset_api_key(&self) -> Result<()> {
+        self.put("/users/reset-api-key", Option::<()>::None)
+            .await?
+            .to_empty()
+            .await
     }
 
     pub async fn ws_get(
         &self,
         path: impl AsRef<str>,
     ) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
-        let ws_url = self.api_url.clone().replace("http", "ws");
+        let ws_url = self
+            .api_url
+            .clone()
+            .replacen("http://", "ws://", 1)
+            .replacen("https://", "wss://", 1);
         let url = format!("{ws_url}{}", path.as_ref());
         let mut req = url.into_client_request()?;
 

api-client/src/lib.rs

@@ -1,47 +1,95 @@
 use anyhow::{Context, Result};
 use async_trait::async_trait;
+use bytes::Bytes;
 use http::StatusCode;
 use serde::de::DeserializeOwned;
 use shuttle_common::models::error::ApiError;
 
-/// A to_json wrapper for handling our error states
+/// Helpers for consuming and parsing response bodies and handling parsing of an ApiError if the response is 4xx/5xx
 #[async_trait]
-pub trait ToJson {
+pub trait ToBodyContent {
     async fn to_json<T: DeserializeOwned>(self) -> Result<T>;
+    async fn to_text(self) -> Result<String>;
+    async fn to_bytes(self) -> Result<Bytes>;
+    async fn to_empty(self) -> Result<()>;
+}
+
+fn into_api_error(body: &str, status_code: StatusCode) -> ApiError {
+    #[cfg(feature = "tracing")]
+    tracing::trace!("Parsing response as API error");
+
+    let res: ApiError = match serde_json::from_str(body) {
+        Ok(res) => res,
+        _ => ApiError::new(
+            format!("Failed to parse error response from the server:\n{}", body),
+            status_code,
+        ),
+    };
+
+    res
+}
+
+/// Tries to convert bytes to string. If not possible, returns a string symbolizing the bytes and the length
+fn bytes_to_string_with_fallback(bytes: Bytes) -> String {
+    String::from_utf8(bytes.to_vec()).unwrap_or_else(|_| format!("[{} bytes]", bytes.len()))
 }
 
 #[async_trait]
-impl ToJson for reqwest::Response {
+impl ToBodyContent for reqwest::Response {
     async fn to_json<T: DeserializeOwned>(self) -> Result<T> {
         let status_code = self.status();
         let bytes = self.bytes().await?;
-        let string = String::from_utf8(bytes.to_vec())
-            .unwrap_or_else(|_| format!("[{} bytes]", bytes.len()));
+        let string = bytes_to_string_with_fallback(bytes);
 
         #[cfg(feature = "tracing")]
         tracing::trace!(response = %string, "Parsing response as JSON");
 
-        if matches!(
-            status_code,
-            StatusCode::OK | StatusCode::SWITCHING_PROTOCOLS
-        ) {
-            serde_json::from_str(&string).context("failed to parse a successful response")
-        } else {
-            #[cfg(feature = "tracing")]
-            tracing::trace!("Parsing response as API error");
-
-            let res: ApiError = match serde_json::from_str(&string) {
-                Ok(res) => res,
-                _ => ApiError::new(
-                    format!(
-                        "Failed to parse error response from the server:\n{}",
-                        string
-                    ),
-                    status_code,
-                ),
-            };
-
-            Err(res.into())
+        if status_code.is_client_error() || status_code.is_server_error() {
+            return Err(into_api_error(&string, status_code).into());
+        }
+
+        serde_json::from_str(&string).context("failed to parse a successful response")
+    }
+
+    async fn to_text(self) -> Result<String> {
+        let status_code = self.status();
+        let bytes = self.bytes().await?;
+        let string = bytes_to_string_with_fallback(bytes);
+
+        #[cfg(feature = "tracing")]
+        tracing::trace!(response = %string, "Parsing response as text");
+
+        if status_code.is_client_error() || status_code.is_server_error() {
+            return Err(into_api_error(&string, status_code).into());
         }
+
+        Ok(string)
+    }
+
+    async fn to_bytes(self) -> Result<Bytes> {
+        let status_code = self.status();
+        let bytes = self.bytes().await?;
+
+        #[cfg(feature = "tracing")]
+        tracing::trace!(response_length = bytes.len(), "Got response bytes");
+
+        if status_code.is_client_error() || status_code.is_server_error() {
+            let string = bytes_to_string_with_fallback(bytes);
+            return Err(into_api_error(&string, status_code).into());
+        }
+
+        Ok(bytes)
+    }
+
+    async fn to_empty(self) -> Result<()> {
+        let status_code = self.status();
+
+        if status_code.is_client_error() || status_code.is_server_error() {
+            let bytes = self.bytes().await?;
+            let string = bytes_to_string_with_fallback(bytes);
+            return Err(into_api_error(&string, status_code).into());
+        }
+
+        Ok(())
     }
 }

api-client/src/util.rs

@@ -862,7 +862,8 @@ impl Shuttle {
 
     async fn logout(&mut self, logout_args: LogoutArgs) -> Result<()> {
         if logout_args.reset_api_key {
-            self.reset_api_key().await?;
+            let client = self.client.as_ref().unwrap();
+            client.reset_api_key().await.context("Resetting API key")?;
             eprintln!("Successfully reset the API key.");
         }
         self.ctx.clear_api_key()?;
@@ -872,17 +873,6 @@ impl Shuttle {
         Ok(())
     }
 
-    async fn reset_api_key(&self) -> Result<()> {
-        let client = self.client.as_ref().unwrap();
-        client.reset_api_key().await.and_then(|res| {
-            if res.status().is_success() {
-                Ok(())
-            } else {
-                Err(anyhow!("Resetting API key failed."))
-            }
-        })
-    }
-
     async fn stop(&self, tracking_args: DeploymentTrackingArgs) -> Result<()> {
         let client = self.client.as_ref().unwrap();
         let pid = self.ctx.project_id();

cargo-shuttle/src/lib.rs

@@ -12,10 +12,14 @@
 
 ## Usage
 
-Start by installing the Shuttle CLI by running the following in a terminal:
+Start by installing the [Shuttle CLI](https://crates.io/crates/cargo-shuttle) by running the following in a terminal ([more installation options](https://docs.shuttle.dev/getting-started/installation)):
 
 ```bash
-cargo install cargo-shuttle
+# Linux / macOS
+curl -sSfL https://www.shuttle.dev/install | bash
+
+# Windows (Powershell)
+iwr https://www.shuttle.dev/install-win | iex
 ```
 
 Now that Shuttle is installed, you can initialize a project with Axum boilerplate: