refactor sql result formatting (#4760)

This commit is contained in:
HugoCasa
2024-11-20 14:53:39 +01:00
committed by GitHub
parent f100f68b2f
commit bb185e7318
5 changed files with 186 additions and 101 deletions
+5
View File
@@ -0,0 +1,5 @@
INSERT INTO public.completed_job (
id, workspace_id, created_by, created_at, duration_ms, success, flow_status, result, job_kind, language
) VALUES (
'1eecb96a-c8b0-4a3d-b1b6-087878c55e41', 'test-workspace', 'test-user', '2023-01-01 00:00:00', 1000, true, '{"_metadata": {"column_order": ["b", "a"]}}', '[{"a": "second", "b": "first"}]', 'script', 'postgresql'
)
+69
View File
@@ -1,3 +1,4 @@
use serde::de::DeserializeOwned;
use std::str::FromStr;
use windmill_api_client::types::{NewScript, NewScriptLanguage};
@@ -3612,3 +3613,71 @@ def main():
run_deployed_relative_imports(&db, content.clone(), ScriptLang::Python3).await;
run_preview_relative_imports(&db, content, ScriptLang::Python3).await;
}
#[sqlx::test(fixtures("base", "result_format"))]
async fn test_result_format(db: Pool<Postgres>) {
let ordered_result_job_id = "1eecb96a-c8b0-4a3d-b1b6-087878c55e41";
set_jwt_secret().await;
let server = ApiServer::start(db.clone()).await;
let port = server.addr.port();
let token = windmill_worker::create_token_for_owner(
&db,
"test-workspace",
"u/test-user",
"",
100,
"",
&Uuid::nil(),
)
.await
.unwrap();
#[derive(Debug, Deserialize)]
struct JobResponse {
result: Option<Box<serde_json::value::RawValue>>,
}
async fn get_result<T: DeserializeOwned>(url: String) -> T {
reqwest::get(url)
.await
.unwrap()
.error_for_status()
.unwrap()
.json()
.await
.unwrap()
}
let correct_result = r#"[{"b":"first","a":"second"}]"#;
let job_response: JobResponse = get_result(format!("http://localhost:{port}/api/w/test-workspace/jobs_u/get/{ordered_result_job_id}?token={token}&no_logs=true")).await;
assert_eq!(job_response.result.unwrap().get(), correct_result);
let job_response: JobResponse = get_result(format!("http://localhost:{port}/api/w/test-workspace/jobs_u/completed/get_result_maybe/{ordered_result_job_id}?token={token}&no_logs=true")).await;
assert_eq!(job_response.result.unwrap().get(), correct_result);
let job_result: Box<serde_json::value::RawValue> = get_result(format!("http://localhost:{port}/api/w/test-workspace/jobs_u/completed/get_result/{ordered_result_job_id}?token={token}&no_logs=true")).await;
assert_eq!(job_result.get(), correct_result);
let response = windmill_api::jobs::run_wait_result(
&db,
Uuid::parse_str(ordered_result_job_id).unwrap(),
"test-workspace".to_string(),
None,
"test-user",
)
.await
.unwrap();
let result: Box<serde_json::value::RawValue> = serde_json::from_slice(
&axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.unwrap()
.to_vec(),
)
.unwrap();
assert_eq!(result.get(), correct_result);
}
+94 -67
View File
@@ -21,10 +21,7 @@ use tokio::time::Instant;
use tower::ServiceBuilder;
use windmill_common::error::JsonResult;
use windmill_common::flow_status::{JobResult, RestartedFrom};
use windmill_common::jobs::{
format_completed_job_result, format_result, CompletedJobWithFormattedResult, FormattedResult,
ENTRYPOINT_OVERRIDE,
};
use windmill_common::jobs::{format_completed_job_result, format_result, ENTRYPOINT_OVERRIDE};
use windmill_common::worker::{CLOUD_HOSTED, TMP_DIR};
#[cfg(all(feature = "enterprise", feature = "parquet"))]
@@ -738,7 +735,11 @@ pub async fn get_queued_job_ex(
// first optional is if authed need to be checked, second is the opt_authed itself
opt_authed: Option<&Option<ApiAuthed>>,
) -> error::Result<Option<JobExtended<QueuedJob>>> {
let query = if no_logs { &*GET_QUEUED_JOB_QUERY_NO_LOGS } else { &*GET_QUEUED_JOB_QUERY };
let query = if no_logs {
&*GET_QUEUED_JOB_QUERY_NO_LOGS
} else {
&*GET_QUEUED_JOB_QUERY
};
let job = sqlx::query_as::<_, JobExtended<QueuedJob>>(query)
.bind(job_id)
.bind(workspace_id)
@@ -763,7 +764,11 @@ pub async fn get_completed_job_ex(
// first optional is if authed need to be checked, second is the opt_authed itself
opt_authed: Option<&Option<ApiAuthed>>,
) -> error::Result<Option<JobExtended<CompletedJob>>> {
let query = if no_logs { &*GET_COMPLETED_JOB_QUERY_NO_LOGS } else { &*GET_COMPLETED_JOB_QUERY };
let query = if no_logs {
&*GET_COMPLETED_JOB_QUERY_NO_LOGS
} else {
&*GET_COMPLETED_JOB_QUERY
};
let cjob = sqlx::query_as::<_, JobExtended<CompletedJob>>(query)
.bind(job_id)
.bind(workspace_id)
@@ -779,13 +784,7 @@ pub async fn get_completed_job_ex(
}
if let Some(mut cjob) = cjob {
let CompletedJobWithFormattedResult { mut cj, result } = format_completed_job_result(cjob.inner);
cj.result = match result {
Some(FormattedResult::RawValue(rv)) => rv,
Some(FormattedResult::Vec(v)) => Some(to_raw_value(&v)),
None => None,
}.map(sqlx::types::Json);
cjob.inner = cj;
cjob.inner = format_completed_job_result(cjob.inner);
return Ok(Some(cjob));
}
@@ -965,10 +964,14 @@ async fn get_job_logs(
log_job_view(&db, opt_authed.as_ref(), &w_id, &id).await?;
#[cfg(all(feature = "enterprise", feature = "parquet"))]
if let Some(r) = get_logs_from_store(text.log_offset.unwrap_or(0), &logs, &text.log_file_index).await {
if let Some(r) =
get_logs_from_store(text.log_offset.unwrap_or(0), &logs, &text.log_file_index).await
{
return r.map(content_plain);
}
if let Some(r) = get_logs_from_disk(text.log_offset.unwrap_or(0), &logs, &text.log_file_index).await {
if let Some(r) =
get_logs_from_disk(text.log_offset.unwrap_or(0), &logs, &text.log_file_index).await
{
return r.map(content_plain);
}
@@ -1601,8 +1604,7 @@ async fn count_completed_jobs_detail(
Query(query): Query<CountCompletedJobsQuery>,
) -> error::JsonResult<i64> {
let mut sqlb = SqlBuilder::select_from("completed_job");
sqlb
.field("COUNT(*) as count");
sqlb.field("COUNT(*) as count");
if !query.all_workspaces.unwrap_or(false) {
sqlb.and_where_eq("workspace_id", "?".bind(&w_id));
@@ -1610,7 +1612,10 @@ async fn count_completed_jobs_detail(
if let Some(after_s_ago) = query.completed_after_s_ago {
let after = Utc::now() - chrono::Duration::seconds(after_s_ago);
sqlb.and_where_gt("started_at + duration_ms / 1000 * interval '1 second'", "?".bind(&after.to_rfc3339()));
sqlb.and_where_gt(
"started_at + duration_ms / 1000 * interval '1 second'",
"?".bind(&after.to_rfc3339()),
);
}
if let Some(success) = query.success {
@@ -1618,19 +1623,21 @@ async fn count_completed_jobs_detail(
}
if let Some(tags) = query.tags {
sqlb.and_where_in("tag", &tags.split(",").map(|t| format!("'{}'", t)).collect::<Vec<_>>());
sqlb.and_where_in(
"tag",
&tags
.split(",")
.map(|t| format!("'{}'", t))
.collect::<Vec<_>>(),
);
}
let sql = sqlb.sql()?;
let stats = sqlx::query_scalar::<_, i64>(&sql)
.fetch_one(&db)
.await?;
let stats = sqlx::query_scalar::<_, i64>(&sql).fetch_one(&db).await?;
Ok(Json(stats))
}
async fn count_completed_jobs(
Extension(db): Extension<DB>,
Path(w_id): Path<String>,
@@ -2316,8 +2323,19 @@ pub struct JobExtended<T> {
}
impl<T> JobExtended<T> {
pub fn new(self_wait_time_ms: Option<i64>, aggregate_wait_time_ms: Option<i64>, inner: T) -> Self {
Self { inner, raw_code: None, raw_lock: None, raw_flow: None, self_wait_time_ms, aggregate_wait_time_ms }
pub fn new(
self_wait_time_ms: Option<i64>,
aggregate_wait_time_ms: Option<i64>,
inner: T,
) -> Self {
Self {
inner,
raw_code: None,
raw_lock: None,
raw_flow: None,
self_wait_time_ms,
aggregate_wait_time_ms,
}
}
}
@@ -2619,7 +2637,10 @@ impl UnifiedJob {
impl<'a> From<UnifiedJob> for Job {
fn from(uj: UnifiedJob) -> Self {
match uj.typ.as_ref() {
"CompletedJob" => Job::CompletedJob(JobExtended::new(uj.self_wait_time_ms, uj.aggregate_wait_time_ms, CompletedJob {
"CompletedJob" => Job::CompletedJob(JobExtended::new(
uj.self_wait_time_ms,
uj.aggregate_wait_time_ms,
CompletedJob {
workspace_id: uj.workspace_id,
id: uj.id,
parent_job: uj.parent_job,
@@ -2650,8 +2671,12 @@ impl<'a> From<UnifiedJob> for Job {
tag: uj.tag,
priority: uj.priority,
labels: uj.labels,
})),
"QueuedJob" => Job::QueuedJob(JobExtended::new(uj.self_wait_time_ms, uj.aggregate_wait_time_ms, QueuedJob {
},
)),
"QueuedJob" => Job::QueuedJob(JobExtended::new(
uj.self_wait_time_ms,
uj.aggregate_wait_time_ms,
QueuedJob {
workspace_id: uj.workspace_id,
id: uj.id,
parent_job: uj.parent_job,
@@ -2689,7 +2714,8 @@ impl<'a> From<UnifiedJob> for Job {
flow_step_id: None,
cache_ttl: None,
priority: uj.priority,
})),
},
)),
t => panic!("job type {} not valid", t),
}
}
@@ -3313,7 +3339,7 @@ pub struct WindmillCompositeResult {
windmill_content_type: Option<String>,
result: Option<Box<RawValue>>,
}
async fn run_wait_result(
pub async fn run_wait_result(
db: &DB,
uuid: Uuid,
w_id: String,
@@ -3360,15 +3386,13 @@ async fn run_wait_result(
.bind(&w_id)
.fetch_optional(db)
.await?;
if let Some(raw_result) = row {
result = match format_result(
if let Some(mut raw_result) = row {
format_result(
raw_result.language.as_ref(),
raw_result.flow_status.map(|x| x.0),
raw_result.result.map(|x| x.0),
) {
FormattedResult::RawValue(rv) => rv,
FormattedResult::Vec(v) => Some(to_raw_value(&v)),
};
raw_result.flow_status.as_ref(),
raw_result.result.as_mut(),
);
result = raw_result.result.map(|x| x.0);
}
}
@@ -4367,7 +4391,7 @@ async fn add_batch_jobs(
raw_code,
raw_lock,
raw_flow,
flow_status
flow_status,
) = match batch_info.kind.as_str() {
"script" => {
if let Some(path) = batch_info.path {
@@ -4458,19 +4482,19 @@ async fn add_batch_jobs(
add_virtual_items_if_necessary(&mut value.modules);
let flow_status = FlowStatus::new(&value);
(
None, // script_hash
path, // script_path
job_kind, // job_kind
None, // language
None, // dedicated_worker
value.concurrency_key.clone(), // custom_concurrency_key
value.concurrent_limit.clone(), // concurrent_limit
value.concurrency_time_window_s, // concurrency_time_window_s
None, // timeout
None, // raw_code
None, // raw_lock
Some(value), // raw_flow
Some(flow_status), // flow_status
None, // script_hash
path, // script_path
job_kind, // job_kind
None, // language
None, // dedicated_worker
value.concurrency_key.clone(), // custom_concurrency_key
value.concurrent_limit.clone(), // concurrent_limit
value.concurrency_time_window_s, // concurrency_time_window_s
None, // timeout
None, // raw_code
None, // raw_lock
Some(value), // raw_flow
Some(flow_status), // flow_status
)
}
"noop" => (
@@ -4516,7 +4540,11 @@ async fn add_batch_jobs(
(id, workspace_id, raw_code, raw_lock, raw_flow)
(SELECT uuid, $1, $2, $3, $4 FROM uuid_table)
RETURNING id"#,
w_id, raw_code, raw_lock, raw_flow.map(sqlx::types::Json) as Option<sqlx::types::Json<FlowValue>>, n
w_id,
raw_code,
raw_lock,
raw_flow.map(sqlx::types::Json) as Option<sqlx::types::Json<FlowValue>>,
n
)
.fetch_all(&db)
.await?;
@@ -5122,8 +5150,7 @@ async fn get_completed_job<'a>(
Extension(db): Extension<DB>,
Path((w_id, id)): Path<(String, Uuid)>,
) -> error::Result<Response> {
let job_o = get_completed_job_ex(&db, &w_id, id, false, Some(&opt_authed))
.await?;
let job_o = get_completed_job_ex(&db, &w_id, id, false, Some(&opt_authed)).await?;
let cj = not_found_if_none(job_o, "Completed Job", id.to_string())?;
let response = Json(cj).into_response();
@@ -5186,7 +5213,7 @@ async fn get_completed_job_result(
.await?
};
let raw_result = not_found_if_none(result_o, "Completed Job", id.to_string())?;
let mut raw_result = not_found_if_none(result_o, "Completed Job", id.to_string())?;
if opt_authed.is_none() && raw_result.created_by.unwrap_or_default() != "anonymous" {
match (suspended_job, resume_id, approver, secret) {
@@ -5226,15 +5253,15 @@ async fn get_completed_job_result(
}
}
let result = format_result(
format_result(
raw_result.language.as_ref(),
raw_result.flow_status.map(|x| x.0),
raw_result.result.map(|x| x.0),
raw_result.flow_status.as_ref(),
raw_result.result.as_mut(),
);
log_job_view(&db, opt_authed.as_ref(), &w_id, &id).await?;
Ok(Json(result).into_response())
Ok(Json(raw_result.result).into_response())
}
#[derive(Deserialize)]
@@ -5280,7 +5307,7 @@ struct CompletedJobResult {
started: Option<bool>,
success: Option<bool>,
completed: bool,
result: Option<FormattedResult>,
result: Option<sqlx::types::Json<Box<RawValue>>>,
}
#[derive(Deserialize)]
@@ -5302,11 +5329,11 @@ async fn get_completed_job_result_maybe(
.fetch_optional(&db)
.await?;
if let Some(res) = result_o {
let result = format_result(
if let Some(mut res) = result_o {
format_result(
res.language.as_ref(),
res.flow_status.map(|x| x.0),
res.result.map(|x| x.0),
res.flow_status.as_ref(),
res.result.as_mut(),
);
if opt_authed.is_none() && res.created_by != "anonymous" {
return Err(Error::BadRequest(
@@ -5320,7 +5347,7 @@ async fn get_completed_job_result_maybe(
started: Some(true),
success: Some(res.success),
completed: true,
result: Some(result),
result: res.result,
})
.into_response())
} else if get_started.is_some_and(|x| x) {
+15 -31
View File
@@ -466,20 +466,6 @@ pub async fn get_payload_tag_from_prefixed_path<'e, E: sqlx::Executor<'e, Databa
Ok((payload, tag))
}
#[derive(Serialize, Debug)]
#[serde(untagged)]
pub enum FormattedResult {
RawValue(Option<Box<RawValue>>),
Vec(Vec<Box<RawValue>>),
}
#[derive(Serialize, Debug)]
pub struct CompletedJobWithFormattedResult {
#[serde(flatten)]
pub cj: CompletedJob,
pub result: Option<FormattedResult>,
}
#[derive(Deserialize)]
struct FlowStatusMetadata {
column_order: Vec<String>,
@@ -493,7 +479,7 @@ struct FlowStatusWithMetadataOnly {
pub fn order_columns(
rows: Option<Vec<Box<RawValue>>>,
column_order: Vec<String>,
) -> Option<Vec<Box<RawValue>>> {
) -> Option<Box<RawValue>> {
if let Some(mut rows) = rows {
if let Some(first_row) = rows.get(0) {
let first_row = serde_json::from_str::<HashMap<String, Box<RawValue>>>(first_row.get());
@@ -508,7 +494,7 @@ pub fn order_columns(
rows[0] = new_row_as_raw_value;
return Some(rows);
return Some(to_raw_value(&rows));
}
}
}
@@ -518,9 +504,9 @@ pub fn order_columns(
pub fn format_result(
language: Option<&ScriptLang>,
flow_status: Option<Box<RawValue>>,
result: Option<Box<RawValue>>,
) -> FormattedResult {
flow_status: Option<&sqlx::types::Json<Box<RawValue>>>,
result: Option<&mut sqlx::types::Json<Box<RawValue>>>,
) -> () {
match language {
Some(&ScriptLang::Postgresql)
| Some(&ScriptLang::Mysql)
@@ -531,27 +517,26 @@ pub fn format_result(
{
if let Some(result) = result {
let rows = serde_json::from_str::<Vec<Box<RawValue>>>(result.get()).ok();
match order_columns(rows, flow_status._metadata.column_order) {
Some(rows) => return FormattedResult::Vec(rows),
None => return FormattedResult::RawValue(Some(result)),
if let Some(ordered_result) =
order_columns(rows, flow_status._metadata.column_order)
{
*result = sqlx::types::Json(ordered_result);
}
}
}
}
_ => {}
}
FormattedResult::RawValue(result)
}
pub fn format_completed_job_result(mut cj: CompletedJob) -> CompletedJobWithFormattedResult {
let sql_result = format_result(
pub fn format_completed_job_result(mut cj: CompletedJob) -> CompletedJob {
format_result(
cj.language.as_ref(),
cj.flow_status.clone().map(|x| x.0),
cj.result.map(|x| x.0),
cj.flow_status.as_ref(),
cj.result.as_mut(),
);
cj.result = None; // very important to avoid sending the result twice
CompletedJobWithFormattedResult { cj, result: Some(sql_result) }
cj
}
pub async fn get_logs_from_disk(
@@ -598,7 +583,6 @@ pub async fn get_logs_from_store(
if log_offset > 0 {
if let Some(file_index) = log_file_index.clone() {
if let Some(os) = OBJECT_STORE_CACHE_SETTINGS.read().await.clone() {
let logs = logs.to_string();
let stream = async_stream::stream! {
for file_p in file_index.clone() {
@@ -533,7 +533,7 @@
/>
</div>
{:else if !forceJson && resultKind === 'pdf'}
<div class="h-96 mt-2">
<div class="h-96 mt-2 border">
<PdfViewer
allowFullscreen
source="data:application/pdf;base64,{contentOrRootString(result.pdf)}"
@@ -679,7 +679,7 @@
/>
</div>
{:else if result?.s3?.endsWith('.pdf')}
<div class="h-96 mt-2">
<div class="h-96 mt-2 border">
<PdfViewer
allowFullscreen
source={`/api/w/${workspaceId}/job_helpers/load_image_preview?file_key=${result.s3}` +
@@ -755,7 +755,7 @@
</button>
{/if}
{:else if s3object?.s3?.endsWith('.pdf')}
<div class="h-96 mt-2">
<div class="h-96 mt-2 border">
<PdfViewer
allowFullscreen
source={`/api/w/${workspaceId}/job_helpers/load_image_preview?file_key=${s3object.s3}` +