starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-10-21 08:03:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

query args from job_args table first

Browse Source
This commit is contained in:
yacine Bouraroui
2024-10-01 19:47:07 +02:00
parent f4df3affa3
commit 045a52cded
4 changed files with 76 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-api/src/inputs.rs
+13 -19
View File
@@ -21,11 +21,7 @@ use std::{
vec,
};
use windmill_common::{
db::UserDB,
error::JsonResult,
jobs::JobKind,
scripts::to_i64,
utils::{not_found_if_none, paginate, Pagination},
db::UserDB, error::JsonResult, jobs::JobKind, query_scalar_with_fallback, scripts::to_i64, utils::{not_found_if_none, paginate, Pagination}
};
pub fn workspaced_service() -> Router {
Router::new()
@@ -178,6 +174,7 @@ struct GetArgs {
input: Option<bool>,
allow_large: Option<bool>,
}
async fn get_args_from_history_or_saved_input(
authed: ApiAuthed,
Extension(user_db): Extension<UserDB>,
@@ -185,6 +182,7 @@ async fn get_args_from_history_or_saved_input(
Path((w_id, job_or_input_id)): Path<(String, Uuid)>,
) -> JsonResult<Option<Value>> {
let mut tx = user_db.begin(&authed).await?;
let result_o = if let Some(input) = g.input {
if input {
sqlx::query_scalar!(
@@ -196,24 +194,20 @@ async fn get_args_from_history_or_saved_input(
.fetch_optional(&mut *tx)
.await?
} else {
sqlx::query_scalar!(
query_scalar_with_fallback!(tx,
"SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM job_args WHERE id = $1 AND workspace_id = $2",
"SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM completed_job WHERE id = $1 AND workspace_id = $2",
job_or_input_id,
job_or_input_id,
w_id,
g.allow_large.unwrap_or(true)
)
.fetch_optional(&mut *tx)
.await?
g.allow_large.unwrap_or(true))?
}
} else {
sqlx::query_scalar!(
"SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM completed_job WHERE id = $1 AND workspace_id = $2 UNION ALL SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM input WHERE id = $1 AND workspace_id = $2",
job_or_input_id,
w_id,
g.allow_large.unwrap_or(true)
)
.fetch_optional(&mut *tx)
.await?
query_scalar_with_fallback!(tx,
"SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM job_args WHERE id = $1 AND workspace_id = $2 UNION ALL SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM input WHERE id = $1 AND workspace_id = $2",
"SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM completed_job WHERE id = $1 AND workspace_id = $2 UNION ALL SELECT CASE WHEN pg_column_size(args) < 40000 OR $3 THEN args ELSE '\"WINDMILL_TOO_BIG\"'::jsonb END as args FROM input WHERE id = $1 AND workspace_id = $2",
job_or_input_id,
w_id,
g.allow_large.unwrap_or(true))?
};
tx.commit().await?;
backend/windmill-common/src/lib.rs
+1
View File
@@ -6,6 +6,7 @@
* LICENSE-AGPL for a copy of the license.
*/
pub mod macros;
use std::{
net::SocketAddr,
sync::{atomic::AtomicBool, Arc},
backend/windmill-common/src/macros.rs
+37
View File
@@ -0,0 +1,37 @@
#[macro_export]
macro_rules! fetch_one_with_fallback {
($db:expr, $query_method:ident, $fetch_method:ident ,$row_type:ty, $query:literal, $table:literal || $fallback_table:literal, $( $param:expr ),* ) => {{
let primary_query = sqlx::$query_method::<_, $row_type>(const_format::formatcp!($query, $table))
$(.bind($param))*
.$fetch_method($db)
.await;
if let Err(sqlx::Error::RowNotFound) = primary_query {
tracing::info!("Data not found in job_params, falling back to fetching from $fallback_table");
sqlx::$query_method::<_, $row_type>(const_format::formatcp!($query, $fallback_table))
$(.bind($param))*
.$fetch_method($db)
.await
} else {
primary_query
}
}};
}
#[macro_export]
macro_rules! query_scalar_with_fallback {
($tx:expr, $query:literal, $fallback_query:literal,$( $param:expr ),* ) => {{
let primary_query = sqlx::query_scalar!($query,$($param),*)
.fetch_optional(&mut *$tx)
.await;
if let Err(sqlx::Error::RowNotFound) = primary_query {
tracing::info!("Data not found in job_args, falling back to fetching from queue");
sqlx::query_scalar!($fallback_query, $($param),*)
.fetch_optional(&mut *$tx)
.await
} else {
primary_query
}
}};
}
backend/windmill-worker/src/worker_flow.rs
+25 -36
View File
@@ -29,7 +29,6 @@ use sqlx::FromRow;
use tokio::sync::mpsc::Sender;
use tracing::instrument;
use uuid::Uuid;
use windmill_common::add_time;
use windmill_common::auth::JobPerms;
#[cfg(feature = "benchmark")]
use windmill_common::bench::BenchmarkIter;
@@ -43,6 +42,7 @@ use windmill_common::jobs::{
RawCode, ENTRYPOINT_OVERRIDE,
};
use windmill_common::worker::to_raw_value;
use windmill_common::{add_time, fetch_one_with_fallback};
use windmill_common::{
error::{self, to_anyhow, Error},
flow_status::{
@@ -61,32 +61,14 @@ type DB = sqlx::Pool<sqlx::Postgres>;
use windmill_queue::{canceled_job_to_result, get_queued_job_tx, push, QueueTransaction};
macro_rules! fetch_one_with_fallback {
($db:expr, $fetch_method:ident, $row_type:ty, $query:literal, $table:literal ,$( $param:expr ),* ) => {{
let primary_query = sqlx::$fetch_method::<_, $row_type>(const_format::formatcp!($query, $table))
$(.bind($param))*
.fetch_one($db)
.await;
if let Err(sqlx::Error::RowNotFound) = primary_query {
tracing::info!("Data not found in job_params, falling back to fetching from queue");
sqlx::$fetch_method::<_, $row_type>(const_format::formatcp!($query, "queue"))
$(.bind($param))*
.fetch_one($db)
.await
} else {
primary_query
}
}};
}
async fn get_args_from_job_id(db: &DB, id: &Uuid) -> Result<RowArgs, Error> {
let args = fetch_one_with_fallback!(
db,
query_as,
fetch_one,
RowArgs,
"SELECT args FROM {} WHERE id = $1",
"job_args",
"job_args" || "queue",
id
);
@@ -334,7 +316,12 @@ pub async fn update_flow_status_after_job_completion_internal<
};
let is_flow = if let Some(step) = step {
fetch_one_with_fallback!(db, query_scalar, Option<bool>, "SELECT raw_flow->'modules'->($1)->'value'->>'type' = 'flow' FROM {} WHERE id = $2","job_params",
fetch_one_with_fallback!(db,
query_scalar,
fetch_one,
Option<bool>,
"SELECT raw_flow->'modules'->($1)->'value'->>'type' = 'flow' FROM {} WHERE id = $2",
"job_params" || "queue",
step as i32,
&flow).map_err(|e| {
Error::InternalErr(format!("error during retrieval of step's type: {e:#}"))
@@ -459,7 +446,7 @@ pub async fn update_flow_status_after_job_completion_internal<
None
};
let nindex = if let Some(position) = position {
let nindex = if let Some(position) = position {
sqlx::query_scalar!(
"UPDATE queue
SET flow_status = JSONB_SET(
@@ -1283,9 +1270,10 @@ async fn has_failure_module<'c>(flow: Uuid, db: &DB) -> Result<bool, Error> {
fetch_one_with_fallback!(
db,
query_scalar,
fetch_one,
Option<bool>,
"SELECT raw_flow->'failure_module' != 'null'::jsonb FROM {} WHERE id = $1",
"job_params",
"job_params" || "queue",
flow
)
.map_err(|e| {
@@ -2233,18 +2221,19 @@ async fn push_next_flow_job<R: rsmq_async::RsmqConnection + Send + Sync + Clone>
);
Ok(Marc::new(hm))
} else if let Some(id) = get_args_from_id {
let row = sqlx::query_as::<_, RawArgs>(
"SELECT args FROM completed_job WHERE id = $1 AND workspace_id = $2",
)
.bind(id)
.bind(&flow_job.workspace_id)
.fetch_optional(db)
.await?;
let row = fetch_one_with_fallback!(
db,
query_as,
fetch_optional,
RowArgs,
"SELECT args FROM {} WHERE id = $1 AND workspace_id = $2",
"job_args" || "completed_job",
id,
&flow_job.workspace_id
)?;
if let Some(raw_args) = row {
Ok(Marc::new(
raw_args.args.map(|x| x.0).unwrap_or_else(HashMap::new),
))
Ok(Marc::new(raw_args.args.map(|x| x.0).unwrap_or_default()))
} else {
Ok(Marc::new(HashMap::new()))
}
@@ -2361,7 +2350,7 @@ async fn push_next_flow_job<R: rsmq_async::RsmqConnection + Send + Sync + Clone>
let mut tx: QueueTransaction<'_, R> = (rsmq.clone(), db.begin().await?).into();
let nargs = args.as_ref();
for i in (0..len).into_iter() {
for i in 0..len {
if i % 100 == 0 && i != 0 {
tracing::info!(id = %flow_job.id, root_id = %job_root, "pushed (non-commited yet) first {i} subflows of {len}");
sqlx::query!(