perf: optimize number of queries needed for job run (#5504)

* optPerf

* update sqlx

* update sqlx
This commit is contained in:
Ruben Fiszel
2025-03-24 20:12:26 +01:00
committed by GitHub
parent a3f6db7dca
commit 3edca4bc91
44 changed files with 375 additions and 382 deletions
@@ -39,7 +39,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -130,28 +130,28 @@
},
{
"ordinal": 25,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 26,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 27,
"name": "teams_team_name",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "ai_models",
"type_info": "VarcharArray"
},
{
"ordinal": 29,
"ordinal": 26,
"name": "code_completion_model",
"type_info": "Varchar"
},
{
"ordinal": 27,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 29,
"name": "teams_team_name",
"type_info": "Text"
}
],
"parameters": {
@@ -185,10 +185,10 @@
true,
true,
true,
true,
true,
true,
false,
true,
true,
true,
true
]
},
@@ -68,7 +68,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -1,50 +1,35 @@
{
"db_name": "PostgreSQL",
"query": "SELECT * FROM job_perms WHERE job_id = $1 AND workspace_id = $2",
"query": "SELECT email, username, is_admin, is_operator, groups, folders FROM job_perms WHERE job_id = $1 AND workspace_id = $2",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "job_id",
"type_info": "Uuid"
},
{
"ordinal": 1,
"name": "email",
"type_info": "Varchar"
},
{
"ordinal": 2,
"ordinal": 1,
"name": "username",
"type_info": "Varchar"
},
{
"ordinal": 3,
"ordinal": 2,
"name": "is_admin",
"type_info": "Bool"
},
{
"ordinal": 4,
"ordinal": 3,
"name": "is_operator",
"type_info": "Bool"
},
{
"ordinal": 5,
"name": "created_at",
"type_info": "Timestamp"
},
{
"ordinal": 6,
"name": "workspace_id",
"type_info": "Varchar"
},
{
"ordinal": 7,
"ordinal": 4,
"name": "groups",
"type_info": "TextArray"
},
{
"ordinal": 8,
"ordinal": 5,
"name": "folders",
"type_info": "JsonbArray"
}
@@ -61,11 +46,8 @@
false,
false,
false,
false,
false,
false,
false
]
},
"hash": "d2def87d7f7901eebc65082f7df5e0a33e5702b25c3db3affa06155e90480e42"
"hash": "2a510a8bec98055796f987d86c344ca116895d71de09db338ec09e425dcebe5e"
}
@@ -68,7 +68,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -135,7 +135,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -130,28 +130,28 @@
},
{
"ordinal": 25,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 26,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 27,
"name": "teams_team_name",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "ai_models",
"type_info": "VarcharArray"
},
{
"ordinal": 29,
"ordinal": 26,
"name": "code_completion_model",
"type_info": "Varchar"
},
{
"ordinal": 27,
"name": "teams_command_script",
"type_info": "Text"
},
{
"ordinal": 28,
"name": "teams_team_id",
"type_info": "Text"
},
{
"ordinal": 29,
"name": "teams_team_name",
"type_info": "Text"
}
],
"parameters": {
@@ -185,10 +185,10 @@
true,
true,
true,
true,
true,
true,
false,
true,
true,
true,
true
]
},
@@ -59,7 +59,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -59,7 +59,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -1,29 +0,0 @@
{
"db_name": "PostgreSQL",
"query": "SELECT v2_job_completed.id AS \"id!\", flow_status AS \"flow_status!: Json<FlowStatus>\"\n FROM v2_job_completed\n INNER JOIN v2_job ON (v2_job_completed.id = v2_job.id)\n WHERE parent_job = $1 AND v2_job_completed.workspace_id = $2 AND flow_status IS NOT NULL",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "id!",
"type_info": "Uuid"
},
{
"ordinal": 1,
"name": "flow_status!: Json<FlowStatus>",
"type_info": "Jsonb"
}
],
"parameters": {
"Left": [
"Uuid",
"Text"
]
},
"nullable": [
false,
true
]
},
"hash": "7e93c924e3fc51f8c26df26e5d09d60e3a3a40b90421aaf589c4c3bcc5a45ec8"
}
@@ -32,7 +32,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -59,7 +59,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -63,7 +63,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -40,7 +40,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -54,7 +54,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
@@ -63,7 +63,8 @@
"rust",
"ansible",
"csharp",
"oracledb"
"oracledb",
"nu"
]
}
}
+1
View File
@@ -1803,6 +1803,7 @@ async fn handle_zombie_jobs(db: &Pool<Postgres>, base_internal_url: &str, worker
*SCRIPT_TOKEN_EXPIRY,
&job.email,
&job.id,
None,
)
.await
.expect("could not create job token");
+3 -2
View File
@@ -315,7 +315,7 @@ mod suspend_resume {
let second = completed.next().await.unwrap();
// print_job(second, &db).await;
let token = windmill_worker::create_token_for_owner(&db, "test-workspace", "u/test-user", "", 100, "", &Uuid::nil()).await.unwrap();
let token = windmill_worker::create_token_for_owner(&db, "test-workspace", "u/test-user", "", 100, "", &Uuid::nil(), None).await.unwrap();
let secret = reqwest::get(format!(
"http://localhost:{port}/api/w/test-workspace/jobs/job_signature/{second}/0?token={token}&approver=ruben"
))
@@ -418,7 +418,7 @@ mod suspend_resume {
/* ... and send a request resume it. */
let second = completed.next().await.unwrap();
let token = windmill_worker::create_token_for_owner(&db, "test-workspace", "u/test-user", "", 100, "", &Uuid::nil()).await.unwrap();
let token = windmill_worker::create_token_for_owner(&db, "test-workspace", "u/test-user", "", 100, "", &Uuid::nil(), None).await.unwrap();
let secret = reqwest::get(format!(
"http://localhost:{port}/api/w/test-workspace/jobs/job_signature/{second}/0?token={token}"
))
@@ -3806,6 +3806,7 @@ async fn test_result_format(db: Pool<Postgres>) {
100,
"",
&Uuid::nil(),
None,
)
.await
.unwrap();
+5 -1
View File
@@ -1208,7 +1208,11 @@ async fn update_resource_type(
Ok(format!("resource_type {} updated", name))
}
#[cfg(any(feature = "postgres_trigger", feature = "mqtt_trigger", all(feature = "sqs_trigger", feature = "enterprise")))]
#[cfg(any(
feature = "postgres_trigger",
feature = "mqtt_trigger",
all(feature = "sqs_trigger", feature = "enterprise")
))]
pub async fn try_get_resource_from_db_as<T>(
authed: ApiAuthed,
user_db: Option<UserDB>,
-3
View File
@@ -24,15 +24,12 @@ pub struct JWTAuthClaims {
#[derive(Deserialize)]
pub struct JobPerms {
pub workspace_id: String,
pub job_id: String,
pub email: String,
pub username: String,
pub is_admin: bool,
pub is_operator: bool,
pub groups: Vec<String>,
pub folders: Vec<serde_json::Value>,
pub created_at: chrono::NaiveDateTime,
}
impl From<JobPerms> for Authed {
+13 -9
View File
@@ -142,16 +142,20 @@ fn format_pull_query(peek: String) -> String {
raw_flow, script_entrypoint_override, preprocessed
FROM v2_job
WHERE id = (SELECT id FROM peek)
) SELECT id, workspace_id, parent_job, created_by, started_at, scheduled_for,
runnable_id, runnable_path, args, canceled_by,
canceled_reason, kind, trigger, trigger_kind, permissioned_as, permissioned_as_email,
flow_status, script_lang,
same_worker, pre_run_error, visible_to_owner,
tag, concurrent_limit, concurrency_time_window_s, flow_innermost_root_job,
timeout, flow_step_id, cache_ttl, priority, raw_code, raw_lock, raw_flow,
script_entrypoint_override, preprocessed
) SELECT j.id, j.workspace_id, j.parent_job, j.created_by, started_at, scheduled_for,
j.runnable_id, j.runnable_path, j.args, canceled_by,
canceled_reason, j.kind, j.trigger, j.trigger_kind, j.permissioned_as, j.permissioned_as_email,
flow_status, j.script_lang,
j.same_worker, j.pre_run_error, j.visible_to_owner,
j.tag, j.concurrent_limit, j.concurrency_time_window_s, j.flow_innermost_root_job,
j.timeout, j.flow_step_id, j.cache_ttl, j.priority, j.raw_code, j.raw_lock, j.raw_flow,
j.script_entrypoint_override, j.preprocessed, pj.runnable_path as parent_runnable_path,
p.email as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin,
p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders
FROM q, j
LEFT JOIN v2_job_status f USING (id)",
LEFT JOIN v2_job_status f USING (id)
LEFT JOIN job_perms p ON p.job_id = j.id
LEFT JOIN v2_job pj ON j.parent_job = pj.id",
peek
);
tracing::debug!("pull query: {}", r);
+7
View File
@@ -2002,6 +2002,13 @@ pub struct PulledJob {
pub raw_code: Option<String>,
pub raw_lock: Option<String>,
pub raw_flow: Option<Json<Box<RawValue>>>,
pub parent_runnable_path: Option<String>,
pub permissioned_as_email: Option<String>,
pub permissioned_as_username: Option<String>,
pub permissioned_as_is_admin: Option<bool>,
pub permissioned_as_is_operator: Option<bool>,
pub permissioned_as_groups: Option<Vec<String>>,
pub permissioned_as_folders: Option<Vec<serde_json::Value>>,
}
impl std::ops::Deref for PulledJob {
@@ -26,8 +26,8 @@ use crate::{
},
handle_child::handle_child,
python_executor::{create_dependencies_dir, handle_python_reqs, uv_pip_compile, PyVersion},
AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NSJAIL_PATH, PATH_ENV,
PROXY_ENVS, PY_INSTALL_DIR, TZ_ENV,
AuthedClient, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NSJAIL_PATH, PATH_ENV, PROXY_ENVS,
PY_INSTALL_DIR, TZ_ENV,
};
lazy_static::lazy_static! {
@@ -185,7 +185,8 @@ pub async fn handle_ansible_job(
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
inner_content: &String,
shared_mount: &str,
base_internal_url: &str,
@@ -263,7 +264,6 @@ pub async fn handle_ansible_job(
})
.unwrap_or_else(|| vec![]);
let authed_client = client.get_authed().await;
let mut nsjail_extra_mounts = vec![];
if let Some(r) = reqs {
nsjail_extra_mounts = create_file_resources(
@@ -272,7 +272,7 @@ pub async fn handle_ansible_job(
job_dir,
interpolated_args.as_ref(),
&r,
&authed_client,
&client,
db,
)
.await?;
@@ -311,7 +311,8 @@ remote_tmp={job_dir}/.ansible/tmp
);
write_file(job_dir, "ansible.cfg", &ansible_cfg_content)?;
let mut reserved_variables = get_reserved_variables(job, &authed_client.token, db).await?;
let mut reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
let additional_python_paths_folders = additional_python_paths.join(":");
if !*DISABLE_NSJAIL {
+9 -7
View File
@@ -46,7 +46,7 @@ use crate::{
OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NSJAIL_PATH, PATH_ENV,
AuthedClient, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NSJAIL_PATH, PATH_ENV,
POWERSHELL_CACHE_DIR, POWERSHELL_PATH, PROXY_ENVS, TZ_ENV,
};
@@ -64,7 +64,8 @@ pub async fn handle_bash_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
content: &str,
job_dir: &str,
shared_mount: &str,
@@ -135,8 +136,8 @@ exit $exit_status
);
write_file(job_dir, "wrapper.sh", &script)?;
let token = client.get_token().await;
let mut reserved_variables = get_reserved_variables(job, &token, db).await?;
let mut reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
reserved_variables.insert("RUST_LOG".to_string(), "info".to_string());
let args = build_args_map(job, client, db).await?.map(Json);
@@ -471,7 +472,8 @@ pub async fn handle_powershell_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
content: &str,
job_dir: &str,
shared_mount: &str,
@@ -653,8 +655,8 @@ $env:PSModulePath = \"{};$PSModulePathBackup\"",
),
)?;
let token = client.get_token().await;
let mut reserved_variables = get_reserved_variables(job, &token, db).await?;
let mut reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
reserved_variables.insert("RUST_LOG".to_string(), "info".to_string());
let _ = write_file(job_dir, "result.json", "")?;
@@ -18,7 +18,7 @@ use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
use crate::{
common::{build_args_values, resolve_job_timeout},
AuthedClientBackgroundTask,
AuthedClient,
};
use gcp_auth::{AuthenticationManager, CustomServiceAccount};
@@ -207,7 +207,7 @@ use windmill_queue::MiniPulledJob;
pub async fn do_bigquery(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -223,8 +223,6 @@ pub async fn do_bigquery(
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
@@ -267,13 +265,13 @@ pub async fn do_bigquery(
.map_err(|x| Error::ExecutionErr(x.to_string()))?
.args;
let (query, args_to_skip) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &bigquery_args)?;
let (query, args_to_skip) =
&sanitize_and_interpolate_unsafe_sql_args(query, &sig, &bigquery_args)?;
let queries = parse_sql_blocks(query);
let mut statement_values: HashMap<String, Value> = HashMap::new();
for arg in &sig {
if args_to_skip.contains(&arg.name) {
continue;
+9 -9
View File
@@ -20,9 +20,9 @@ use crate::{
read_file_content, read_result, start_child_process, write_file_binary, OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, BUNFIG_INSTALL_SCOPES, BUN_BUNDLE_CACHE_DIR, BUN_CACHE_DIR,
BUN_PATH, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NODE_BIN_PATH, NODE_PATH,
NPM_CONFIG_REGISTRY, NPM_PATH, NSJAIL_PATH, PATH_ENV, PROXY_ENVS, TZ_ENV,
AuthedClient, BUNFIG_INSTALL_SCOPES, BUN_BUNDLE_CACHE_DIR, BUN_CACHE_DIR, BUN_PATH,
DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NODE_BIN_PATH, NODE_PATH, NPM_CONFIG_REGISTRY,
NPM_PATH, NSJAIL_PATH, PATH_ENV, PROXY_ENVS, TZ_ENV,
};
#[cfg(windows)]
@@ -825,7 +825,8 @@ pub async fn handle_bun_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
job_dir: &str,
inner_content: &String,
base_internal_url: &str,
@@ -934,7 +935,7 @@ pub async fn handle_bun_job(
&job.id,
&job.workspace_id,
Some(db),
&client.get_token().await,
&client.token,
job.runnable_path(),
job_dir,
base_internal_url,
@@ -1108,8 +1109,7 @@ try {{
Ok(()) as Result<()>
};
let reserved_variables_f = async {
let client = client.get_authed().await;
let vars = get_reserved_variables(job, &client.token, db).await?;
let vars = get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
Ok(vars) as Result<HashMap<String, String>>
};
let (_, reserved_variables) = tokio::try_join!(args_and_out_f, reserved_variables_f)?;
@@ -1127,7 +1127,7 @@ try {{
build_loader(
job_dir,
base_internal_url,
&client.get_token().await,
&client.token,
&job.workspace_id,
job.runnable_path(),
if annotation.nodejs {
@@ -1145,7 +1145,7 @@ try {{
build_loader(
job_dir,
base_internal_url,
&client.get_token().await,
&client.token,
&job.workspace_id,
job.runnable_path(),
if annotation.nodejs {
+15 -45
View File
@@ -39,14 +39,11 @@ use windmill_common::{variables, DB};
use tokio::{io::AsyncWriteExt, process::Child, time::Instant};
use crate::{
AuthedClient, AuthedClientBackgroundTask, JOB_DEFAULT_TIMEOUT, MAX_RESULT_SIZE,
MAX_TIMEOUT_DURATION, PATH_ENV,
};
use crate::{AuthedClient, JOB_DEFAULT_TIMEOUT, MAX_RESULT_SIZE, MAX_TIMEOUT_DURATION, PATH_ENV};
pub async fn build_args_map<'a>(
job: &'a MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
db: &Pool<Postgres>,
) -> error::Result<Option<HashMap<String, Box<RawValue>>>> {
if let Some(args) = &job.args {
@@ -74,7 +71,7 @@ pub fn check_executor_binary_exists(
pub async fn build_args_values(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
db: &Pool<Postgres>,
) -> error::Result<HashMap<String, serde_json::Value>> {
if let Some(args) = &job.args {
@@ -86,7 +83,7 @@ pub async fn build_args_values(
#[tracing::instrument(level = "trace", skip_all)]
pub async fn create_args_and_out_file(
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
job: &MiniPulledJob,
job_dir: &str,
db: &Pool<Postgres>,
@@ -126,7 +123,7 @@ lazy_static::lazy_static! {
}
pub async fn transform_json<'a>(
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
workspace: &str,
vs: &'a HashMap<String, Box<RawValue>>,
job: &MiniPulledJob,
@@ -150,9 +147,7 @@ pub async fn transform_json<'a>(
let value = serde_json::from_str(inner_vs).map_err(|e| {
error::Error::internal_err(format!("Error while parsing inner arg: {e:#}"))
})?;
let transformed =
transform_json_value(&k, &client.get_authed().await, workspace, value, job, db)
.await?;
let transformed = transform_json_value(&k, &client, workspace, value, job, db).await?;
let as_raw = serde_json::from_value(transformed).map_err(|e| {
error::Error::internal_err(format!("Error while parsing inner arg: {e:#}"))
})?;
@@ -165,7 +160,7 @@ pub async fn transform_json<'a>(
}
pub async fn transform_json_as_values<'a>(
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
workspace: &str,
vs: &'a HashMap<String, Box<RawValue>>,
job: &MiniPulledJob,
@@ -178,9 +173,7 @@ pub async fn transform_json_as_values<'a>(
let value = serde_json::from_str(inner_vs).map_err(|e| {
error::Error::internal_err(format!("Error while parsing inner arg: {e:#}"))
})?;
let transformed =
transform_json_value(&k, &client.get_authed().await, workspace, value, job, db)
.await?;
let transformed = transform_json_value(&k, &client, workspace, value, job, db).await?;
let as_raw = serde_json::from_value(transformed).map_err(|e| {
error::Error::internal_err(format!("Error while parsing inner arg: {e:#}"))
})?;
@@ -283,40 +276,14 @@ pub async fn transform_json_value(
// let path = y.strip_prefix("$res:").unwrap();
}
Value::String(y) if y.starts_with("$") => {
let flow_path = if let Some(uuid) = job.parent_job {
sqlx::query_scalar!("SELECT runnable_path FROM v2_job WHERE id = $1", uuid)
.fetch_optional(db)
.await?
.flatten()
} else {
None
};
let variables = variables::get_reserved_variables(
db,
&job.workspace_id,
&client.token,
&job.permissioned_as_email,
&job.created_by,
&job.id.to_string(),
&job.permissioned_as,
job.runnable_path.clone(),
job.parent_job.map(|x| x.to_string()),
flow_path,
job.schedule_path(),
job.flow_step_id.clone(),
job.flow_innermost_root_job.clone().map(|x| x.to_string()),
None,
Some(job.scheduled_for.clone()),
)
.await;
let variables = get_reserved_variables(job, &client.token, &db, None).await?;
let name = y.strip_prefix("$").unwrap();
let value = variables
.iter()
.find(|x| x.name == name)
.map(|x| x.value.clone())
.find(|x| x.0 == name)
.map(|x| x.1.clone())
.unwrap_or_else(|| y);
Ok(json!(value))
}
@@ -417,8 +384,11 @@ pub async fn get_reserved_variables(
job: &MiniPulledJob,
token: &str,
db: &sqlx::Pool<sqlx::Postgres>,
parent_runnable_path: Option<String>,
) -> Result<HashMap<String, String>, Error> {
let flow_path = if let Some(uuid) = job.parent_job {
let flow_path = if parent_runnable_path.is_some() {
parent_runnable_path
} else if let Some(uuid) = job.parent_job {
sqlx::query_scalar!("SELECT runnable_path FROM v2_job WHERE id = $1", uuid)
.fetch_optional(db)
.await?
@@ -36,7 +36,7 @@ use crate::{
};
use crate::common::OccupancyMetrics;
use crate::AuthedClientBackgroundTask;
use crate::AuthedClient;
#[cfg(windows)]
use crate::SYSTEM_ROOT;
@@ -433,7 +433,8 @@ pub async fn handle_csharp_job(
_canceled_by: &mut Option<CanceledBy>,
_job: &MiniPulledJob,
_db: &sqlx::Pool<sqlx::Postgres>,
_client: &AuthedClientBackgroundTask,
_client: &AuthedClient,
_parent_runnable_path: Option<String>,
_inner_content: &str,
_job_dir: &str,
_requirements_o: Option<&String>,
@@ -452,7 +453,8 @@ pub async fn handle_csharp_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
inner_content: &str,
job_dir: &str,
requirements_o: Option<&String>,
@@ -534,8 +536,8 @@ pub async fn handle_csharp_job(
let logs2 = format!("{cache_logs}\n\n--- C# CODE EXECUTION ---\n");
append_logs(&job.id, &job.workspace_id, format!("{}\n", logs2), db).await;
let client = &client.get_authed().await;
let reserved_variables = get_reserved_variables(job, &client.token, db).await?;
let reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
let child = if !*DISABLE_NSJAIL {
write_file(
+10 -9
View File
@@ -11,8 +11,8 @@ use crate::{
start_child_process, OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, DENO_CACHE_DIR, DENO_PATH, DISABLE_NSJAIL, HOME_ENV,
NPM_CONFIG_REGISTRY, PATH_ENV, TZ_ENV,
AuthedClient, DENO_CACHE_DIR, DENO_PATH, DISABLE_NSJAIL, HOME_ENV, NPM_CONFIG_REGISTRY,
PATH_ENV, TZ_ENV,
};
use tokio::{fs::File, io::AsyncReadExt, process::Command};
use windmill_common::error::{self};
@@ -179,7 +179,8 @@ pub async fn handle_deno_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
job_dir: &str,
inner_content: &String,
base_internal_url: &str,
@@ -318,21 +319,21 @@ try {{
Ok(()) as Result<()>
};
let reserved_variables_f = async {
let client = client.get_authed().await;
let vars = get_reserved_variables(job, &client.token, db).await?;
Ok((vars, client.token)) as Result<(HashMap<String, String>, String)>
let vars = get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
Ok(vars) as Result<HashMap<String, String>>
};
let (_, reserved_variables) = tokio::try_join!(args_and_out_f, reserved_variables_f)?;
Ok(reserved_variables) as error::Result<(HashMap<String, String>, String)>
Ok(reserved_variables) as error::Result<HashMap<String, String>>
};
let ((reserved_variables, token), _, _) = tokio::try_join!(
let (reserved_variables, _, _) = tokio::try_join!(
reserved_variables_args_out_f,
write_wrapper_f,
write_import_map_f
)?;
let mut common_deno_proc_envs = get_common_deno_proc_envs(&token, base_internal_url).await;
let mut common_deno_proc_envs =
get_common_deno_proc_envs(&client.token, base_internal_url).await;
if !*DISABLE_NSJAIL {
common_deno_proc_envs.insert("HOME".to_string(), job_dir.to_string());
}
+6 -6
View File
@@ -19,8 +19,8 @@ use crate::{
start_child_process, OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, GOPRIVATE, GOPROXY,
GO_BIN_CACHE_DIR, GO_CACHE_DIR, HOME_ENV, NSJAIL_PATH, PATH_ENV, TZ_ENV,
AuthedClient, DISABLE_NSJAIL, DISABLE_NUSER, GOPRIVATE, GOPROXY, GO_BIN_CACHE_DIR,
GO_CACHE_DIR, HOME_ENV, NSJAIL_PATH, PATH_ENV, TZ_ENV,
};
const GO_REQ_SPLITTER: &str = "//go.sum\n";
@@ -37,7 +37,8 @@ pub async fn handle_go_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
inner_content: &str,
job_dir: &str,
requirements_o: Option<&String>,
@@ -248,9 +249,8 @@ func Run(req Req) (interface{{}}, error){{
let logs2 = format!("{cache_logs}\n\n--- GO CODE EXECUTION ---\n");
append_logs(&job.id, &job.workspace_id, logs2, db).await;
let client = &client.get_authed().await;
let reserved_variables = get_reserved_variables(job, &client.token, db).await?;
let reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
let child = if !*DISABLE_NSJAIL {
let _ = write_file(
@@ -13,7 +13,7 @@ use serde::Deserialize;
use crate::common::{build_http_client, resolve_job_timeout, OccupancyMetrics};
use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::{common::build_args_map, AuthedClientBackgroundTask};
use crate::{common::build_args_map, AuthedClient};
#[derive(Deserialize)]
struct GraphqlApi {
@@ -35,7 +35,7 @@ struct GraphqlError {
pub async fn do_graphql(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -18,7 +18,7 @@ use windmill_queue::{append_logs, CanceledBy};
use crate::common::{build_args_values, OccupancyMetrics};
use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
use crate::AuthedClientBackgroundTask;
use crate::AuthedClient;
#[derive(Deserialize)]
struct MssqlDatabase {
@@ -36,7 +36,7 @@ lazy_static::lazy_static! {
pub async fn do_mssql(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -51,8 +51,6 @@ pub async fn do_mssql(
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
@@ -133,7 +131,8 @@ pub async fn do_mssql(
.map_err(|x| Error::ExecutionErr(x.to_string()))?
.args;
let (query, args_to_skip) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &mssql_args)?;
let (query, args_to_skip) =
&sanitize_and_interpolate_unsafe_sql_args(query, &sig, &mssql_args)?;
let mut prepared_query = Query::new(query.to_owned());
for arg in &sig {
+15 -11
View File
@@ -21,7 +21,10 @@ use windmill_queue::CanceledBy;
use windmill_queue::MiniPulledJob;
use crate::{
common::{build_args_values, OccupancyMetrics}, handle_child::run_future_with_polling_update_job_poller, sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args, AuthedClientBackgroundTask
common::{build_args_values, OccupancyMetrics},
handle_child::run_future_with_polling_update_job_poller,
sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args,
AuthedClient,
};
#[derive(Deserialize)]
@@ -101,7 +104,7 @@ pub fn do_mysql_inner<'a>(
pub async fn do_mysql(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -115,14 +118,14 @@ pub async fn do_mysql(
let inline_db_res_path = parse_db_resource(&query);
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
)
.await?)
Some(
client
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
)
.await?,
)
} else {
job_args.get("database").cloned()
};
@@ -171,7 +174,8 @@ pub async fn do_mysql(
}
let arg_t = arg.otyp.clone().unwrap_or_else(|| "text".to_string());
let arg_n = arg.name.clone();
let mysql_v = match job_args.get(arg.name.as_str())
let mysql_v = match job_args
.get(arg.name.as_str())
.unwrap_or_else(|| &json!(null))
{
Value::Null => mysql_async::Value::NULL,
+6 -5
View File
@@ -13,8 +13,7 @@ use crate::{
create_args_and_out_file, get_reserved_variables, read_result, start_child_process,
OccupancyMetrics,
},
handle_child, AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, NSJAIL_PATH, PATH_ENV,
PROXY_ENVS,
handle_child, AuthedClient, DISABLE_NSJAIL, DISABLE_NUSER, NSJAIL_PATH, PATH_ENV, PROXY_ENVS,
};
const NSJAIL_CONFIG_RUN_NU_CONTENT: &str = include_str!("../nsjail/run.nu.config.proto");
@@ -29,7 +28,8 @@ lazy_static::lazy_static! {
pub(crate) struct JobHandlerInput<'a> {
pub base_internal_url: &'a str,
pub canceled_by: &'a mut Option<CanceledBy>,
pub client: &'a AuthedClientBackgroundTask,
pub client: &'a AuthedClient,
pub parent_runnable_path: Option<String>,
pub db: &'a sqlx::Pool<sqlx::Postgres>,
pub envs: HashMap<String, String>,
pub inner_content: &'a str,
@@ -221,14 +221,15 @@ async fn run<'a>(
job_dir,
shared_mount,
client,
parent_runnable_path,
envs,
base_internal_url,
..
}: &mut JobHandlerInput<'a>,
// plugins: Vec<&'a str>,
) -> Result<(), Error> {
let client = &client.get_authed().await;
let reserved_variables = get_reserved_variables(job, &client.token, db).await?;
let reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path.clone()).await?;
let child = if !cfg!(windows) && !*DISABLE_NSJAIL {
append_logs(
&job.id,
@@ -23,7 +23,7 @@ use crate::{
common::{build_args_values, check_executor_binary_exists, OccupancyMetrics},
handle_child::run_future_with_polling_update_job_poller,
sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args,
AuthedClientBackgroundTask,
AuthedClient,
};
#[derive(Deserialize)]
@@ -233,7 +233,8 @@ fn get_statement_values(
}
let arg_t = arg.otyp.clone().unwrap_or_else(|| "text".to_string());
let arg_n = arg.name.clone();
let oracle_v: Box<dyn ToSql + Send + Sync> = match job_args.get(arg.name.as_str())
let oracle_v: Box<dyn ToSql + Send + Sync> = match job_args
.get(arg.name.as_str())
.unwrap_or_else(|| &json!(null))
{
// Value::Null => todo!(),
@@ -293,7 +294,7 @@ fn get_statement_values(
pub async fn do_oracledb(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -315,8 +316,6 @@ pub async fn do_oracledb(
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
+3 -6
View File
@@ -37,7 +37,7 @@ use windmill_queue::{CanceledBy, MiniPulledJob};
use crate::common::{build_args_values, sizeof_val, OccupancyMetrics};
use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
use crate::{AuthedClientBackgroundTask, MAX_RESULT_SIZE};
use crate::{AuthedClient, MAX_RESULT_SIZE};
use bytes::Buf;
use lazy_static::lazy_static;
use urlencoding::encode;
@@ -159,7 +159,7 @@ fn do_postgresql_inner<'a>(
pub async fn do_postgresql(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -175,8 +175,6 @@ pub async fn do_postgresql(
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
@@ -293,8 +291,7 @@ pub async fn do_postgresql(
let sig = parse_pgsql_sig(&query).map_err(|x| Error::ExecutionErr(x.to_string()))?;
let (query, _) =
&sanitize_and_interpolate_unsafe_sql_args(query, &sig.args, &pg_args)?;
let (query, _) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig.args, &pg_args)?;
let queries = parse_sql_blocks(query);
+6 -5
View File
@@ -20,8 +20,8 @@ use crate::{
read_result, start_child_process, OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, COMPOSER_CACHE_DIR, COMPOSER_PATH, DISABLE_NSJAIL, DISABLE_NUSER,
NSJAIL_PATH, PHP_PATH,
AuthedClient, COMPOSER_CACHE_DIR, COMPOSER_PATH, DISABLE_NSJAIL, DISABLE_NUSER, NSJAIL_PATH,
PHP_PATH,
};
const NSJAIL_CONFIG_RUN_PHP_CONTENT: &str = include_str!("../nsjail/run.php.config.proto");
@@ -139,7 +139,8 @@ pub async fn handle_php_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
job_dir: &str,
inner_content: &String,
base_internal_url: &str,
@@ -265,8 +266,8 @@ try {{
Ok(()) as Result<()>
};
let reserved_variables_f = async {
let client = client.get_authed().await;
let vars = get_reserved_variables(job, &client.token, db).await?;
let vars = get_reserved_variables(job, &client.token, db, parent_runnable_path.clone())
.await?;
Ok(vars) as Result<HashMap<String, String>>
};
let (_, reserved_variables) = tokio::try_join!(args_and_out_f, reserved_variables_f)?;
@@ -76,9 +76,8 @@ use crate::{
start_child_process, OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, INSTANCE_PYTHON_VERSION,
NSJAIL_PATH, PATH_ENV, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, PROXY_ENVS, PY_INSTALL_DIR, TZ_ENV,
UV_CACHE_DIR,
AuthedClient, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, INSTANCE_PYTHON_VERSION, NSJAIL_PATH,
PATH_ENV, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, PROXY_ENVS, PY_INSTALL_DIR, TZ_ENV, UV_CACHE_DIR,
};
// To change latest stable version:
@@ -839,7 +838,8 @@ pub async fn handle_python_job(
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
inner_content: &String,
shared_mount: &str,
base_internal_url: &str,
@@ -1024,8 +1024,8 @@ except BaseException as e:
tracing::debug!("Finished writing wrapper");
let client = client.get_authed().await;
let mut reserved_variables = get_reserved_variables(job, &client.token, db).await?;
let mut reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
// Add /tmp/windmill/cache/python_xyz/global-site-packages to PYTHONPATH.
// Usefull if certain wheels needs to be preinstalled before execution.
@@ -235,7 +235,7 @@ async fn send_job_completed(
canceled_by: Option<CanceledBy>,
success: bool,
cached_res_path: Option<String>,
token: String,
token: &str,
duration: Option<i64>,
) {
let jc = JobCompleted {
@@ -246,7 +246,7 @@ async fn send_job_completed(
canceled_by,
success,
cached_res_path,
token,
token: token.to_string(),
duration,
};
job_completed_tx
@@ -264,7 +264,7 @@ pub async fn process_result(
mem_peak: i32,
canceled_by: Option<CanceledBy>,
cached_res_path: Option<String>,
token: String,
token: &str,
column_order: Option<Vec<String>>,
new_args: Option<HashMap<String, Box<RawValue>>>,
db: &DB,
+6 -5
View File
@@ -19,8 +19,8 @@ use crate::{
read_result, start_child_process, OccupancyMetrics,
},
handle_child::handle_child,
AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NSJAIL_PATH, PATH_ENV,
PROXY_ENVS, RUST_CACHE_DIR, TZ_ENV,
AuthedClient, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, NSJAIL_PATH, PATH_ENV, PROXY_ENVS,
RUST_CACHE_DIR, TZ_ENV,
};
#[cfg(windows)]
@@ -277,7 +277,8 @@ pub async fn handle_rust_job(
canceled_by: &mut Option<CanceledBy>,
job: &MiniPulledJob,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
inner_content: &str,
job_dir: &str,
requirements_o: Option<&String>,
@@ -343,8 +344,8 @@ pub async fn handle_rust_job(
let logs2 = format!("{cache_logs}\n\n--- RUST CODE EXECUTION ---\n");
append_logs(&job.id, &job.workspace_id, logs2, db).await;
let client = &client.get_authed().await;
let reserved_variables = get_reserved_variables(job, &client.token, db).await?;
let reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
let child = if !*DISABLE_NSJAIL {
let _ = write_file(
@@ -10,7 +10,10 @@ use windmill_parser_sql::{SANITIZED_ENUM_STR, SANITIZED_RAW_STRING_STR};
/// a number, that can contain underscores
fn sanitize_identifier(arg: &Arg, input: &str) -> Result<(), error::Error> {
if input.is_empty() {
return Err(error::Error::BadRequest(format!("Interpolated argument `{}` cannot be empty", arg.name)));
return Err(error::Error::BadRequest(format!(
"Interpolated argument `{}` cannot be empty",
arg.name
)));
}
if input
.chars()
@@ -19,7 +19,7 @@ use serde::{Deserialize, Serialize};
use crate::common::{build_http_client, resolve_job_timeout, OccupancyMetrics};
use crate::handle_child::run_future_with_polling_update_job_poller;
use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
use crate::{common::build_args_values, AuthedClientBackgroundTask};
use crate::{common::build_args_values, AuthedClient};
#[derive(Serialize)]
struct Claims {
@@ -247,7 +247,7 @@ fn do_snowflake_inner<'a>(
pub async fn do_snowflake(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
query: &str,
db: &sqlx::Pool<sqlx::Postgres>,
mem_peak: &mut i32,
@@ -263,8 +263,6 @@ pub async fn do_snowflake(
let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
Some(
client
.get_authed()
.await
.get_resource_value_interpolated::<serde_json::Value>(
&inline_db_res_path,
Some(job.id.to_string()),
+143 -108
View File
@@ -160,20 +160,10 @@ use crate::bench::{benchmark_init, BenchmarkInfo, BenchmarkIter};
use windmill_common::add_time;
pub async fn create_token_for_owner_in_bg(
db: &Pool<Postgres>,
job: &MiniPulledJob,
) -> Arc<RwLock<String>> {
let rw_lock = Arc::new(RwLock::new(String::new()));
// struct Permission
pub async fn create_token(db: &DB, job: &MiniPulledJob, perms: Option<JobPerms>) -> String {
// skipping test runs
if job.workspace_id != "" {
let mut locked = rw_lock.clone().write_owned().await;
let db = db.clone();
let w_id = job.workspace_id.clone();
let owner = job.permissioned_as.clone();
let email = job.permissioned_as_email.clone();
let job_id = job.id.clone();
let label = if job.permissioned_as != format!("u/{}", job.created_by)
&& job.permissioned_as != job.created_by
{
@@ -181,23 +171,22 @@ pub async fn create_token_for_owner_in_bg(
} else {
"ephemeral-script".to_string()
};
tokio::spawn(async move {
let token = create_token_for_owner(
&db.clone(),
&w_id,
&owner,
&label,
*SCRIPT_TOKEN_EXPIRY,
&email,
&job_id,
)
.warn_after_seconds(5)
.await
.expect("could not create job token");
*locked = token;
});
};
return rw_lock;
create_token_for_owner(
&db,
&job.workspace_id,
&job.permissioned_as,
&label,
*SCRIPT_TOKEN_EXPIRY,
&job.permissioned_as_email,
&job.id,
perms,
)
.warn_after_seconds(5)
.await
.expect("could not create job token")
} else {
return "".to_string();
}
}
#[tracing::instrument(level = "trace", skip_all)]
@@ -209,21 +198,27 @@ pub async fn create_token_for_owner(
expires_in: u64,
email: &str,
job_id: &Uuid,
perms: Option<JobPerms>,
) -> error::Result<String> {
// TODO: Bad implementation. We should not have access to this DB here.
if let Some(token) = JOB_TOKEN.as_ref() {
return Ok(token.clone());
}
let job_authed = match sqlx::query_as!(
JobPerms,
"SELECT * FROM job_perms WHERE job_id = $1 AND workspace_id = $2",
job_id,
w_id
)
let job_perms = if perms.is_some() {
Ok(perms)
} else {
sqlx::query_as!(
JobPerms,
"SELECT email, username, is_admin, is_operator, groups, folders FROM job_perms WHERE job_id = $1 AND workspace_id = $2",
job_id,
w_id
)
.fetch_optional(db)
.await
{
};
let job_authed = match job_perms {
Ok(Some(jp)) => jp.into(),
_ => {
tracing::warn!("Could not get permissions for job {job_id} from job_perms table, getting permissions directly...");
@@ -460,25 +455,6 @@ pub const MAX_RESULT_SIZE: usize = 1024 * 1024 * 2; // 2MB
pub const INIT_SCRIPT_TAG: &str = "init_script";
pub struct AuthedClientBackgroundTask {
pub base_internal_url: String,
pub workspace: String,
pub token: Arc<RwLock<String>>,
}
impl AuthedClientBackgroundTask {
pub async fn get_authed(&self) -> AuthedClient {
return AuthedClient {
base_internal_url: self.base_internal_url.clone(),
workspace: self.workspace.clone(),
token: self.get_token().await,
force_client: None,
};
}
pub async fn get_token(&self) -> String {
return self.token.read().await.clone();
}
}
#[derive(Clone)]
pub struct AuthedClient {
pub base_internal_url: String,
@@ -1322,35 +1298,43 @@ pub async fn run_worker(
v2_job.created_by,
v2_job_queue.started_at,
scheduled_for,
runnable_path,
kind,
runnable_id,
canceled_reason,
canceled_by,
permissioned_as,
permissioned_as_email,
flow_status,
v2_job.runnable_path,
v2_job.kind,
v2_job.runnable_id,
v2_job_queue.canceled_reason,
v2_job_queue.canceled_by,
v2_job.permissioned_as,
v2_job.permissioned_as_email,
v2_job_status.flow_status,
v2_job.tag,
script_lang,
same_worker,
pre_run_error,
concurrent_limit,
concurrency_time_window_s,
flow_innermost_root_job,
timeout,
flow_step_id,
cache_ttl,
v2_job.script_lang,
v2_job.same_worker,
v2_job.pre_run_error,
v2_job.concurrent_limit,
v2_job.concurrency_time_window_s,
v2_job.flow_innermost_root_job,
v2_job.timeout,
v2_job.flow_step_id,
v2_job.cache_ttl,
v2_job_queue.priority,
preprocessed,
script_entrypoint_override,
trigger,
trigger_kind,
visible_to_owner,
raw_code,
raw_lock,
raw_flow
FROM v2_job_queue INNER JOIN v2_job ON v2_job.id = v2_job_queue.id LEFT JOIN v2_job_status ON v2_job_status.id = v2_job_queue.id WHERE v2_job_queue.id = $1
",
v2_job.preprocessed,
v2_job.script_entrypoint_override,
v2_job.trigger,
v2_job.trigger_kind,
v2_job.visible_to_owner,
v2_job.raw_code,
v2_job.raw_lock,
v2_job.raw_flow,
pj.runnable_path as parent_runnable_path,
p.email as permissioned_as_email, p.username as permissioned_as_username, p.is_admin as permissioned_as_is_admin,
p.is_operator as permissioned_as_is_operator, p.groups as permissioned_as_groups, p.folders as permissioned_as_folders
FROM v2_job_queue
INNER JOIN v2_job ON v2_job.id = v2_job_queue.id
LEFT JOIN v2_job_status ON v2_job_status.id = v2_job_queue.id
LEFT JOIN job_perms p ON p.job_id = v2_job.id
LEFT JOIN v2_job pj ON v2_job.parent_job = pj.id
WHERE v2_job_queue.id = $1
",
)
.bind(same_worker_job.job_id)
.fetch_optional(db)
@@ -1526,7 +1510,6 @@ pub async fn run_worker(
.expect("send job completed END");
add_time!(bench, "sent job completed");
} else {
let token = create_token_for_owner_in_bg(&db, &job).await;
add_outstanding_wait_time(&job, db, OUTSTANDING_WAIT_TIME_THRESHOLD_MS);
#[cfg(feature = "prometheus")]
@@ -1627,17 +1610,57 @@ pub async fn run_worker(
.expect("could not create shared dir");
}
let authed_client = AuthedClientBackgroundTask {
base_internal_url: base_internal_url.to_string(),
token,
workspace: job.workspace_id.to_string(),
};
#[cfg(feature = "prometheus")]
let tag = job.tag.clone();
let is_init_script: bool = job.tag.as_str() == INIT_SCRIPT_TAG;
let PulledJob { job, raw_code, raw_lock, raw_flow } = job;
let PulledJob {
job,
raw_code,
raw_lock,
raw_flow,
parent_runnable_path,
permissioned_as_email,
permissioned_as_username,
permissioned_as_is_admin,
permissioned_as_is_operator,
permissioned_as_groups,
permissioned_as_folders,
} = job;
let job_perms = match (
permissioned_as_email,
permissioned_as_username,
permissioned_as_is_admin,
permissioned_as_is_operator,
permissioned_as_groups,
permissioned_as_folders,
) {
(
Some(email),
Some(username),
Some(is_admin),
Some(is_operator),
Some(groups),
Some(folders),
) => Some(JobPerms {
email,
username,
is_admin,
is_operator,
groups,
folders,
}),
_ => None,
};
let token = create_token(&db, &job, job_perms).await;
let authed_client = AuthedClient {
base_internal_url: base_internal_url.to_string(),
token,
workspace: job.workspace_id.to_string(),
force_client: None,
};
let arc_job = Arc::new(job);
add_time!(bench, "handle_queued_job START");
@@ -1678,6 +1701,7 @@ pub async fn run_worker(
raw_code,
raw_lock,
raw_flow,
parent_runnable_path,
db,
&authed_client,
&hostname,
@@ -1698,7 +1722,7 @@ pub async fn run_worker(
Err(err) => {
handle_job_error(
db,
&authed_client.get_authed().await,
&authed_client,
arc_job.as_ref(),
0,
None,
@@ -1923,7 +1947,7 @@ pub struct JobCompleted {
async fn do_nativets(
job: &MiniPulledJob,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
env_code: String,
code: String,
db: &Pool<Postgres>,
@@ -1968,8 +1992,9 @@ async fn handle_queued_job(
raw_code: Option<String>,
raw_lock: Option<String>,
raw_flow: Option<Json<Box<RawValue>>>,
parent_runnable_path: Option<String>,
db: &DB,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
hostname: &str,
worker_name: &str,
worker_dir: &str,
@@ -2081,17 +2106,15 @@ async fn handle_queued_job(
};
let cached_res_path = if job.cache_ttl.is_some() {
Some(cached_result_path(db, &client.get_authed().await, &job, preview_data.as_ref()).await)
Some(cached_result_path(db, &client, &job, preview_data.as_ref()).await)
} else {
None
};
if let Some(cached_res_path) = cached_res_path.as_ref() {
let authed_client = client.get_authed().await;
let cached_result_maybe = get_cached_resource_value_if_valid(
db,
&authed_client,
&client,
&job.id,
&job.workspace_id,
&cached_res_path,
@@ -2113,7 +2136,7 @@ async fn handle_queued_job(
canceled_by: None,
success: true,
cached_res_path: None,
token: authed_client.token,
token: client.token.clone(),
duration: None,
})
.await
@@ -2132,7 +2155,7 @@ async fn handle_queued_job(
job,
&flow_data,
db,
&client.get_authed().await,
&client,
None,
same_worker_tx,
worker_dir,
@@ -2193,7 +2216,7 @@ async fn handle_queued_job(
worker_name,
worker_dir,
base_internal_url,
&client.get_token().await,
&client.token,
occupancy_metrics,
)
.await
@@ -2209,7 +2232,7 @@ async fn handle_queued_job(
worker_name,
worker_dir,
base_internal_url,
&client.get_token().await,
&client.token,
occupancy_metrics,
)
.await
@@ -2223,7 +2246,7 @@ async fn handle_queued_job(
worker_name,
worker_dir,
base_internal_url,
&client.get_token().await,
&client.token,
occupancy_metrics,
)
.await
@@ -2246,6 +2269,7 @@ async fn handle_queued_job(
preview_data,
db,
client,
parent_runnable_path,
job_dir,
worker_dir,
&mut mem_peak,
@@ -2284,7 +2308,7 @@ async fn handle_queued_job(
mem_peak,
canceled_by,
cached_res_path,
client.get_token().await,
&client.token,
column_order,
new_args,
db,
@@ -2421,9 +2445,7 @@ async fn try_validate_schema(
};
let sv = match job.runnable_id {
Some(hash)
if job.kind != JobKind::Preview && job.kind != JobKind::FlowPreview =>
{
Some(hash) if job.kind != JobKind::Preview && job.kind != JobKind::FlowPreview => {
sv_fut.cached(validators_cache, (sub_key, hash)).await?
}
_ => sv_fut.await?,
@@ -2464,7 +2486,8 @@ async fn handle_code_execution_job(
job: &MiniPulledJob,
preview: Option<Arc<ScriptData>>,
db: &sqlx::Pool<sqlx::Postgres>,
client: &AuthedClientBackgroundTask,
client: &AuthedClient,
parent_runnable_path: Option<String>,
job_dir: &str,
#[allow(unused_variables)] worker_dir: &str,
mem_peak: &mut i32,
@@ -2754,7 +2777,8 @@ async fn handle_code_execution_job(
)
.await;
let reserved_variables = get_reserved_variables(job, &client.get_token().await, db).await?;
let reserved_variables =
get_reserved_variables(job, &client.token, db, parent_runnable_path).await?;
let env_code = format!(
"const process = {{ env: {{}} }};\nconst BASE_URL = '{base_internal_url}';\nconst BASE_INTERNAL_URL = '{base_internal_url}';\nprocess.env['BASE_URL'] = BASE_URL;process.env['BASE_INTERNAL_URL'] = BASE_INTERNAL_URL;\n{}",
@@ -2839,6 +2863,7 @@ mount {{
canceled_by,
db,
client,
parent_runnable_path,
&code,
&shared_mount,
base_internal_url,
@@ -2856,6 +2881,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
job_dir,
&code,
base_internal_url,
@@ -2875,6 +2901,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
job_dir,
&code,
base_internal_url,
@@ -2893,6 +2920,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
&code,
job_dir,
lock.as_ref(),
@@ -2911,6 +2939,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
&code,
job_dir,
&shared_mount,
@@ -2929,6 +2958,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
&code,
job_dir,
&shared_mount,
@@ -2953,6 +2983,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
job_dir,
&code,
base_internal_url,
@@ -2976,6 +3007,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
&code,
job_dir,
lock.as_ref(),
@@ -3004,6 +3036,7 @@ mount {{
canceled_by,
db,
client,
parent_runnable_path,
&code,
&shared_mount,
base_internal_url,
@@ -3019,6 +3052,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
&code,
job_dir,
lock.as_ref(),
@@ -3043,6 +3077,7 @@ mount {{
job,
db,
client,
parent_runnable_path,
inner_content: &code,
job_dir,
requirements_o: lock.as_ref(),
+6 -4
View File
@@ -2651,7 +2651,7 @@ async fn push_next_flow_job(
{
sqlx::query_as!(
JobPerms,
"SELECT * FROM job_perms WHERE job_id = $1 AND workspace_id = $2",
"SELECT email, username, is_admin, is_operator, groups, folders FROM job_perms WHERE job_id = $1 AND workspace_id = $2",
root_job,
flow_job.workspace_id,
)
@@ -2940,15 +2940,17 @@ async fn push_next_flow_job(
.execute(&mut *tx)
.await?;
tx.commit().warn_after_seconds(3).await?;
tracing::info!(id = %flow_job.id, root_id = %job_root, "all next flow jobs pushed: {uuids:?}");
if continue_on_same_worker {
if !is_one_uuid {
return Err(Error::BadRequest(
"Cannot continue on same worker with multiple jobs, parallel cannot be used in conjunction with same_worker".to_string(),
));
}
}
tx.commit().warn_after_seconds(3).await?;
tracing::info!(id = %flow_job.id, root_id = %job_root, "all next flow jobs pushed: {uuids:?}");
if continue_on_same_worker {
same_worker_tx
.send(SameWorkerPayload { job_id: first_uuid, recoverable: true })
.await