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

fix: prevent bigquery/snowflake against abuse timeout

Browse Source
This commit is contained in:
Ruben Fiszel
2024-02-20 01:56:04 +01:00
parent 66fc78f233
commit 3761de874e
5 changed files with 250 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/windmill-worker/src/bigquery_executor.rs
+106 -84
View File
@@ -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<sqlx::Postgres>,
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
worker_name: &str,
) -> windmill_common::error::Result<Box<RawValue>> {
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::<BigqueryResponse>().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::<BigqueryResponse>().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::<i64>()
.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::<i64>()
.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::<Vec<_>>();
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::<Vec<_>>();
return Ok(to_raw_value(&rows));
}
Err(e) => match response.json::<BigqueryErrorResponse>().await {
Ok(bq_err) => {
return Err(Error::ExecutionErr(format!(
Err(e) => match response.json::<BigqueryErrorResponse>().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 {
backend/windmill-worker/src/graphql_executor.rs
+5 -2
View File
@@ -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)
}
backend/windmill-worker/src/mssql_executor.rs
+40 -18
View File
@@ -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<sqlx::Postgres>,
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
worker_name: &str,
) -> error::Result<Box<RawValue>> {
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::<Result<Vec<Map<String, Value>>, Error>>();
result
})
.collect::<Result<Vec<Vec<Map<String, Value>>>, 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::<Result<Vec<Map<String, Value>>, Error>>();
result
})
.collect::<Result<Vec<Vec<Map<String, Value>>>, 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));
}
backend/windmill-worker/src/snowflake_executor.rs
+69 -48
View File
@@ -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<sqlx::Postgres>,
mem_peak: &mut i32,
canceled_by: &mut Option<CanceledBy>,
worker_name: &str,
) -> windmill_common::error::Result<Box<RawValue>> {
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::<SnowflakeResponse>()
.await
.map_err(|e| Error::ExecutionErr(e.to_string()))?;
match response.error_for_status_ref() {
Ok(_) => {
let result = response
.json::<SnowflakeResponse>()
.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::<Vec<_>>(),
);
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::<Vec<_>>(),
);
Ok(rows)
}
Err(e) => {
let resp = response.text().await.unwrap_or("".to_string());
match serde_json::from_str::<SnowflakeError>(&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::<SnowflakeError>(&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 {
backend/windmill-worker/src/worker.rs
+30 -3
View File
@@ -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(