Generic retry and error parsing (#3558)

Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Michael Neale <michael.neale@gmail.com>
2025-08-05 10:37:04 +02:00
parent 33f57da4c9
commit 918faddf03
26 changed files with 1190 additions and 1748 deletions
@@ -0,0 +1,225 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use reqwest::{
+    header::{HeaderMap, HeaderName, HeaderValue},
+    Client, Response, StatusCode,
+};
+use serde_json::Value;
+use std::fmt;
+use std::time::Duration;
+
+pub struct ApiClient {
+    client: Client,
+    host: String,
+    auth: AuthMethod,
+    default_headers: HeaderMap,
+    timeout: Duration,
+}
+
+pub enum AuthMethod {
+    BearerToken(String),
+    ApiKey {
+        header_name: String,
+        key: String,
+    },
+    #[allow(dead_code)]
+    OAuth(OAuthConfig),
+    Custom(Box<dyn AuthProvider>),
+}
+
+pub struct OAuthConfig {
+    pub host: String,
+    pub client_id: String,
+    pub redirect_url: String,
+    pub scopes: Vec<String>,
+}
+
+#[async_trait]
+pub trait AuthProvider: Send + Sync {
+    async fn get_auth_header(&self) -> Result<(String, String)>;
+}
+
+pub struct ApiResponse {
+    pub status: StatusCode,
+    pub payload: Option<Value>,
+}
+
+impl fmt::Debug for AuthMethod {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            AuthMethod::BearerToken(_) => f.debug_tuple("BearerToken").field(&"[hidden]").finish(),
+            AuthMethod::ApiKey { header_name, .. } => f
+                .debug_struct("ApiKey")
+                .field("header_name", header_name)
+                .field("key", &"[hidden]")
+                .finish(),
+            AuthMethod::OAuth(_) => f.debug_tuple("OAuth").field(&"[config]").finish(),
+            AuthMethod::Custom(_) => f.debug_tuple("Custom").field(&"[provider]").finish(),
+        }
+    }
+}
+
+impl ApiResponse {
+    pub async fn from_response(response: Response) -> Result<Self> {
+        let status = response.status();
+        let payload = response.json().await.ok();
+        Ok(Self { status, payload })
+    }
+}
+
+pub struct ApiRequestBuilder<'a> {
+    client: &'a ApiClient,
+    path: &'a str,
+    headers: HeaderMap,
+}
+
+impl ApiClient {
+    pub fn new(host: String, auth: AuthMethod) -> Result<Self> {
+        Self::with_timeout(host, auth, Duration::from_secs(600))
+    }
+
+    pub fn with_timeout(host: String, auth: AuthMethod, timeout: Duration) -> Result<Self> {
+        Ok(Self {
+            client: Client::builder().timeout(timeout).build()?,
+            host,
+            auth,
+            default_headers: HeaderMap::new(),
+            timeout,
+        })
+    }
+
+    pub fn with_headers(mut self, headers: HeaderMap) -> Result<Self> {
+        self.default_headers = headers;
+        self.client = Client::builder()
+            .timeout(self.timeout)
+            .default_headers(self.default_headers.clone())
+            .build()?;
+        Ok(self)
+    }
+
+    pub fn with_header(mut self, key: &str, value: &str) -> Result<Self> {
+        let header_name = HeaderName::from_bytes(key.as_bytes())?;
+        let header_value = HeaderValue::from_str(value)?;
+        self.default_headers.insert(header_name, header_value);
+        self.client = Client::builder()
+            .timeout(self.timeout)
+            .default_headers(self.default_headers.clone())
+            .build()?;
+        Ok(self)
+    }
+
+    pub fn request<'a>(&'a self, path: &'a str) -> ApiRequestBuilder<'a> {
+        ApiRequestBuilder {
+            client: self,
+            path,
+            headers: HeaderMap::new(),
+        }
+    }
+
+    pub async fn api_post(&self, path: &str, payload: &Value) -> Result<ApiResponse> {
+        self.request(path).api_post(payload).await
+    }
+
+    pub async fn response_post(&self, path: &str, payload: &Value) -> Result<Response> {
+        self.request(path).response_post(payload).await
+    }
+
+    pub async fn api_get(&self, path: &str) -> Result<ApiResponse> {
+        self.request(path).api_get().await
+    }
+
+    pub async fn response_get(&self, path: &str) -> Result<Response> {
+        self.request(path).response_get().await
+    }
+
+    fn build_url(&self, path: &str) -> Result<url::Url> {
+        use url::Url;
+        let base_url =
+            Url::parse(&self.host).map_err(|e| anyhow::anyhow!("Invalid base URL: {}", e))?;
+        base_url
+            .join(path)
+            .map_err(|e| anyhow::anyhow!("Failed to construct URL: {}", e))
+    }
+
+    async fn get_oauth_token(&self, config: &OAuthConfig) -> Result<String> {
+        super::oauth::get_oauth_token_async(
+            &config.host,
+            &config.client_id,
+            &config.redirect_url,
+            &config.scopes,
+        )
+        .await
+    }
+}
+
+impl<'a> ApiRequestBuilder<'a> {
+    pub fn header(mut self, key: &str, value: &str) -> Result<Self> {
+        let header_name = HeaderName::from_bytes(key.as_bytes())?;
+        let header_value = HeaderValue::from_str(value)?;
+        self.headers.insert(header_name, header_value);
+        Ok(self)
+    }
+
+    #[allow(dead_code)]
+    pub fn headers(mut self, headers: HeaderMap) -> Self {
+        self.headers.extend(headers);
+        self
+    }
+
+    pub async fn api_post(self, payload: &Value) -> Result<ApiResponse> {
+        let response = self.response_post(payload).await?;
+        ApiResponse::from_response(response).await
+    }
+
+    pub async fn response_post(self, payload: &Value) -> Result<Response> {
+        let request = self.send_request(|url, client| client.post(url)).await?;
+        Ok(request.json(payload).send().await?)
+    }
+
+    pub async fn api_get(self) -> Result<ApiResponse> {
+        let response = self.response_get().await?;
+        ApiResponse::from_response(response).await
+    }
+
+    pub async fn response_get(self) -> Result<Response> {
+        let request = self.send_request(|url, client| client.get(url)).await?;
+        Ok(request.send().await?)
+    }
+
+    async fn send_request<F>(&self, request_builder: F) -> Result<reqwest::RequestBuilder>
+    where
+        F: FnOnce(url::Url, &Client) -> reqwest::RequestBuilder,
+    {
+        let url = self.client.build_url(self.path)?;
+        let mut request = request_builder(url, &self.client.client);
+        request = request.headers(self.headers.clone());
+
+        request = match &self.client.auth {
+            AuthMethod::BearerToken(token) => {
+                request.header("Authorization", format!("Bearer {}", token))
+            }
+            AuthMethod::ApiKey { header_name, key } => request.header(header_name.as_str(), key),
+            AuthMethod::OAuth(config) => {
+                let token = self.client.get_oauth_token(config).await?;
+                request.header("Authorization", format!("Bearer {}", token))
+            }
+            AuthMethod::Custom(provider) => {
+                let (header_name, header_value) = provider.get_auth_header().await?;
+                request.header(header_name, header_value)
+            }
+        };
+
+        Ok(request)
+    }
+}
+
+impl fmt::Debug for ApiClient {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ApiClient")
+            .field("host", &self.host)
+            .field("auth", &"[auth method]")
+            .field("timeout", &self.timeout)
+            .field("default_headers", &self.default_headers)
+            .finish_non_exhaustive()
+    }
+}