From 3761de874e03e7ade952925ebd57b75b92a4aa0e Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 20 Feb 2024 01:56:04 +0100 Subject: [PATCH] fix: prevent bigquery/snowflake against abuse timeout --- .../windmill-worker/src/bigquery_executor.rs | 190 ++++++++++-------- .../windmill-worker/src/graphql_executor.rs | 7 +- backend/windmill-worker/src/mssql_executor.rs | 58 ++++-- .../windmill-worker/src/snowflake_executor.rs | 117 ++++++----- backend/windmill-worker/src/worker.rs | 33 ++- 5 files changed, 250 insertions(+), 155 deletions(-) diff --git a/backend/windmill-worker/src/bigquery_executor.rs b/backend/windmill-worker/src/bigquery_executor.rs index 2eef2be433..50c7e3c156 100644 --- a/backend/windmill-worker/src/bigquery_executor.rs +++ b/backend/windmill-worker/src/bigquery_executor.rs @@ -1,11 +1,14 @@ +use futures::TryFutureExt; use serde_json::{json, value::RawValue, Value}; +use windmill_common::error::to_anyhow; use windmill_common::jobs::QueuedJob; use windmill_common::{error::Error, worker::to_raw_value}; use windmill_parser_sql::parse_bigquery_sig; -use windmill_queue::HTTP_CLIENT; +use windmill_queue::{CanceledBy, HTTP_CLIENT}; use serde::Deserialize; +use crate::common::run_future_with_polling_update_job_poller; use crate::{ common::{build_args_values, resolve_job_timeout}, AuthedClientBackgroundTask, @@ -59,6 +62,9 @@ pub async fn do_bigquery( client: &AuthedClientBackgroundTask, query: &str, db: &sqlx::Pool, + mem_peak: &mut i32, + canceled_by: &mut Option, + worker_name: &str, ) -> windmill_common::error::Result> { let bigquery_args = build_args_values(job, client, db).await?; @@ -137,103 +143,119 @@ pub async fn do_bigquery( ) .unwrap_or(200000); - let response = HTTP_CLIENT - .post( - "https://bigquery.googleapis.com/bigquery/v2/projects/".to_string() - + authentication_manager - .project_id() - .await - .map_err(|e| Error::ExecutionErr(e.to_string()))? - .as_str() - + "/queries", - ) - .bearer_auth(token.as_str()) - .json(&json!({ - "query": query, - "useLegacySql": false, - "maxResults": 10000, - "timeoutMs": timeout_ms, - "queryParameters": statement_values, - })) - .send() - .await - .map_err(|e| Error::ExecutionErr(format!("Could not send query to BigQuery API: {}", e)))?; - - match response.error_for_status_ref() { - Ok(_) => { - let result = response.json::().await.map_err(|e| { - Error::ExecutionErr(format!( - "BigQuery API response could not be parsed: {}", - e.to_string() - )) + let result_f = async { + let response = HTTP_CLIENT + .post( + "https://bigquery.googleapis.com/bigquery/v2/projects/".to_string() + + authentication_manager + .project_id() + .await + .map_err(|e| Error::ExecutionErr(e.to_string()))? + .as_str() + + "/queries", + ) + .bearer_auth(token.as_str()) + .json(&json!({ + "query": query, + "useLegacySql": false, + "maxResults": 10000, + "timeoutMs": timeout_ms, + "queryParameters": statement_values, + })) + .send() + .await + .map_err(|e| { + Error::ExecutionErr(format!("Could not send query to BigQuery API: {}", e)) })?; - if !result.jobComplete { - return Err(Error::ExecutionErr( - "BigQuery API did not answer query in time".to_string(), - )); - } + match response.error_for_status_ref() { + Ok(_) => { + let result = response.json::().await.map_err(|e| { + Error::ExecutionErr(format!( + "BigQuery API response could not be parsed: {}", + e.to_string() + )) + })?; - if result.rows.is_none() || result.rows.as_ref().unwrap().len() == 0 { - return Ok(serde_json::from_str("[]").unwrap()); - } + if !result.jobComplete { + return Err(Error::ExecutionErr( + "BigQuery API did not answer query in time".to_string(), + )); + } - if result.schema.is_none() { - return Err(Error::ExecutionErr( - "Incomplete response from BigQuery API".to_string(), - )); - } + if result.rows.is_none() || result.rows.as_ref().unwrap().len() == 0 { + return Ok(serde_json::from_str("[]").unwrap()); + } - if result - .totalRows - .unwrap_or(json!("")) - .as_str() - .unwrap_or("") - .parse::() - .unwrap_or(0) - > 10000 - { - return Err(Error::ExecutionErr( + if result.schema.is_none() { + return Err(Error::ExecutionErr( + "Incomplete response from BigQuery API".to_string(), + )); + } + + if result + .totalRows + .unwrap_or(json!("")) + .as_str() + .unwrap_or("") + .parse::() + .unwrap_or(0) + > 10000 + { + return Err(Error::ExecutionErr( "More than 10000 rows were requested, use LIMIT 10000 to limit the number of rows".to_string(), )); + } + + let rows = result + .rows + .unwrap() + .iter() + .map(|row| { + let mut row_map = serde_json::Map::new(); + row.f + .iter() + .zip(result.schema.as_ref().unwrap().fields.iter()) + .for_each(|(field, schema)| { + row_map.insert( + schema.name.clone(), + parse_val(&field.v, &schema.r#type, &schema), + ); + }); + Value::from(row_map) + }) + .collect::>(); + + return Ok(to_raw_value(&rows)); } - - let rows = result - .rows - .unwrap() - .iter() - .map(|row| { - let mut row_map = serde_json::Map::new(); - row.f - .iter() - .zip(result.schema.as_ref().unwrap().fields.iter()) - .for_each(|(field, schema)| { - row_map.insert( - schema.name.clone(), - parse_val(&field.v, &schema.r#type, &schema), - ); - }); - Value::from(row_map) - }) - .collect::>(); - - return Ok(to_raw_value(&rows)); - } - Err(e) => match response.json::().await { - Ok(bq_err) => { - return Err(Error::ExecutionErr(format!( + Err(e) => match response.json::().await { + Ok(bq_err) => Err(Error::ExecutionErr(format!( "Error from BigQuery API: {}", bq_err.error.message ))) - } - Err(_) => { - return Err(Error::ExecutionErr(format!( + .map_err(to_anyhow)?, + Err(_) => Err(Error::ExecutionErr(format!( "Error from BigQuery API could not be parsed: {}", e.to_string() ))) - } - }, - } + .map_err(to_anyhow)?, + }, + } + }; + let r = run_future_with_polling_update_job_poller( + job.id, + job.timeout, + db, + mem_peak, + canceled_by, + result_f.map_err(to_anyhow), + worker_name, + &job.workspace_id, + ) + .await?; + + *mem_peak = (r.get().len() / 1000) as i32; + Ok(r) } fn convert_val(arg_t: String, arg_v: Value) -> Value { diff --git a/backend/windmill-worker/src/graphql_executor.rs b/backend/windmill-worker/src/graphql_executor.rs index e91fb5a233..1663b69edd 100644 --- a/backend/windmill-worker/src/graphql_executor.rs +++ b/backend/windmill-worker/src/graphql_executor.rs @@ -134,7 +134,7 @@ pub async fn do_graphql( .unwrap_or_else(|| serde_json::from_str("{}").unwrap())) }; - Ok(run_future_with_polling_update_job_poller( + let r = run_future_with_polling_update_job_poller( job.id, job.timeout, db, @@ -144,5 +144,8 @@ pub async fn do_graphql( worker_name, &job.workspace_id, ) - .await?) + .await?; + + *mem_peak = (r.get().len() / 1000) as i32; + Ok(r) } diff --git a/backend/windmill-worker/src/mssql_executor.rs b/backend/windmill-worker/src/mssql_executor.rs index b91f949089..d3ea0d964b 100644 --- a/backend/windmill-worker/src/mssql_executor.rs +++ b/backend/windmill-worker/src/mssql_executor.rs @@ -1,5 +1,6 @@ use base64::{engine::general_purpose, Engine as _}; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; +use futures::TryFutureExt; use serde::Deserialize; use serde_json::value::RawValue; use serde_json::{Map, Value}; @@ -11,8 +12,9 @@ use windmill_common::error::{self, Error}; use windmill_common::worker::to_raw_value; use windmill_common::{error::to_anyhow, jobs::QueuedJob}; use windmill_parser_sql::parse_mssql_sig; +use windmill_queue::CanceledBy; -use crate::common::build_args_values; +use crate::common::{build_args_values, run_future_with_polling_update_job_poller}; use crate::AuthedClientBackgroundTask; #[derive(Deserialize)] @@ -30,6 +32,9 @@ pub async fn do_mssql( client: &AuthedClientBackgroundTask, query: &str, db: &sqlx::Pool, + mem_peak: &mut i32, + canceled_by: &mut Option, + worker_name: &str, ) -> error::Result> { let mssql_args = build_args_values(job, client, db).await?; @@ -84,23 +89,40 @@ pub async fn do_mssql( json_value_to_sql(&mut prepared_query, &arg_v, &arg_t)?; } - // A response to a query is a stream of data, that must be - // polled to the end before querying again. Using streams allows - // fetching data in an asynchronous manner, if needed. - let stream = prepared_query.query(&mut client).await.map_err(to_anyhow)?; - let rows = stream - .into_results() - .await - .map_err(to_anyhow)? - .into_iter() - .map(|rows| { - let result = rows - .into_iter() - .map(|row| row_to_json(row)) - .collect::>, Error>>(); - result - }) - .collect::>>, Error>>()?; + let result_f = async { + // A response to a query is a stream of data, that must be + // polled to the end before querying again. Using streams allows + // fetching data in an asynchronous manner, if needed. + let stream = prepared_query.query(&mut client).await.map_err(to_anyhow)?; + stream + .into_results() + .await + .map_err(to_anyhow)? + .into_iter() + .map(|rows| { + let result = rows + .into_iter() + .map(|row| row_to_json(row)) + .collect::>, Error>>(); + result + }) + .collect::>>, Error>>() + }; + + let rows = run_future_with_polling_update_job_poller( + job.id, + job.timeout, + db, + mem_peak, + canceled_by, + result_f.map_err(to_anyhow), + worker_name, + &job.workspace_id, + ) + .await?; + + let r = to_raw_value(&rows); + *mem_peak = (r.get().len() / 1000) as i32; return Ok(to_raw_value(&rows)); } diff --git a/backend/windmill-worker/src/snowflake_executor.rs b/backend/windmill-worker/src/snowflake_executor.rs index df084fddcb..9c6a15fd8f 100644 --- a/backend/windmill-worker/src/snowflake_executor.rs +++ b/backend/windmill-worker/src/snowflake_executor.rs @@ -1,17 +1,20 @@ use base64::{engine, Engine as _}; use core::fmt::Write; +use futures::TryFutureExt; use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; use pem; use serde_json::{json, value::RawValue, Value}; use sha2::{Digest, Sha256}; +use windmill_common::error::to_anyhow; use windmill_common::jobs::QueuedJob; use windmill_common::{error::Error, worker::to_raw_value}; use windmill_parser_sql::parse_snowflake_sig; -use windmill_queue::HTTP_CLIENT; +use windmill_queue::{CanceledBy, HTTP_CLIENT}; use serde::{Deserialize, Serialize}; +use crate::common::run_future_with_polling_update_job_poller; use crate::{common::build_args_values, AuthedClientBackgroundTask}; #[derive(Serialize)] @@ -65,6 +68,9 @@ pub async fn do_snowflake( client: &AuthedClientBackgroundTask, query: &str, db: &sqlx::Pool, + mem_peak: &mut i32, + canceled_by: &mut Option, + worker_name: &str, ) -> windmill_common::error::Result> { let snowflake_args = build_args_values(job, client, db).await?; @@ -153,60 +159,75 @@ pub async fn do_snowflake( body.insert("bindings".to_string(), json!(bindings)); } - let response = HTTP_CLIENT - .post(format!( - "https://{}.snowflakecomputing.com/api/v2/statements/", - database.account_identifier.to_uppercase() - )) - .bearer_auth(token) - .header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT") - .json(&body) - .send() - .await - .map_err(|e| Error::ExecutionErr(e.to_string()))?; + let result_f = async { + let response = HTTP_CLIENT + .post(format!( + "https://{}.snowflakecomputing.com/api/v2/statements/", + database.account_identifier.to_uppercase() + )) + .bearer_auth(token) + .header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT") + .json(&body) + .send() + .await + .map_err(|e| Error::ExecutionErr(e.to_string()))?; - match response.error_for_status_ref() { - Ok(_) => { - let result = response - .json::() - .await - .map_err(|e| Error::ExecutionErr(e.to_string()))?; + match response.error_for_status_ref() { + Ok(_) => { + let result = response + .json::() + .await + .map_err(|e| Error::ExecutionErr(e.to_string()))?; - if result.resultSetMetaData.numRows > 10000 { - return Err(Error::ExecutionErr( + if result.resultSetMetaData.numRows > 10000 { + return Err(Error::ExecutionErr( "More than 10000 rows were requested, use LIMIT 10000 to limit the number of rows".to_string(), )); + } + + let rows = to_raw_value( + &result + .data + .iter() + .map(|row| { + let mut row_map = serde_json::Map::new(); + row.iter() + .zip(result.resultSetMetaData.rowType.iter()) + .for_each(|(val, row_type)| { + row_map.insert( + row_type.name.clone(), + parse_val(&val, &row_type.r#type), + ); + }); + Value::from(row_map) + }) + .collect::>(), + ); + + Ok(rows) } - - let rows = to_raw_value( - &result - .data - .iter() - .map(|row| { - let mut row_map = serde_json::Map::new(); - row.iter() - .zip(result.resultSetMetaData.rowType.iter()) - .for_each(|(val, row_type)| { - row_map.insert( - row_type.name.clone(), - parse_val(&val, &row_type.r#type), - ); - }); - Value::from(row_map) - }) - .collect::>(), - ); - - Ok(rows) - } - Err(e) => { - let resp = response.text().await.unwrap_or("".to_string()); - match serde_json::from_str::(&resp) { - Ok(sf_err) => Err(Error::ExecutionErr(sf_err.message)), - Err(_) => Err(Error::ExecutionErr(e.to_string())), + Err(e) => { + let resp = response.text().await.unwrap_or("".to_string()); + match serde_json::from_str::(&resp) { + Ok(sf_err) => Err(Error::ExecutionErr(sf_err.message)), + Err(_) => Err(Error::ExecutionErr(e.to_string())), + } } } - } + }; + let r = run_future_with_polling_update_job_poller( + job.id, + job.timeout, + db, + mem_peak, + canceled_by, + result_f.map_err(to_anyhow), + worker_name, + &job.workspace_id, + ) + .await?; + *mem_peak = (r.get().len() / 1000) as i32; + Ok(r) } fn convert_typ_val(arg_t: String, arg_v: Value) -> Value { diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 5d9551bb01..7bbb6c6ec9 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -2928,7 +2928,16 @@ async fn handle_code_execution_job( #[cfg(feature = "enterprise")] { - return do_bigquery(job, &client, &inner_content, db).await; + return do_bigquery( + job, + &client, + &inner_content, + db, + mem_peak, + canceled_by, + worker_name, + ) + .await; } } else if language == Some(ScriptLang::Snowflake) { #[cfg(not(feature = "enterprise"))] @@ -2940,7 +2949,16 @@ async fn handle_code_execution_job( #[cfg(feature = "enterprise")] { - return do_snowflake(job, &client, &inner_content, db).await; + return do_snowflake( + job, + &client, + &inner_content, + db, + mem_peak, + canceled_by, + worker_name, + ) + .await; } } else if language == Some(ScriptLang::Mssql) { #[cfg(not(feature = "enterprise"))] @@ -2952,7 +2970,16 @@ async fn handle_code_execution_job( #[cfg(feature = "enterprise")] { - return do_mssql(job, &client, &inner_content, db).await; + return do_mssql( + job, + &client, + &inner_content, + db, + mem_peak, + canceled_by, + worker_name, + ) + .await; } } else if language == Some(ScriptLang::Graphql) { return do_graphql(