trs-mlflow 0.7.0

//! This module contains some http related stuff which we are using
//! in a client.
//! use anyhow::
use crate::error::ApiError;
use anyhow::{anyhow, Context, Result};
use reqwest::StatusCode;
use serde::de::DeserializeOwned;
use serde_json::from_str;
use std::any::type_name;

/// Wrapper around reqwest::Response that is used to provided better
/// error handling of response errors from mlflow.
pub struct Response(pub reqwest::Response);

impl Response {
    /// Tries to parse json and returns error with WHOLE body if
    /// it cannot do that.
    pub async fn json<T: DeserializeOwned>(self) -> Result<T> {
        let body = self.0.text().await.context("Cannot load body as a text")?;

        from_str(&body).map_err(|error| {
            anyhow!(error).context(format!("Cannot parse {} from {}", type_name::<T>(), body))
        })
    }

    /// Tries to parse error response in form of json (or plain text)
    /// from mlflow response if status is not OK.
    pub async fn error_for_status(self) -> Result<Self> {
        let Err(error) = self.0.error_for_status_ref() else {
            return Ok(self);
        };

        let status = self.0.status();

        // we try to load information about error from body
        let body = self
            .0
            .text()
            .await
            .context("Cannot load error response from server")?;

        // Mlflow server doesn't return ApiError if we are not authorized
        // so we need to handle it manually
        let api_error = if status == StatusCode::UNAUTHORIZED {
            ApiError {
                error_code: "UNAUTHORIZED".to_owned(),
                message: body,
            }
        } else {
            from_str::<ApiError>(&body)
                .with_context(|| format!("Cannot parse ApiError from {}", body))?
        };

        Err(anyhow!(error).context(api_error))
    }
}