windmill/backend/windmill-common/src/client.rs

use anyhow::Context;
use reqwest::{Body, Response};
use serde::de::DeserializeOwned;

use crate::utils::{HTTP_CLIENT, HTTP_CLIENT_STREAMING};

#[derive(Clone)]
pub struct AuthedClient {
    pub base_internal_url: String,
    pub workspace: String,
    pub token: String,
    pub force_client: Option<reqwest::Client>,
}

impl AuthedClient {
    pub fn new(
        base_internal_url: String,
        workspace: String,
        token: String,
        force_client: Option<reqwest::Client>,
    ) -> AuthedClient {
        AuthedClient { base_internal_url, workspace, token, force_client }
    }

    pub async fn get(&self, url: &str, query: Vec<(&str, String)>) -> anyhow::Result<Response> {
        self.force_client
            .as_ref()
            .unwrap_or(&HTTP_CLIENT)
            .get(url)
            .query(&query)
            .header(
                reqwest::header::ACCEPT,
                reqwest::header::HeaderValue::from_static("application/json"),
            )
            .header(
                reqwest::header::AUTHORIZATION,
                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))?,
            )
            .send()
            .await
            .map_err(|e| {
                tracing::error!("Error executing get request from authed http client to {url} with query {query:?}: {e:#?}");
                anyhow::anyhow!("Error executing get request from authed http client to {url} with query {query:?}: {e:#?}")
            })
    }

    /// Like [`AuthedClient::get`], but for a response whose size decides how
    /// long it takes. `HTTP_CLIENT`'s total timeout would cut off a large one
    /// partway through, so this uses the streaming client, which bounds only
    /// the connect.
    pub async fn get_streaming(
        &self,
        url: &str,
        query: Vec<(&str, String)>,
    ) -> anyhow::Result<Response> {
        self.force_client
            .as_ref()
            .unwrap_or(&HTTP_CLIENT_STREAMING)
            .get(url)
            .query(&query)
            .header(
                reqwest::header::AUTHORIZATION,
                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))?,
            )
            .send()
            .await
            .map_err(|e| {
                tracing::error!("Error streaming get request from authed http client to {url} with query {query:?}: {e:#?}");
                anyhow::anyhow!("Error streaming get request from authed http client to {url} with query {query:?}: {e:#?}")
            })
    }

    pub async fn get_id_token(&self, audience: &str) -> anyhow::Result<String> {
        let url = format!(
            "{}/api/w/{}/oidc/token/{}",
            self.base_internal_url, self.workspace, audience
        );
        let response = self
            .force_client
            .as_ref()
            .unwrap_or(&HTTP_CLIENT)
            .post(&url)
            .header(
                reqwest::header::AUTHORIZATION,
                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))?,
            )
            .send()
            .await
            .map_err(|e| {
                tracing::error!("Error requesting oidc token from {url}: {e:#?}");
                anyhow::anyhow!("Error requesting oidc token from {url}: {e:#?}")
            })?;

        match response.status().as_u16() {
            200u16 => Ok(response.text().await.context("reading oidc token body")?),
            status => {
                let body = response.text().await.unwrap_or_default();
                Err(anyhow::anyhow!(
                    "oidc token request to {url} failed with status {status}: {body}"
                ))
            }
        }
    }

    pub async fn get_resource_value<T: DeserializeOwned>(&self, path: &str) -> anyhow::Result<T> {
        let url = format!(
            "{}/api/w/{}/resources/get_value/{}",
            self.base_internal_url, self.workspace, path
        );
        make_basic_get_request(self, &url, None, Some("decoding resource value as json")).await
    }

    pub async fn get_variable_value(&self, path: &str) -> anyhow::Result<String> {
        let url = format!(
            "{}/api/w/{}/variables/get_value/{}",
            self.base_internal_url, self.workspace, path
        );
        make_basic_get_request(self, &url, None, Some("decoding variable value as json")).await
    }

    pub async fn get_resource_value_interpolated<T: DeserializeOwned>(
        &self,
        path: &str,
        job_id: Option<String>,
    ) -> anyhow::Result<T> {
        let url = format!(
            "{}/api/w/{}/resources/get_value_interpolated/{}",
            self.base_internal_url, self.workspace, path
        );
        let mut query = Vec::with_capacity(1usize);
        if let Some(v) = &job_id {
            query.push(("job_id", v.to_string()));
        }
        let response = self.get(&url, query).await?;
        match response.status().as_u16() {
            200u16 => Ok(response
                .json::<T>()
                .await
                .context("decoding interpolated resource value as json")?),
            _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())),
        }
    }

    /// Whether the workspace configures this dbt warehouse. Resolves nothing.
    pub async fn dbt_warehouse_exists(&self, name: &str) -> anyhow::Result<()> {
        let url = format!(
            "{}/api/w/{}/dbt/warehouse_exists/{}",
            self.base_internal_url, self.workspace, name
        );
        let response = self.get(&url, vec![]).await?;
        match response.status().as_u16() {
            200u16 => Ok(()),
            _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())),
        }
    }

    /// Record a run's settled dbt nodes, for a worker with no database.
    ///
    /// Posted with the JOB's token, not the agent's: the route takes the job
    /// from the token, and the agent's own credential only authenticates
    /// against the agent surface.
    pub async fn record_dbt_run_progress(
        &self,
        req: &[crate::dbt_manifest::DbtRunProgressRequest],
    ) -> anyhow::Result<()> {
        let url = format!(
            "{}/api/w/{}/dbt/run_progress",
            self.base_internal_url, self.workspace
        );
        let response = self
            .force_client
            .as_ref()
            .unwrap_or(&HTTP_CLIENT)
            .post(&url)
            .header(
                reqwest::header::AUTHORIZATION,
                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))?,
            )
            .json(req)
            .send()
            .await?;
        match response.status().as_u16() {
            200u16 => Ok(()),
            _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())),
        }
    }

    /// Where a dbt warehouse name points, for a worker with no database.
    pub async fn get_dbt_warehouse<T: DeserializeOwned>(&self, name: &str) -> anyhow::Result<T> {
        let url = format!(
            "{}/api/w/{}/dbt/warehouse/{}",
            self.base_internal_url, self.workspace, name
        );
        make_basic_get_request(self, &url, None, Some("decoding the dbt warehouse ref")).await
    }

    pub async fn get_completed_job_result<T: DeserializeOwned>(
        &self,
        path: &str,
        json_path: Option<String>,
    ) -> anyhow::Result<T> {
        let url = format!(
            "{}/api/w/{}/jobs_u/completed/get_result/{}",
            self.base_internal_url, self.workspace, path
        );
        let query = query_from_json_path(json_path);
        make_basic_get_request(
            self,
            &url,
            Some(query),
            Some("decoding completed job result as json"),
        )
        .await
    }

    pub async fn get_result_by_id<T: DeserializeOwned>(
        &self,
        flow_job_id: &str,
        node_id: &str,
        json_path: Option<String>,
    ) -> anyhow::Result<T> {
        let url = format!(
            "{}/api/w/{}/jobs/result_by_id/{}/{}",
            self.base_internal_url, self.workspace, flow_job_id, node_id
        );
        let query = query_from_json_path(json_path);
        make_basic_get_request(
            self,
            &url,
            Some(query),
            Some("decoding result by id as json"),
        )
        .await
    }

    pub async fn get_flow_user_state(
        &self,
        job_id: &str,
        key: &str,
    ) -> anyhow::Result<serde_json::Value> {
        let url = format!(
            "{}/api/w/{}/jobs/flow/user_states/{}/{}",
            self.base_internal_url,
            self.workspace,
            job_id,
            urlencoding::encode(key),
        );
        make_basic_get_request(self, &url, None, Some("decoding flow user state as json")).await
    }

    pub async fn upload_s3_file<S>(
        &self,
        workspace_id: &str,
        object_key: String,
        storage: Option<String>,
        body: S,
    ) -> anyhow::Result<()>
    where
        S: futures::stream::TryStream + Send + 'static,
        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        bytes::Bytes: From<S::Ok>,
    {
        let mut query = vec![("file_key", object_key)];
        if let Some(storage) = storage {
            query.push(("storage", storage));
        }
        let url = format!(
            "{}/api/w/{}/job_helpers/upload_s3_file",
            self.base_internal_url, workspace_id
        );
        // Use the streaming HTTP client (no total request timeout) because the
        // upload body is streamed and total time depends on data size.
        let response = self
            .force_client
            .as_ref()
            .unwrap_or(&HTTP_CLIENT_STREAMING)
            .post(&url)
            .query(&query)
            .header(
                reqwest::header::ACCEPT,
                reqwest::header::HeaderValue::from_static("application/json"),
            )
            .header(
                reqwest::header::AUTHORIZATION,
                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))
                    .map_err(|e| anyhow::anyhow!(e.to_string()))?,
            )
            .body(Body::wrap_stream(body))
            .send()
            .await
            .context(format!("Failed to send upload_s3_file request to {url}"))?;

        match response.status().as_u16() {
            200u16 => Ok(()),
            _ => {
                let status = response.status();
                let body = response.text().await.unwrap_or_default();
                Err(anyhow::anyhow!(
                    "upload_s3_file request to {url} failed with status {status}: {body}"
                ))
            }
        }
    }

    pub async fn download_s3_file(
        &self,
        workspace_id: &str,
        file_key: &str,
        storage: Option<String>,
    ) -> anyhow::Result<bytes::Bytes> {
        let mut query = vec![("file_key", file_key.to_string())];
        if let Some(storage) = storage {
            query.push(("storage", storage));
        }
        let response = self
            .force_client
            .as_ref()
            .unwrap_or(&HTTP_CLIENT)
            .get(&format!(
                "{}/api/w/{}/job_helpers/download_s3_file",
                self.base_internal_url, workspace_id
            ))
            .query(&query)
            .header(
                reqwest::header::AUTHORIZATION,
                reqwest::header::HeaderValue::from_str(&format!("Bearer {}", self.token))
                    .map_err(|e| anyhow::anyhow!(e.to_string()))?,
            )
            .send()
            .await
            .context("Failed to send download_s3_file request")
            .map_err(|e| anyhow::anyhow!(e.to_string()))?;

        match response.status().as_u16() {
            200u16 => Ok(response
                .bytes()
                .await
                .context("Failed to read response bytes")?),
            _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())),
        }
    }
}

#[inline]
fn query_from_json_path(json_path: Option<String>) -> Vec<(&'static str, String)> {
    json_path
        .map(|json_path| vec![("json_path", json_path)])
        .unwrap_or_else(|| Vec::new())
}

#[inline]
async fn make_basic_get_request<T: DeserializeOwned>(
    client: &AuthedClient,
    url: &str,
    query: Option<Vec<(&'static str, String)>>,
    context: Option<&'static str>,
) -> anyhow::Result<T> {
    let response = client
        .get(&url, query.unwrap_or_else(|| Vec::new()))
        .await?;

    match response.status().as_u16() {
        200u16 => {
            let json_body = response
                .json::<T>()
                .await
                .context(context.unwrap_or("error decoding body as json"))?;
            Ok(json_body)
        }
        _ => Err(anyhow::anyhow!(response.text().await.unwrap_or_default())),
    }
}