From 36db047a869bc4b06fc2e3fe8bf9c5d03bd20312 Mon Sep 17 00:00:00 2001 From: Guillaume Bouvignies Date: Mon, 4 Dec 2023 13:19:40 +0100 Subject: [PATCH] feat: Custom content type for script and flow results (#2767) --- backend/windmill-api/src/jobs.rs | 68 ++++++++++++++++++++++------ backend/windmill-common/src/error.rs | 56 +++++++++-------------- 2 files changed, 76 insertions(+), 48 deletions(-) diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index c498df065d..422095b6dd 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -6,6 +6,7 @@ * LICENSE-AGPL for a copy of the license. */ +use axum::http::HeaderValue; use serde_json::value::RawValue; use std::collections::HashMap; use windmill_common::flow_status::RestartedFrom; @@ -1874,15 +1875,16 @@ impl Drop for Guard { } #[derive(Deserialize)] -pub struct WindmillStatusCode { +pub struct WindmillCompositeResult { windmill_status_code: Option, + windmill_content_type: Option, result: Option>, } async fn run_wait_result( db: &DB, uuid: Uuid, Path((w_id, _)): Path<(String, T)>, - node_id: Option, + node_id_for_empty_return: Option, ) -> error::Result { let mut result; let timeout = SERVER_CONFIG.read().await.timeout_wait_result.clone(); @@ -1898,10 +1900,16 @@ async fn run_wait_result( let mut accumulated_delay = 0 as u64; loop { - if let Some(node_id) = node_id.as_ref() { - result = get_result_by_id_from_running_flow(&db, &w_id, &uuid, node_id, None) - .await - .ok(); + if let Some(node_id_for_empty_return) = node_id_for_empty_return.as_ref() { + result = get_result_by_id_from_running_flow( + &db, + &w_id, + &uuid, + node_id_for_empty_return, + None, + ) + .await + .ok(); } else { let row = sqlx::query("SELECT result FROM completed_job WHERE id = $1 AND workspace_id = $2") @@ -1934,13 +1942,47 @@ async fn run_wait_result( if let Some(result) = result { g.done = true; - let status_code = serde_json::from_str::(result.get()); - match status_code { - Ok(WindmillStatusCode { windmill_status_code: Some(status_code), result }) => { - Err(Error::CustomStatusCode( - StatusCode::from_u16(status_code).unwrap_or_else(|_| StatusCode::IM_A_TEAPOT), - result, - )) + let composite_result = serde_json::from_str::(result.get()); + match composite_result { + Ok(WindmillCompositeResult { + windmill_status_code, + windmill_content_type, + result: result_value, + }) => { + if windmill_content_type.is_none() && windmill_status_code.is_none() { + return Ok(Json(result).into_response()); + } + + let status_code_or_default = windmill_status_code + .map(|val| match StatusCode::from_u16(val) { + Ok(sc) => Ok(sc), + Err(_) => Err(Error::ExecutionErr("Invalid status code".to_string())), + }) + .unwrap_or(if result_value.is_some() { + Ok(StatusCode::OK) + } else { + Ok(StatusCode::NO_CONTENT) + })?; + + if windmill_content_type.is_some() { + let serialized_result = result_value + .map(|val| val.get().to_owned()) + .unwrap_or_else(String::new); + return Ok(( + status_code_or_default, + [( + http::header::CONTENT_TYPE, + HeaderValue::from_str(windmill_content_type.unwrap().as_str()).unwrap(), + )], + serialized_result, + ) + .into_response()); + } + return Ok(( + status_code_or_default, + Json(result_value), // default to JSON result if no content type is provided + ) + .into_response()); } _ => Ok(Json(result).into_response()), } diff --git a/backend/windmill-common/src/error.rs b/backend/windmill-common/src/error.rs index a17a47c5e7..394bd8c23a 100644 --- a/backend/windmill-common/src/error.rs +++ b/backend/windmill-common/src/error.rs @@ -15,7 +15,6 @@ use axum::{ }; use hyper::StatusCode; -use serde_json::value::RawValue; #[cfg(feature = "sqlx")] use sqlx::migrate::MigrateError; use thiserror::Error; @@ -65,8 +64,6 @@ pub enum Error { Anyhow(#[from] anyhow::Error), #[error("Error: {0:#?}")] JsonErr(serde_json::Value), - #[error("Custom Status Code: {0:#?}")] - CustomStatusCode(StatusCode, Option>), #[error("{0}")] OpenAIError(String), } @@ -85,41 +82,30 @@ pub fn to_anyhow(e: T) -> anyhow:: #[cfg(feature = "axum")] impl IntoResponse for Error { fn into_response(self) -> axum::response::Response { - match self { - Self::CustomStatusCode(code, result) => { - let mut res = Json(result).into_response(); - let status_mut = res.status_mut(); - *status_mut = code; - res + let e = &self; + let body = body::boxed(body::Full::from(e.to_string())); + + let status = match self { + Self::NotFound(_) => axum::http::StatusCode::NOT_FOUND, + Self::NotAuthorized(_) => axum::http::StatusCode::UNAUTHORIZED, + Self::RequireAdmin(_) => axum::http::StatusCode::FORBIDDEN, + Self::SqlErr(_) | Self::BadRequest(_) | Self::OpenAIError(_) => { + axum::http::StatusCode::BAD_REQUEST } - _ => { - let e = &self; - let body = body::boxed(body::Full::from(e.to_string())); + _ => axum::http::StatusCode::INTERNAL_SERVER_ERROR, + }; - let status = match self { - Self::NotFound(_) => axum::http::StatusCode::NOT_FOUND, - Self::NotAuthorized(_) => axum::http::StatusCode::UNAUTHORIZED, - Self::RequireAdmin(_) => axum::http::StatusCode::FORBIDDEN, - Self::CustomStatusCode(code, _) => code, - Self::SqlErr(_) | Self::BadRequest(_) | Self::OpenAIError(_) => { - axum::http::StatusCode::BAD_REQUEST - } - _ => axum::http::StatusCode::INTERNAL_SERVER_ERROR, - }; + if matches!(status, axum::http::StatusCode::NOT_FOUND) { + tracing::warn!(not_found = e.to_string()); + } else { + tracing::error!(error = e.to_string()); + }; - if matches!(status, axum::http::StatusCode::NOT_FOUND) { - tracing::warn!(not_found = e.to_string()); - } else { - tracing::error!(error = e.to_string()); - }; - - axum::response::Response::builder() - .header("Content-Type", "text/plain") - .status(status) - .body(body) - .unwrap() - } - } + axum::response::Response::builder() + .header("Content-Type", "text/plain") + .status(status) + .body(body) + .unwrap() } }