diff --git a/backend/Cargo.lock b/backend/Cargo.lock index f6730c98c4..02a5178079 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -10689,7 +10689,6 @@ dependencies = [ name = "windmill-sql-datatype-parser-wasm" version = "1.367.2" dependencies = [ - "serde", "wasm-bindgen", "wasm-bindgen-test", "windmill-parser", diff --git a/backend/parsers/windmill-parser-bash/src/lib.rs b/backend/parsers/windmill-parser-bash/src/lib.rs index 5eeb854700..fa5eb3973d 100644 --- a/backend/parsers/windmill-parser-bash/src/lib.rs +++ b/backend/parsers/windmill-parser-bash/src/lib.rs @@ -58,6 +58,7 @@ fn parse_bash_file(code: &str) -> anyhow::Result>> { default: default.clone().map(|x| json!(x)), otyp: None, has_default: default.is_some(), + oidx: None, }); } else { break; @@ -94,6 +95,7 @@ fn parse_powershell_file(code: &str) -> anyhow::Result>> { default: default.clone(), otyp: None, has_default: default.is_some(), + oidx: None, }); } } @@ -130,35 +132,40 @@ non_required="${5:-}" name: "token".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "image".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "digest".to_string(), typ: Typ::Str(None), default: Some(json!("latest with spaces")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "text".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "non_required".to_string(), typ: Typ::Str(None), default: Some(json!("")), - has_default: true + has_default: true, + oidx: None } ], no_main_func: None diff --git a/backend/parsers/windmill-parser-go/src/lib.rs b/backend/parsers/windmill-parser-go/src/lib.rs index 4e86ae2cb5..c2564dbbd2 100644 --- a/backend/parsers/windmill-parser-go/src/lib.rs +++ b/backend/parsers/windmill-parser-go/src/lib.rs @@ -27,7 +27,14 @@ pub fn parse_go_sig(code: &str) -> anyhow::Result { .iter() .map(|param| { let (otyp, typ) = parse_go_typ(¶m.typ); - Arg { name: get_name(param), otyp, typ, default: None, has_default: false } + Arg { + name: get_name(param), + otyp, + typ, + default: None, + has_default: false, + oidx: None, + } }) .collect_vec(); Ok(MainArgSignature { @@ -179,28 +186,32 @@ func main(x int, y string, z bool, l []string, o struct { Name string `json:"nam name: "x".to_string(), typ: Typ::Int, has_default: false, - default: None + default: None, + oidx: None }, Arg { otyp: Some("string".to_string()), name: "y".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: Some("bool".to_string()), name: "z".to_string(), typ: Typ::Bool, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: Some("[]string".to_string()), name: "l".to_string(), typ: Typ::List(Box::new(Typ::Str(None))), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: Some("struct { Name string `json:\"name\"` }".to_string()), @@ -210,21 +221,24 @@ func main(x int, y string, z bool, l []string, o struct { Name string `json:"nam typ: Box::new(Typ::Str(None)) },]), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: Some("interface{}".to_string()), name: "n".to_string(), typ: Typ::Object(vec![]), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: Some("map[string]interface{}".to_string()), name: "m".to_string(), typ: Typ::Object(vec![]), default: None, - has_default: false + has_default: false, + oidx: None }, ], no_main_func: Some(false) diff --git a/backend/parsers/windmill-parser-graphql/src/lib.rs b/backend/parsers/windmill-parser-graphql/src/lib.rs index 2ab6872762..99eb8db261 100644 --- a/backend/parsers/windmill-parser-graphql/src/lib.rs +++ b/backend/parsers/windmill-parser-graphql/src/lib.rs @@ -51,6 +51,7 @@ fn parse_graphql_file(code: &str) -> anyhow::Result>> { default: parsed_default, otyp: Some(typ.unwrap()), has_default, + oidx: None, }); } @@ -93,21 +94,24 @@ query($i: Int, $arr: [String]!, $wahoo: String = "wahoo") { name: "i".to_string(), typ: Typ::Int, default: None, - has_default: true + has_default: true, + oidx: None }, Arg { otyp: Some("[String]".to_string()), name: "arr".to_string(), typ: Typ::List(Box::new(Typ::Str(None))), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: Some("String".to_string()), name: "wahoo".to_string(), typ: Typ::Str(None), default: Some(json!("wahoo")), - has_default: true + has_default: true, + oidx: None } ], no_main_func: None diff --git a/backend/parsers/windmill-parser-php/src/lib.rs b/backend/parsers/windmill-parser-php/src/lib.rs index 228fc78e75..27befc1ea7 100644 --- a/backend/parsers/windmill-parser-php/src/lib.rs +++ b/backend/parsers/windmill-parser-php/src/lib.rs @@ -85,6 +85,7 @@ pub fn parse_php_signature( typ, has_default: default.is_some(), default, + oidx: None, } }) .collect(); @@ -135,35 +136,40 @@ function main(string $input1 = \"hey\", bool $input2 = false, int $input3 = 3, f name: "input1".to_string(), typ: Typ::Str(None), has_default: true, - default: Some(Value::String("hey".to_string())) + default: Some(Value::String("hey".to_string())), + oidx: None }, Arg { otyp: None, name: "input2".to_string(), typ: Typ::Bool, has_default: true, - default: Some(Value::Bool(false)) + default: Some(Value::Bool(false)), + oidx: None }, Arg { otyp: None, name: "input3".to_string(), typ: Typ::Int, has_default: true, - default: Some(Value::Number(Number::from(3))) + default: Some(Value::Number(Number::from(3))), + oidx: None }, Arg { otyp: None, name: "input4".to_string(), typ: Typ::Float, has_default: true, - default: Some(Value::Number(Number::from_f64(f64::from(4.5)).unwrap())) + default: Some(Value::Number(Number::from_f64(f64::from(4.5)).unwrap())), + oidx: None }, Arg { otyp: None, name: "resource".to_string(), typ: Typ::Resource("stripe".to_string()), has_default: false, - default: None + default: None, + oidx: None } ], no_main_func: Some(false) diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index a7e3e431e1..3adde32185 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/backend/parsers/windmill-parser-py/src/lib.rs @@ -114,6 +114,7 @@ pub fn parse_python_signature( typ, has_default: default.is_some(), default, + oidx: None, } }) .collect(), @@ -271,49 +272,56 @@ def main(test1: str, name: datetime.datetime = datetime.now(), byte: bytes = byt name: "test1".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "name".to_string(), typ: Typ::Unknown, default: Some(json!("")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "byte".to_string(), typ: Typ::Bytes, default: Some(json!("")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "f".to_string(), typ: Typ::Str(None), default: Some(json!("wewe")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "g".to_string(), typ: Typ::Int, default: Some(json!(21)), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "h".to_string(), typ: Typ::List(Box::new(Typ::Int)), default: Some(json!([1, 2])), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "i".to_string(), typ: Typ::Bool, default: Some(json!(true)), - has_default: true + has_default: true, + oidx: None }, ], no_main_func: Some(false), @@ -352,28 +360,32 @@ def main(test1: str, name: "test1".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "name".to_string(), typ: Typ::Unknown, default: Some(json!("")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "byte".to_string(), typ: Typ::Bytes, default: Some(json!("")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "resource".to_string(), typ: Typ::Resource("postgresql".to_string()), default: Some(json!("$res:g/all/resource")), - has_default: true + has_default: true, + oidx: None } ], no_main_func: Some(false), @@ -407,28 +419,32 @@ def main(test1: str, name: "test1".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "s3o".to_string(), typ: Typ::Resource("S3Object".to_string()), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "name".to_string(), typ: Typ::Str(None), default: Some(json!("test")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "byte".to_string(), typ: Typ::Bytes, default: Some(json!("")), - has_default: true + has_default: true, + oidx: None } ], no_main_func: Some(false), @@ -459,7 +475,8 @@ def main(test1: Literal["foo", "bar"], test2: List[Literal["foo", "bar"]]): retu name: "test1".to_string(), typ: Typ::Str(Some(vec!["foo".to_string(), "bar".to_string()])), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, @@ -469,7 +486,8 @@ def main(test1: Literal["foo", "bar"], test2: List[Literal["foo", "bar"]]): retu "bar".to_string() ])))), default: None, - has_default: false + has_default: false, + oidx: None } ], no_main_func: Some(false), @@ -499,7 +517,8 @@ def main(test1: DynSelect_foo): return name: "test1".to_string(), typ: Typ::DynSelect("foo".to_string()), default: None, - has_default: false + has_default: false, + oidx: None }], no_main_func: Some(false), } diff --git a/backend/parsers/windmill-parser-sql/src/lib.rs b/backend/parsers/windmill-parser-sql/src/lib.rs index b883ed6b28..c97e291648 100644 --- a/backend/parsers/windmill-parser-sql/src/lib.rs +++ b/backend/parsers/windmill-parser-sql/src/lib.rs @@ -4,7 +4,11 @@ use anyhow::anyhow; use regex::Regex; use serde_json::json; -use std::collections::HashMap; +use std::{ + collections::{HashMap, HashSet}, + iter::Peekable, + str::CharIndices, +}; pub use windmill_parser::{Arg, MainArgSignature, Typ}; pub fn parse_mysql_sig(code: &str) -> anyhow::Result { @@ -62,9 +66,32 @@ pub fn parse_db_resource(code: &str) -> Option { cap.map(|x| x.get(1).map(|x| x.as_str().to_string()).unwrap()) } +pub fn parse_sql_blocks(code: &str) -> Vec<&str> { + let mut blocks = vec![]; + let mut last_idx = 0; + + run_on_sql_statement_matches( + code, + |char, _| char == ';', + |idx, _| { + blocks.push(&code[last_idx..=idx]); + last_idx = idx + 1; + }, + ); + if last_idx < code.len() { + let last_block = &code[last_idx..]; + if RE_NONEMPTY_SQL_BLOCK.is_match(last_block) { + blocks.push(last_block); + } + } + blocks +} + lazy_static::lazy_static! { static ref RE_CODE_PGSQL: Regex = Regex::new(r#"(?m)\$(\d+)(?:::(\w+(?:\[\])?))?"#).unwrap(); + static ref RE_NONEMPTY_SQL_BLOCK: Regex = Regex::new(r#"(?m)^\s*[^\s](?:[^-]|$)"#).unwrap(); + static ref RE_DB: Regex = Regex::new(r#"(?m)^-- database (\S+) *(?:\r|\n|$)"#).unwrap(); // -- $1 name (type) = default @@ -117,6 +144,7 @@ fn parse_mysql_file(code: &str) -> anyhow::Result>> { default: parsed_default, otyp: Some(typ), has_default, + oidx: None, }); } @@ -140,6 +168,7 @@ fn parse_mysql_file(code: &str) -> anyhow::Result>> { default: parsed_default, otyp: Some(typ), has_default, + oidx: None, }); } } @@ -147,7 +176,105 @@ fn parse_mysql_file(code: &str) -> anyhow::Result>> { Ok(Some(args)) } +enum ParserState { + Normal, + InSingleQuote, + InDoubleQuote, + InSingleLineComment, + InMultiLineComment, +} + +fn run_on_sql_statement_matches< + F1: FnMut(char, &mut Peekable) -> bool, + F2: FnMut(usize, &mut Peekable) -> (), +>( + code: &str, + mut cond: F1, + mut case: F2, +) { + let mut chars = code.char_indices().peekable(); + let mut state = ParserState::Normal; + while let Some((idx, char)) = chars.next() { + match (&state, char) { + (ParserState::Normal, '\'') => { + state = ParserState::InSingleQuote; + } + (ParserState::Normal, '"') => { + state = ParserState::InDoubleQuote; + } + (ParserState::Normal, '-') + if chars.peek().is_some_and(|&(_, next_char)| next_char == '-') => + { + state = ParserState::InSingleLineComment; + } + (ParserState::Normal, '/') + if chars.peek().is_some_and(|&(_, next_char)| next_char == '*') => + { + state = ParserState::InMultiLineComment; + } + (ParserState::Normal, _) if cond(char, &mut chars) => { + case(idx, &mut chars); + } + (ParserState::InSingleQuote, '\'') => { + if chars + .peek() + .is_some_and(|&(_, next_char)| next_char == '\'') + { + chars.next(); // skip the escaped single quote + } else { + state = ParserState::Normal; + } + } + (ParserState::InDoubleQuote, '"') => { + if chars.peek().is_some_and(|&(_, next_char)| next_char == '"') { + chars.next(); // skip the escaped single quote + } else { + state = ParserState::Normal; + } + } + (ParserState::InSingleLineComment, '\n') => { + state = ParserState::Normal; + } + (ParserState::InMultiLineComment, '*') + if chars.peek().is_some_and(|&(_, next_char)| next_char == '/') => + { + state = ParserState::Normal; + } + _ => {} + } + } +} + +pub fn parse_pg_statement_arg_indices(code: &str) -> HashSet { + let mut arg_indices = HashSet::new(); + run_on_sql_statement_matches( + code, + |char, chars| { + char == '$' + && chars + .peek() + .is_some_and(|&(_, next_char)| next_char.is_ascii_digit()) + }, + |_, chars| { + let mut arg_idx = String::new(); + while let Some(&(_, char)) = chars.peek() { + if char.is_ascii_digit() { + arg_idx.push(char); + chars.next(); + } else { + break; + } + } + if let Ok(arg_idx) = arg_idx.parse::() { + arg_indices.insert(arg_idx); + } + }, + ); + arg_indices +} + fn parse_pg_file(code: &str) -> anyhow::Result>> { + let mut args = vec![]; let mut hm: HashMap = HashMap::new(); for cap in RE_CODE_PGSQL.captures_iter(code) { hm.insert( @@ -159,39 +286,74 @@ fn parse_pg_file(code: &str) -> anyhow::Result>> { .unwrap_or_else(|| "text".to_string()), ); } - let mut args = vec![]; - for i in 1..50 { - if hm.contains_key(&i) { - let typ = hm.get(&i).unwrap().to_lowercase(); - args.push(Arg { - name: format!("${}", i), - typ: parse_pg_typ(typ.as_str()), - default: None, - otyp: Some(typ), - has_default: false, - }); - } else { - break; - } + for (i, v) in hm.iter() { + let typ = v.to_lowercase(); + args.push(Arg { + name: format!("${}", i), + typ: parse_pg_typ(typ.as_str()), + default: None, + otyp: Some(typ), + has_default: false, + oidx: Some(*i), + }); } + args.sort_by(|a, b| a.oidx.unwrap().cmp(&b.oidx.unwrap())); for cap in RE_ARG_PGSQL.captures_iter(code) { - let i = cap.get(1).and_then(|x| x.as_str().parse::().ok()); - if i.is_none() || i.unwrap() as usize > args.len() { - continue; - } - let name = cap.get(2).map(|x| x.as_str().to_string()).unwrap(); - let default = cap.get(3).map(|x| x.as_str().to_string()); - let has_default = default.is_some(); - let oarg = args[(i.unwrap() - 1) as usize].clone(); - let parsed_default = default.and_then(|x| parsed_default(&oarg.typ, x)); + let i = cap + .get(1) + .and_then(|x| x.as_str().parse::().ok()) + .map(|x| x); - args[(i.unwrap() - 1) as usize] = - Arg { name, typ: oarg.typ, default: parsed_default, otyp: oarg.otyp, has_default }; + if let Some(arg_pos) = args + .iter() + .position(|x| i.is_some_and(|i| x.oidx.unwrap() == i)) + { + let name = cap.get(2).map(|x| x.as_str().to_string()).unwrap(); + let default = cap.get(3).map(|x| x.as_str().to_string()); + let has_default = default.is_some(); + let oarg = args[arg_pos].clone(); + let parsed_default = default.and_then(|x| parsed_default(&oarg.typ, x)); + + args[arg_pos] = Arg { + name, + typ: oarg.typ, + default: parsed_default, + otyp: oarg.otyp, + has_default, + oidx: oarg.oidx, + }; + } } Ok(Some(args)) } +pub fn parse_sql_statement_named_params(code: &str, prefix: char) -> HashSet { + let mut arg_names = HashSet::new(); + run_on_sql_statement_matches( + code, + |char, chars| { + char == prefix + && chars + .peek() + .is_some_and(|&(_, next_char)| next_char.is_alphanumeric()) + }, + |_, chars| { + let mut arg_name = String::new(); + while let Some(&(_, char)) = chars.peek() { + if char.is_alphanumeric() { + arg_name.push(char); + chars.next(); + } else { + break; + } + } + arg_names.insert(arg_name); + }, + ); + arg_names +} + fn parse_bigquery_file(code: &str) -> anyhow::Result>> { let mut args: Vec = vec![]; @@ -213,6 +375,7 @@ fn parse_bigquery_file(code: &str) -> anyhow::Result>> { default: parsed_default, otyp: Some(typ), has_default, + oidx: None, }); } @@ -240,6 +403,7 @@ fn parse_snowflake_file(code: &str) -> anyhow::Result>> { default: parsed_default, otyp: Some(typ), has_default, + oidx: None, }); } @@ -267,6 +431,7 @@ fn parse_mssql_file(code: &str) -> anyhow::Result>> { default: parsed_default, otyp: Some(typ), has_default, + oidx: None, }); } @@ -372,7 +537,7 @@ mod tests { use super::*; #[test] - fn test_parse_sql_sig() -> anyhow::Result<()> { + fn test_parse_pgsql_sig() -> anyhow::Result<()> { let code = r#" SELECT * FROM table WHERE token=$1::TEXT AND image=$2::BIGINT "#; @@ -388,14 +553,237 @@ SELECT * FROM table WHERE token=$1::TEXT AND image=$2::BIGINT name: "$1".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: Some(1), }, Arg { otyp: Some("bigint".to_string()), name: "$2".to_string(), typ: Typ::Int, default: None, - has_default: false + has_default: false, + oidx: Some(2), + }, + ], + no_main_func: None + } + ); + + Ok(()) + } + + #[test] + fn test_parse_pgsql_mutli_sig() -> anyhow::Result<()> { + let code = r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT $3::TEXT, $1::BIGINT; +SELECT $2::TEXT; +"#; + //println!("{}", serde_json::to_string()?); + assert_eq!( + parse_pgsql_sig(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + otyp: Some("bigint".to_string()), + name: "param1".to_string(), + typ: Typ::Int, + default: None, + has_default: false, + oidx: Some(1), + }, + Arg { + otyp: Some("text".to_string()), + name: "param2".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: Some(2), + }, + Arg { + otyp: Some("text".to_string()), + name: "param3".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: Some(3), + }, + ], + no_main_func: None + } + ); + + Ok(()) + } + + #[test] + fn test_parse_sql_blocks_multi_2semi() -> anyhow::Result<()> { + let code = r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT; +-- ; +SELECT $2::TEXT; +"#; + assert_eq!(parse_sql_blocks(code).len(), 2); + + Ok(()) + } + + #[test] + fn test_parse_sql_blocks_multi_1semi() -> anyhow::Result<()> { + let code = r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT; +-- ; +SELECT $2::TEXT +"#; + assert_eq!( + parse_sql_blocks(code), + vec![ + r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT;"#, + r#" +-- ; +SELECT $2::TEXT +"# + ] + ); + + Ok(()) + } + + #[test] + fn test_parse_sql_blocks_single_1semi() -> anyhow::Result<()> { + let code = r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT; +-- hey +"#; + assert_eq!( + parse_sql_blocks(code), + vec![ + r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT;"#, + ] + ); + + Ok(()) + } + + #[test] + fn test_parse_sql_blocks_single_nosemi() -> anyhow::Result<()> { + let code = r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT +"#; + assert_eq!( + parse_sql_blocks(code), + vec![ + r#" +-- $1 param1 +-- $2 param2 +-- $3 param3 +SELECT '--', ';' $3::TEXT, $1::BIGINT +"# + ] + ); + + Ok(()) + } + + #[test] + fn test_parse_mysql_positional_sig() -> anyhow::Result<()> { + let code = r#" +-- ? param1 (int) = 3 +-- ? param2 (text) +SELECT ?, ?; +"#; + assert_eq!( + parse_mysql_sig(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + otyp: Some("int".to_string()), + name: "param1".to_string(), + typ: Typ::Int, + default: Some(json!(3)), + has_default: true, + oidx: None, + }, + Arg { + otyp: Some("text".to_string()), + name: "param2".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, + }, + ], + no_main_func: None + } + ); + + Ok(()) + } + + #[test] + fn test_parse_mysql_sig() -> anyhow::Result<()> { + let code = r#" +-- :param1 (int) = 3 +-- :param2 (text) +-- :param3 (text) +SELECT :param3, :param1; +SELECT :param2; +"#; + assert_eq!( + parse_mysql_sig(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + otyp: Some("int".to_string()), + name: "param1".to_string(), + typ: Typ::Int, + default: Some(json!(3)), + has_default: true, + oidx: None, + }, + Arg { + otyp: Some("text".to_string()), + name: "param2".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, + }, + Arg { + otyp: Some("text".to_string()), + name: "param3".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, }, ], no_main_func: None @@ -408,9 +796,10 @@ SELECT * FROM table WHERE token=$1::TEXT AND image=$2::BIGINT #[test] fn test_parse_bigquery_sig() -> anyhow::Result<()> { let code = r#" --- @token (string) +-- @token (string) = abc -- @image (int64) -SELECT * FROM table WHERE token=@token AND image=@image +SELECT * FROM table WHERE token=@token AND image=@image; +SELECT @token; "#; //println!("{}", serde_json::to_string()?); assert_eq!( @@ -423,15 +812,111 @@ SELECT * FROM table WHERE token=@token AND image=@image otyp: Some("string".to_string()), name: "token".to_string(), typ: Typ::Str(None), - default: None, - has_default: false + default: Some(json!("abc")), + has_default: true, + oidx: None, }, Arg { otyp: Some("int64".to_string()), name: "image".to_string(), typ: Typ::Int, default: None, - has_default: false + has_default: false, + oidx: None, + }, + ], + no_main_func: None + } + ); + + Ok(()) + } + + #[test] + fn test_parse_snowflake_sig() -> anyhow::Result<()> { + let code = r#" +-- ? param1 (int) = 3 +-- ? param2 (varchar) +SELECT ?, ?; +-- ? param3 (varchar) +SELECT ?; +"#; + assert_eq!( + parse_snowflake_sig(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + otyp: Some("int".to_string()), + name: "param1".to_string(), + typ: Typ::Int, + default: Some(json!(3)), + has_default: true, + oidx: None, + }, + Arg { + otyp: Some("varchar".to_string()), + name: "param2".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, + }, + Arg { + otyp: Some("varchar".to_string()), + name: "param3".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, + } + ], + no_main_func: None + } + ); + + Ok(()) + } + + #[test] + fn test_parse_mssql_sig() -> anyhow::Result<()> { + let code = r#" +-- @p1 param1 (int) = 3 +-- @p2 param2 (varchar) +-- @p3 param3 (varchar) +SELECT @p3, @p1; +SELECT @p2; +"#; + assert_eq!( + parse_mssql_sig(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + otyp: Some("int".to_string()), + name: "param1".to_string(), + typ: Typ::Int, + default: Some(json!(3)), + has_default: true, + oidx: None, + }, + Arg { + otyp: Some("varchar".to_string()), + name: "param2".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, + }, + Arg { + otyp: Some("varchar".to_string()), + name: "param3".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false, + oidx: None, }, ], no_main_func: None diff --git a/backend/parsers/windmill-parser-ts/src/lib.rs b/backend/parsers/windmill-parser-ts/src/lib.rs index 44521b9598..56d645edf4 100644 --- a/backend/parsers/windmill-parser-ts/src/lib.rs +++ b/backend/parsers/windmill-parser-ts/src/lib.rs @@ -204,6 +204,7 @@ fn parse_param( typ, default: None, has_default: ident.id.optional || nullable, + oidx: None, }) } // Pat::Object(ObjectPat { ... }) = todo!() @@ -250,13 +251,13 @@ fn parse_param( if typ == Typ::Unknown && dflt.is_some() { typ = json_to_typ(dflt.as_ref().unwrap()); } - Ok(Arg { otyp: None, name, typ, default: dflt, has_default: true }) + Ok(Arg { otyp: None, name, typ, default: dflt, has_default: true, oidx: None }) } Pat::Object(ObjectPat { type_ann, .. }) => { let (typ, nullable) = eval_type_ann(&type_ann); *counter += 1; let name = format!("anon{}", counter); - Ok(Arg { otyp: None, name, typ, default: None, has_default: nullable }) + Ok(Arg { otyp: None, name, typ, default: None, has_default: nullable, oidx: None }) } _ => Err(anyhow::anyhow!( "parameter syntax unsupported: `{}`: {:#?}", diff --git a/backend/parsers/windmill-parser-wasm/pkg/package.json b/backend/parsers/windmill-parser-wasm/pkg/package.json index 0bd3c3dfdd..0c70400479 100644 --- a/backend/parsers/windmill-parser-wasm/pkg/package.json +++ b/backend/parsers/windmill-parser-wasm/pkg/package.json @@ -3,7 +3,7 @@ "collaborators": [ "Ruben Fiszel " ], - "version": "1.364.4", + "version": "1.367.2", "files": [ "windmill_parser_wasm_bg.wasm", "windmill_parser_wasm.js", @@ -14,4 +14,4 @@ "sideEffects": [ "./snippets/*" ] -} +} \ No newline at end of file diff --git a/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm.js b/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm.js index 7198bec9ec..a945347d31 100644 --- a/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm.js +++ b/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm.js @@ -6,6 +6,20 @@ heap.push(undefined, null, true, false); function getObject(idx) { return heap[idx]; } +let heap_next = heap.length; + +function dropObject(idx) { + if (idx < 132) return; + heap[idx] = heap_next; + heap_next = idx; +} + +function takeObject(idx) { + const ret = getObject(idx); + dropObject(idx); + return ret; +} + let WASM_VECTOR_LEN = 0; let cachedUint8Memory0 = null; @@ -84,20 +98,6 @@ function getInt32Memory0() { return cachedInt32Memory0; } -let heap_next = heap.length; - -function dropObject(idx) { - if (idx < 132) return; - heap[idx] = heap_next; - heap_next = idx; -} - -function takeObject(idx) { - const ret = getObject(idx); - dropObject(idx); - return ret; -} - let cachedFloat64Memory0 = null; function getFloat64Memory0() { @@ -585,6 +585,13 @@ async function __wbg_load(module, imports) { function __wbg_get_imports() { const imports = {}; imports.wbg = {}; + imports.wbg.__wbg_eval_aa725d466edcea2c = function(arg0, arg1) { + const ret = eval(getStringFromWasm0(arg0, arg1)); + return addHeapObject(ret); + }; + imports.wbg.__wbindgen_object_drop_ref = function(arg0) { + takeObject(arg0); + }; imports.wbg.__wbindgen_string_get = function(arg0, arg1) { const obj = getObject(arg1); const ret = typeof(obj) === 'string' ? obj : undefined; @@ -593,13 +600,6 @@ function __wbg_get_imports() { getInt32Memory0()[arg0 / 4 + 1] = len1; getInt32Memory0()[arg0 / 4 + 0] = ptr1; }; - imports.wbg.__wbindgen_object_drop_ref = function(arg0) { - takeObject(arg0); - }; - imports.wbg.__wbg_eval_2ea6d5f9a10f336a = function(arg0, arg1) { - const ret = eval(getStringFromWasm0(arg0, arg1)); - return addHeapObject(ret); - }; imports.wbg.__wbindgen_boolean_get = function(arg0) { const v = getObject(arg0); const ret = typeof(v) === 'boolean' ? (v ? 1 : 0) : 2; diff --git a/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm_bg.wasm b/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm_bg.wasm index 26776f17c1..eb64263b45 100644 Binary files a/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm_bg.wasm and b/backend/parsers/windmill-parser-wasm/pkg/windmill_parser_wasm_bg.wasm differ diff --git a/backend/parsers/windmill-parser-wasm/tests/wasm.rs b/backend/parsers/windmill-parser-wasm/tests/wasm.rs index ddefcae149..9f0feacbed 100644 --- a/backend/parsers/windmill-parser-wasm/tests/wasm.rs +++ b/backend/parsers/windmill-parser-wasm/tests/wasm.rs @@ -28,77 +28,88 @@ export function main(test1?: string, test2: string = \"burkina\", name: "test1".to_string(), typ: Typ::Str(None), default: None, - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "test2".to_string(), typ: Typ::Str(None), default: Some(json!("burkina")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "test3".to_string(), typ: Typ::Resource("postgres".to_string()), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "b64".to_string(), typ: Typ::Bytes, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "ls".to_string(), typ: Typ::List(Box::new(Typ::Bytes)), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "email".to_string(), typ: Typ::Email, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "literal".to_string(), typ: Typ::Str(Some(vec!["test".to_string()])), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "literal_union".to_string(), typ: Typ::Str(Some(vec!["test".to_string(), "test2".to_string()])), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "opt_type".to_string(), typ: Typ::Str(None), default: None, - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "opt_type_union".to_string(), typ: Typ::Str(None), default: None, - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "opt_type_union_union2".to_string(), typ: Typ::Str(None), default: None, - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, @@ -108,21 +119,24 @@ export function main(test1?: string, test2: string = \"burkina\", ObjectProperty { key: "b".to_string(), typ: Box::new(Typ::Float) } ]), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "literals_with_undefined".to_string(), typ: Typ::Str(Some(vec!["foo".to_string(), "bar".to_string()])), default: None, - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "dyn_select".to_string(), typ: Typ::DynSelect("foo".to_string()), default: None, - has_default: false + has_default: false, + oidx: None } ], no_main_func: Some(false) @@ -154,35 +168,40 @@ export function main(test2 = \"burkina\", name: "test2".to_string(), typ: Typ::Str(None), default: Some(json!("burkina")), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "bool".to_string(), typ: Typ::Bool, default: Some(json!(true)), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "float".to_string(), typ: Typ::Float, default: Some(json!(4.2)), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "int".to_string(), typ: Typ::Int, default: Some(json!(42)), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, name: "ls".to_string(), typ: Typ::List(Box::new(Typ::Str(None))), default: Some(json!(["test"])), - has_default: true + has_default: true, + oidx: None }, Arg { otyp: None, @@ -192,7 +211,8 @@ export function main(test2 = \"burkina\", ObjectProperty { key: "b".to_string(), typ: Box::new(Typ::Int) } ]), default: Some(json!({"a": "test", "b": 42})), - has_default: true + has_default: true, + oidx: None } ], no_main_func: Some(false) @@ -224,21 +244,24 @@ export function main(foo: FooBar, {a, b}: FooBar, {c, d}: FooBar = {a: \"foo\", otyp: None, typ: Typ::Resource("foo_bar".to_string()), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { name: "anon1".to_string(), otyp: None, typ: Typ::Resource("foo_bar".to_string()), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { name: "anon2".to_string(), otyp: None, typ: Typ::Resource("foo_bar".to_string()), default: Some(json!({"a": "foo", "b": 42})), - has_default: true + has_default: true, + oidx: None } ], no_main_func: Some(false) @@ -268,7 +291,8 @@ export function main(foo: (\"foo\" | \"bar\")[]) { "bar".to_string() ])))), default: None, - has_default: false + has_default: false, + oidx: None }], no_main_func: Some(false) } @@ -349,56 +373,64 @@ Write-Output 'Testing...' name: "test_none".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_string".to_string(), typ: Typ::Str(None), default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_int".to_string(), typ: Typ::Int, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_decimal".to_string(), typ: Typ::Float, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_double".to_string(), typ: Typ::Float, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_single".to_string(), typ: Typ::Float, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_datetime_lower".to_string(), typ: Typ::Datetime, default: None, - has_default: false + has_default: false, + oidx: None }, Arg { otyp: None, name: "test_datetime_upper".to_string(), typ: Typ::Datetime, default: None, - has_default: false + has_default: false, + oidx: None } ], no_main_func: None, diff --git a/backend/parsers/windmill-parser/Cargo.toml b/backend/parsers/windmill-parser/Cargo.toml index a56721468a..9bf4ad5eba 100644 --- a/backend/parsers/windmill-parser/Cargo.toml +++ b/backend/parsers/windmill-parser/Cargo.toml @@ -9,5 +9,5 @@ name = "windmill_parser" path = "./src/lib.rs" [dependencies] -serde.workspace = true +serde = { version = "1.0", features = ["derive"] } serde_json.workspace = true diff --git a/backend/parsers/windmill-parser/src/lib.rs b/backend/parsers/windmill-parser/src/lib.rs index f6d34cade8..a739dc53d2 100644 --- a/backend/parsers/windmill-parser/src/lib.rs +++ b/backend/parsers/windmill-parser/src/lib.rs @@ -57,6 +57,7 @@ pub struct Arg { pub typ: Typ, pub default: Option, pub has_default: bool, + pub oidx: Option, } pub fn json_to_typ(js: &Value) -> Typ { diff --git a/backend/parsers/windmill-sql-datatype-parser-wasm/Cargo.toml b/backend/parsers/windmill-sql-datatype-parser-wasm/Cargo.toml index 3176ec1ce6..2d0ee556ae 100644 --- a/backend/parsers/windmill-sql-datatype-parser-wasm/Cargo.toml +++ b/backend/parsers/windmill-sql-datatype-parser-wasm/Cargo.toml @@ -16,5 +16,4 @@ wasm-bindgen-test.workspace = true [dependencies] windmill-parser.workspace = true windmill-parser-sql.workspace = true -wasm-bindgen.workspace = true -serde = { version = "1.0", features = ["derive"] } \ No newline at end of file +wasm-bindgen.workspace = true \ No newline at end of file diff --git a/backend/windmill-worker/src/bigquery_executor.rs b/backend/windmill-worker/src/bigquery_executor.rs index 15523211a2..d9194c7475 100644 --- a/backend/windmill-worker/src/bigquery_executor.rs +++ b/backend/windmill-worker/src/bigquery_executor.rs @@ -1,9 +1,14 @@ -use futures::TryFutureExt; +use std::collections::HashMap; + +use futures::future::BoxFuture; +use futures::{FutureExt, 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, parse_db_resource}; +use windmill_parser_sql::{ + parse_bigquery_sig, parse_db_resource, parse_sql_blocks, parse_sql_statement_named_params, +}; use windmill_queue::{CanceledBy, HTTP_CLIENT}; use serde::Deserialize; @@ -57,119 +62,35 @@ struct BigqueryError { message: String, } -pub async fn do_bigquery( - job: &QueuedJob, - client: &AuthedClientBackgroundTask, - query: &str, - db: &sqlx::Pool, - mem_peak: &mut i32, - canceled_by: &mut Option, - worker_name: &str, - column_order: &mut Option>, -) -> windmill_common::error::Result> { - let bigquery_args = build_args_values(job, client, db).await?; +fn do_bigquery_inner<'a>( + query: &'a str, + all_statement_values: &'a HashMap, + project_id: &'a str, + token: &'a str, + timeout_ms: i32, + column_order: Option<&'a mut Option>>, +) -> windmill_common::error::Result>>> { + let param_names = parse_sql_statement_named_params(query, '@'); - let inline_db_res_path = parse_db_resource(&query); + let statement_values = all_statement_values + .iter() + .filter_map(|(name, val)| { + if param_names.contains(name) { + Some(val) + } else { + None + } + }) + .collect::>(); - let db_arg = if let Some(inline_db_res_path) = inline_db_res_path { - Some( - client - .get_authed() - .await - .get_resource_value_interpolated::( - &inline_db_res_path, - Some(job.id.to_string()), - ) - .await?, - ) - } else { - bigquery_args.get("database").cloned() - }; - - let database = if let Some(db) = db_arg { - db.to_string() - } else { - return Err(Error::BadRequest("Missing database argument".to_string())); - }; - - let service_account = CustomServiceAccount::from_json(&database) - .map_err(|e| Error::ExecutionErr(e.to_string()))?; - - let authentication_manager = AuthenticationManager::from(service_account); - let scopes = &["https://www.googleapis.com/auth/bigquery"]; - let token = authentication_manager - .get_token(scopes) - .await - .map_err(|e| Error::ExecutionErr(e.to_string()))?; - - let mut statement_values: Vec = vec![]; - - let sig = parse_bigquery_sig(&query) - .map_err(|x| Error::ExecutionErr(x.to_string()))? - .args; - - for arg in &sig { - let arg_t = arg.otyp.clone().unwrap_or_else(|| "string".to_string()); - let arg_n = arg.clone().name; - let arg_v = bigquery_args.get(&arg.name).cloned().unwrap_or(json!("")); - let bigquery_v = if arg_t.ends_with("[]") { - let base_type = arg_t.strip_suffix("[]").unwrap_or(&arg_t); - json!({ - "name": arg.name, - "parameterType": { - "type": "ARRAY", - "arrayType": { - "type": base_type.to_uppercase() - } - }, - "parameterValue": { - "arrayValues": bigquery_args - .get(&arg.name) - .unwrap_or(&json!([])) - .as_array() - .unwrap_or(&vec![]) - .iter() - .map(|x| { - convert_val(base_type.to_string(), x.clone()) - }) - .collect::>() - } - }) - } else { - json!({ - "name": arg_n, - "parameterType": { - "type": arg_t.to_uppercase() - }, - "parameterValue": { - "value": convert_val(arg_t, arg_v), - } - }) - }; - - statement_values.push(bigquery_v); - } - - let timeout_ms = i32::try_from( - resolve_job_timeout(&db, &job.workspace_id, job.id, job.timeout) - .await - .0 - .as_millis(), - ) - .unwrap_or(200000); - - let result_f = async { + let result_f = async move { 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() + + project_id + "/queries", ) - .bearer_auth(token.as_str()) + .bearer_auth(token) .json(&json!({ "query": query, "useLegacySql": false, @@ -222,16 +143,18 @@ pub async fn do_bigquery( )); } - *column_order = Some( - result - .schema - .as_ref() - .unwrap() - .fields - .iter() - .map(|x| x.name.clone()) - .collect::>(), - ); + if let Some(column_order) = column_order { + *column_order = Some( + result + .schema + .as_ref() + .unwrap() + .fields + .iter() + .map(|x| x.name.clone()) + .collect::>(), + ); + } let rows = result .rows @@ -268,6 +191,154 @@ pub async fn do_bigquery( }, } }; + + Ok(result_f.boxed()) +} + +pub async fn do_bigquery( + job: &QueuedJob, + client: &AuthedClientBackgroundTask, + query: &str, + db: &sqlx::Pool, + mem_peak: &mut i32, + canceled_by: &mut Option, + worker_name: &str, + column_order: &mut Option>, +) -> windmill_common::error::Result> { + let bigquery_args = build_args_values(job, client, db).await?; + + 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::( + &inline_db_res_path, + Some(job.id.to_string()), + ) + .await?, + ) + } else { + bigquery_args.get("database").cloned() + }; + + let database = if let Some(db) = db_arg { + db.to_string() + } else { + return Err(Error::BadRequest("Missing database argument".to_string())); + }; + + let service_account = CustomServiceAccount::from_json(&database) + .map_err(|e| Error::ExecutionErr(e.to_string()))?; + + let authentication_manager = AuthenticationManager::from(service_account); + let scopes = &["https://www.googleapis.com/auth/bigquery"]; + let token = authentication_manager + .get_token(scopes) + .await + .map_err(|e| Error::ExecutionErr(e.to_string()))?; + + let timeout_ms = i32::try_from( + resolve_job_timeout(&db, &job.workspace_id, job.id, job.timeout) + .await + .0 + .as_millis(), + ) + .unwrap_or(200000); + + let project_id = authentication_manager + .project_id() + .await + .map_err(|e| Error::ExecutionErr(e.to_string()))?; + + let queries = parse_sql_blocks(query); + + let mut statement_values: HashMap = HashMap::new(); + + let sig = parse_bigquery_sig(&query) + .map_err(|x| Error::ExecutionErr(x.to_string()))? + .args; + + for arg in &sig { + let arg_t = arg.otyp.clone().unwrap_or_else(|| "string".to_string()); + let arg_n = arg.clone().name; + let arg_v = bigquery_args.get(&arg.name).cloned().unwrap_or(json!("")); + let bigquery_v = if arg_t.ends_with("[]") { + let base_type = arg_t.strip_suffix("[]").unwrap_or(&arg_t); + json!({ + "name": arg.name, + "parameterType": { + "type": "ARRAY", + "arrayType": { + "type": base_type.to_uppercase() + } + }, + "parameterValue": { + "arrayValues": bigquery_args + .get(&arg.name) + .unwrap_or(&json!([])) + .as_array() + .unwrap_or(&vec![]) + .iter() + .map(|x| { + convert_val(base_type.to_string(), x.clone()) + }) + .collect::>() + } + }) + } else { + json!({ + "name": arg_n, + "parameterType": { + "type": arg_t.to_uppercase() + }, + "parameterValue": { + "value": convert_val(arg_t, arg_v), + } + }) + }; + + statement_values.insert(arg_n, bigquery_v); + } + + let result_f = if queries.len() > 1 { + let futures = queries + .iter() + .map(|x| { + do_bigquery_inner( + x, + &statement_values, + &project_id, + token.as_str(), + timeout_ms, + None, + ) + }) + .collect::>>()?; + + let f = async { + let mut res: Vec> = vec![]; + for fut in futures { + let r = fut.await?; + res.push(r); + } + Ok(to_raw_value(&res)) + }; + + f.boxed() + } else { + do_bigquery_inner( + query, + &statement_values, + &project_id, + token.as_str(), + timeout_ms, + Some(column_order), + )? + }; + let r = run_future_with_polling_update_job_poller( job.id, job.timeout, diff --git a/backend/windmill-worker/src/mysql_executor.rs b/backend/windmill-worker/src/mysql_executor.rs index d6d56968de..8c138351e4 100644 --- a/backend/windmill-worker/src/mysql_executor.rs +++ b/backend/windmill-worker/src/mysql_executor.rs @@ -1,17 +1,23 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; use base64::Engine; +use futures::{future::BoxFuture, FutureExt}; +use itertools::Itertools; use mysql_async::{ consts::ColumnType, prelude::*, FromValueError, OptsBuilder, Params, Row, SslOpts, }; use serde::{Deserialize, Serialize}; use serde_json::{json, value::RawValue, Value}; use sqlx::types::Json; +use tokio::sync::Mutex; use windmill_common::{ error::{to_anyhow, Error}, jobs::QueuedJob, }; -use windmill_parser_sql::{parse_db_resource, parse_mysql_sig, RE_ARG_MYSQL_NAMED}; +use windmill_parser_sql::{ + parse_db_resource, parse_mysql_sig, parse_sql_blocks, parse_sql_statement_named_params, + RE_ARG_MYSQL_NAMED, +}; use windmill_queue::CanceledBy; use crate::{ @@ -29,6 +35,59 @@ struct MysqlDatabase { ssl: Option, } +pub fn do_mysql_inner<'a>( + query: &'a str, + all_statement_values: &Params, + conn: Arc>, + column_order: Option<&'a mut Option>>, +) -> windmill_common::error::Result>>> { + let param_names = parse_sql_statement_named_params(query, ':') + .into_iter() + .map(|x| x.into_bytes()) + .collect_vec(); + + let statement_values = if let Params::Named(m) = all_statement_values { + Params::Named( + m.into_iter() + .filter(|(k, _)| param_names.contains(&k)) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + ) + } else { + all_statement_values.clone() + }; + + let result_f = async move { + let rows: Vec = conn + .lock() + .await + .exec(query, statement_values) + .await + .map_err(to_anyhow)?; + + if let Some(column_order) = column_order { + *column_order = Some( + rows.first() + .map(|x| { + x.columns() + .iter() + .map(|x| x.name_str().to_string()) + .collect::>() + }) + .unwrap_or_default(), + ); + } + + Ok(rows + .into_iter() + .map(|x| convert_row_to_value(x)) + .collect::>()) + as Result, anyhow::Error> + }; + + Ok(result_f.boxed()) +} + pub async fn do_mysql( job: &QueuedJob, client: &AuthedClientBackgroundTask, @@ -91,14 +150,11 @@ pub async fn do_mysql( opts }; - let pool = mysql_async::Pool::new(opts); - let mut conn = pool.get_conn().await.map_err(to_anyhow)?; - - let sig = parse_mysql_sig(&query) + let sig = parse_mysql_sig(query) .map_err(|x| Error::ExecutionErr(x.to_string()))? .args; - let using_named_params = RE_ARG_MYSQL_NAMED.captures_iter(&query).count() > 0; + let using_named_params = RE_ARG_MYSQL_NAMED.captures_iter(query).count() > 0; let mut statement_values: Params = match using_named_params { true => Params::Named(HashMap::new()), @@ -106,10 +162,8 @@ pub async fn do_mysql( }; for arg in &sig { let arg_t = arg.otyp.clone().unwrap_or_else(|| "text".to_string()); - let arg_n = arg.clone().name; - let mysql_v = match job - .args - .as_ref() + let arg_n = arg.name.clone(); + let mysql_v = match job_args .and_then(|x| { x.get(arg.name.as_str()) .map(|x| serde_json::from_str::(x.get()).ok()) @@ -172,35 +226,30 @@ pub async fn do_mysql( } } - let result_f = async { - let rows: Vec = conn - .exec( - query, - match statement_values { - Params::Positional(v) => Params::Positional(v), - Params::Named(m) => Params::Named(m), - _ => Params::Empty, - }, - ) - .await - .map_err(to_anyhow)?; + let pool = mysql_async::Pool::new(opts); + let conn = pool.get_conn().await.map_err(to_anyhow)?; + let conn_a = Arc::new(Mutex::new(conn)); - *column_order = Some( - rows.first() - .map(|x| { - x.columns() - .iter() - .map(|x| x.name_str().to_string()) - .collect::>() - }) - .unwrap_or_default(), - ); + let queries = parse_sql_blocks(query); - Ok(rows - .into_iter() - .map(|x| convert_row_to_value(x)) - .collect::>()) - as Result, anyhow::Error> + let result_f = if queries.len() > 1 { + let futures = queries + .iter() + .map(|x| do_mysql_inner(x, &statement_values, conn_a.clone(), None)) + .collect::>>()?; + + let f = async { + let mut res: Vec = vec![]; + for fut in futures { + let r = fut.await?; + res.push(serde_json::to_value(r).map_err(to_anyhow)?); + } + Ok(res) + }; + + f.boxed() + } else { + do_mysql_inner(query, &statement_values, conn_a.clone(), Some(column_order))? }; let result = run_future_with_polling_update_job_poller( @@ -215,7 +264,7 @@ pub async fn do_mysql( ) .await?; - drop(conn); + drop(conn_a); pool.disconnect().await.map_err(to_anyhow)?; diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 981585c7c5..3fb96677fd 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -1,12 +1,15 @@ +use std::collections::HashMap; use std::net::IpAddr; -use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::Duration; use anyhow::Context; use base64::{engine, Engine as _}; use chrono::Utc; -use futures::TryStreamExt; +use futures::future::BoxFuture; +use futures::{FutureExt, TryStreamExt}; +use itertools::Itertools; use native_tls::{Certificate, TlsConnector}; use postgres_native_tls::MakeTlsConnector; use rust_decimal::{prelude::FromPrimitive, Decimal}; @@ -16,6 +19,7 @@ use serde_json::Map; use serde_json::Value; use tokio::sync::Mutex; use tokio_postgres::types::IsNull; +use tokio_postgres::Client; use tokio_postgres::{ types::{to_sql_checked, ToSql}, NoTls, Row, @@ -28,8 +32,10 @@ use uuid::Uuid; use windmill_common::error::{self, Error}; use windmill_common::worker::{to_raw_value, CLOUD_HOSTED}; use windmill_common::{error::to_anyhow, jobs::QueuedJob}; -use windmill_parser::Typ; -use windmill_parser_sql::{parse_db_resource, parse_pgsql_sig}; +use windmill_parser::{Arg, Typ}; +use windmill_parser_sql::{ + parse_db_resource, parse_pg_statement_arg_indices, parse_pgsql_sig, parse_sql_blocks, +}; use windmill_queue::CanceledBy; use crate::common::{build_args_values, run_future_with_polling_update_job_poller, sizeof_val}; @@ -56,6 +62,87 @@ lazy_static! { pub static ref RUNNING: AtomicBool = AtomicBool::new(false); } +fn do_postgresql_inner<'a>( + mut query: String, + param_idx_to_arg_and_value: &HashMap)>, + client: &'a Client, + column_order: Option<&'a mut Option>>, + siz: &'a AtomicUsize, +) -> error::Result>>> { + let mut query_params = vec![]; + + let arg_indices = parse_pg_statement_arg_indices(&query); + + let mut i = 1; + for oidx in arg_indices.iter().sorted() { + if let Some((arg, value)) = param_idx_to_arg_and_value.get(&oidx) { + if *oidx as usize != i { + query = query.replace(&format!("${}", oidx), &format!("${}", i)); + } + let value = value.unwrap_or_else(|| &serde_json::Value::Null); + let arg_t = arg + .otyp + .as_ref() + .ok_or_else(|| anyhow::anyhow!("Missing otyp for pg arg"))?; + let typ = &arg.typ; + let param = convert_val(value, arg_t, typ)?; + query_params.push(param); + i += 1; + } + } + + let result_f = async move { + // Now we can execute a simple statement that just returns its parameter. + let rows = client + .query_raw(&query, query_params) + .await + .map_err(to_anyhow)?; + + let rows = rows.try_collect::>().await.map_err(to_anyhow)?; + + if let Some(column_order) = column_order { + *column_order = Some( + rows.first() + .map(|x| { + x.columns() + .iter() + .map(|x| x.name().to_string()) + .collect::>() + }) + .unwrap_or_default(), + ); + } + + let mut res: Vec = vec![]; + for row in rows.into_iter() { + let r = postgres_row_to_json_value(row); + if let Ok(v) = r.as_ref() { + let size = sizeof_val(v); + siz.fetch_add(size, Ordering::Relaxed); + } + if *CLOUD_HOSTED { + let siz = siz.load(Ordering::Relaxed); + if siz > MAX_RESULT_SIZE * 4 { + return Err(anyhow::anyhow!( + "Query result too large for cloud (size = {} > {})", + siz, + MAX_RESULT_SIZE & 4 + )); + } + } + if let Ok(v) = r { + res.push(v); + } else { + return Err(to_anyhow(r.err().unwrap())); + } + } + + Ok(res) + }; + + Ok(result_f.boxed()) +} + pub async fn do_postgresql( job: &QueuedJob, client: &AuthedClientBackgroundTask, @@ -175,34 +262,7 @@ pub async fn do_postgresql( Some((client, handle)) }; - let mut statement_values: Vec = vec![]; - - let sig = parse_pgsql_sig(&query) - .map_err(|x| Error::ExecutionErr(x.to_string()))? - .args; - - for arg in &sig { - statement_values.push( - pg_args - .get(&arg.name) - .map(|x| x.to_owned()) - .unwrap_or_else(|| serde_json::Value::Null), - ); - } - - let query_params = statement_values - .iter() - .enumerate() - .map(|(i, value)| { - let arg_t = &sig[i] - .otyp - .as_ref() - .ok_or_else(|| anyhow::anyhow!("Missing otyp for pg arg"))? - .to_owned(); - let typ = &sig[i].typ; - convert_val(value, arg_t, typ) - }) - .collect::>>()?; + let queries = parse_sql_blocks(query); let (client, handle) = if let Some((client, handle)) = new_client.as_ref() { (client, Some(handle)) @@ -211,52 +271,49 @@ pub async fn do_postgresql( (client, None) }; - let result_f = async { - // Now we can execute a simple statement that just returns its parameter. - let rows = client - .query_raw(query, query_params) - .await - .map_err(to_anyhow)?; + let sig = parse_pgsql_sig(&query).map_err(|x| Error::ExecutionErr(x.to_string()))?; + let param_idx_to_arg_and_value = sig + .args + .iter() + .filter_map(|x| x.oidx.map(|oidx| (oidx, (x, pg_args.get(&x.name))))) + .collect::>(); - let rows = rows.try_collect::>().await.map_err(to_anyhow)?; + let size = AtomicUsize::new(0); + let result_f = if queries.len() > 1 { + let futures = queries + .iter() + .map(|x| { + do_postgresql_inner( + x.to_string(), + ¶m_idx_to_arg_and_value, + client, + None, + &size, + ) + }) + .collect::>>()?; - *column_order = Some( - rows.first() - .map(|x| { - x.columns() - .iter() - .map(|x| x.name().to_string()) - .collect::>() - }) - .unwrap_or_default(), - ); - - let mut siz = 0; - let mut res: Vec = vec![]; - for row in rows.into_iter() { - let r = postgres_row_to_json_value(row); - if let Ok(v) = r.as_ref() { - let size = sizeof_val(v); - siz += size; + let f = async { + let mut res: Vec = vec![]; + for fut in futures { + let r = fut.await?; + res.push(serde_json::to_value(r).map_err(to_anyhow)?); } - if *CLOUD_HOSTED && siz > MAX_RESULT_SIZE * 4 { - return Err(anyhow::anyhow!( - "Query result too large for cloud (size = {} > {})", - siz, - MAX_RESULT_SIZE & 4 - )); - } - if let Ok(v) = r { - res.push(v); - } else { - return Err(to_anyhow(r.err().unwrap())); - } - } + Ok(res) + }; - Ok((res, siz)) + f.boxed() + } else { + do_postgresql_inner( + query.to_string(), + ¶m_idx_to_arg_and_value, + client, + Some(column_order), + &size, + )? }; - let (result, size) = run_future_with_polling_update_job_poller( + let result = run_future_with_polling_update_job_poller( job.id, job.timeout, db, @@ -268,7 +325,7 @@ pub async fn do_postgresql( ) .await?; - *mem_peak = size as i32; + *mem_peak = size.load(Ordering::Relaxed) as i32; RUNNING.store(false, std::sync::atomic::Ordering::Relaxed); diff --git a/backend/windmill-worker/src/snowflake_executor.rs b/backend/windmill-worker/src/snowflake_executor.rs index e18b7106ae..349a2dcd73 100644 --- a/backend/windmill-worker/src/snowflake_executor.rs +++ b/backend/windmill-worker/src/snowflake_executor.rs @@ -1,15 +1,17 @@ use base64::{engine, Engine as _}; use chrono::Datelike; use core::fmt::Write; -use futures::TryFutureExt; +use futures::future::BoxFuture; +use futures::{FutureExt, TryFutureExt}; use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; use serde_json::{json, value::RawValue, Value}; use sha2::{Digest, Sha256}; +use std::collections::HashMap; 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_db_resource, parse_snowflake_sig}; +use windmill_parser_sql::{parse_db_resource, parse_snowflake_sig, parse_sql_blocks}; use windmill_queue::{CanceledBy, HTTP_CLIENT}; use serde::{Deserialize, Serialize}; @@ -63,6 +65,104 @@ struct SnowflakeError { message: String, } +fn do_snowflake_inner<'a>( + query: &'a str, + job_args: &HashMap, + mut body: serde_json::Map, + account_identifier: &'a str, + token: &'a str, + column_order: Option<&'a mut Option>>, +) -> windmill_common::error::Result>>> { + body.insert("statement".to_string(), json!(query)); + + let mut bindings = serde_json::Map::new(); + let sig = parse_snowflake_sig(&query) + .map_err(|x| Error::ExecutionErr(x.to_string()))? + .args; + + let mut i = 1; + for arg in &sig { + let arg_t = arg.otyp.clone().unwrap_or_else(|| "string".to_string()); + let arg_v = job_args.get(&arg.name).cloned().unwrap_or(json!("")); + let snowflake_v = convert_typ_val(arg_t, arg_v); + + bindings.insert(i.to_string(), snowflake_v); + i += 1; + } + + if i > 1 { + body.insert("bindings".to_string(), json!(bindings)); + } + + let result_f = async move { + let response = HTTP_CLIENT + .post(format!( + "https://{}.snowflakecomputing.com/api/v2/statements/", + 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()))?; + + 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(), + )); + } + if let Some(column_order) = column_order { + *column_order = Some( + result + .resultSetMetaData + .rowType + .iter() + .map(|x| x.name.clone()) + .collect::>(), + ); + } + 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), + ); + }); + 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())), + } + } + } + }; + + Ok(result_f.boxed()) +} + pub async fn do_snowflake( job: &QueuedJob, client: &AuthedClientBackgroundTask, @@ -132,21 +232,6 @@ pub async fn do_snowflake( tracing::debug!("Snowflake token: {}", token); - let mut bindings = serde_json::Map::new(); - let sig = parse_snowflake_sig(&query) - .map_err(|x| Error::ExecutionErr(x.to_string()))? - .args; - - let mut i = 1; - for arg in &sig { - let arg_t = arg.otyp.clone().unwrap_or_else(|| "string".to_string()); - let arg_v = snowflake_args.get(&arg.name).cloned().unwrap_or(json!("")); - let snowflake_v = convert_typ_val(arg_t, arg_v); - - bindings.insert(i.to_string(), snowflake_v); - i += 1; - } - let mut body = serde_json::Map::new(); if database.schema.is_some() { body.insert( @@ -172,79 +257,44 @@ pub async fn do_snowflake( json!(database.database.unwrap().to_uppercase()), ); } - body.insert("statement".to_string(), json!(query)); body.insert("timeout".to_string(), json!(10)); // in seconds - if i > 1 { - body.insert("bindings".to_string(), json!(bindings)); - } + let queries = parse_sql_blocks(query); - 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()))?; + let result_f = if queries.len() > 1 { + let futures = queries + .iter() + .map(|x| { + do_snowflake_inner( + x, + &snowflake_args, + body.clone(), + &database.account_identifier, + &token, + None, + ) + }) + .collect::>>()?; - match response.error_for_status_ref() { - Ok(_) => { - let result = response - .json::() - .await - .map_err(|e| Error::ExecutionErr(e.to_string()))?; - - tracing::info!("Snowflake response: {:?}", result); - - 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(), - )); - } - - *column_order = Some( - result - .resultSetMetaData - .rowType - .iter() - .map(|x| x.name.clone()) - .collect::>(), - ); - - 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), - ); - }); - row_map - }) - .collect::>(), - ); - - Ok(rows) + let f = async { + let mut res: Vec> = vec![]; + for fut in futures { + let r = fut.await?; + res.push(r); } - 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())), - } - } - } + Ok(to_raw_value(&res)) + }; + + f.boxed() + } else { + do_snowflake_inner( + query, + &snowflake_args, + body.clone(), + &database.account_identifier, + &token, + Some(column_order), + )? }; let r = run_future_with_polling_update_job_poller( job.id, diff --git a/frontend/package-lock.json b/frontend/package-lock.json index fb93c6edea..792296121b 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -52,7 +52,7 @@ "vscode-languageclient": "~9.0.1", "vscode-uri": "~3.0.8", "vscode-ws-jsonrpc": "~3.1.0", - "windmill-parser-wasm": "^1.364.4", + "windmill-parser-wasm": "^1.367.2", "windmill-sql-datatype-parser-wasm": "^1.318.0", "y-monaco": "^0.1.4", "y-websocket": "^1.5.0", @@ -10261,9 +10261,9 @@ } }, "node_modules/windmill-parser-wasm": { - "version": "1.364.4", - "resolved": "https://registry.npmjs.org/windmill-parser-wasm/-/windmill-parser-wasm-1.364.4.tgz", - "integrity": "sha512-zAW29GvUXdctOJOXSN8Z+yNfG67J5rbmI5FZD8esw8MyIIrha0ajqgtUYIKeuT+Q2kdBLtVc7AVo9/Se2Q4D2g==" + "version": "1.367.2", + "resolved": "https://registry.npmjs.org/windmill-parser-wasm/-/windmill-parser-wasm-1.367.2.tgz", + "integrity": "sha512-If/IXXXADC0jWq4Vj6IN1IRZ1GTWCpE4oMnyZAkVzZUJ2ZbQLw5maLhAmm4iAMpmMHRb4U4kd68D1h30Iy4MNQ==" }, "node_modules/windmill-sql-datatype-parser-wasm": { "version": "1.318.0", diff --git a/frontend/package.json b/frontend/package.json index 6cc1a16f97..aab6fc157b 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -132,7 +132,7 @@ "vscode-languageclient": "~9.0.1", "vscode-uri": "~3.0.8", "vscode-ws-jsonrpc": "~3.1.0", - "windmill-parser-wasm": "^1.364.4", + "windmill-parser-wasm": "^1.367.2", "windmill-sql-datatype-parser-wasm": "^1.318.0", "y-monaco": "^0.1.4", "y-websocket": "^1.5.0", diff --git a/frontend/src/lib/script_helpers.ts b/frontend/src/lib/script_helpers.ts index 23ee85226a..64c1727da6 100644 --- a/frontend/src/lib/script_helpers.ts +++ b/frontend/src/lib/script_helpers.ts @@ -160,32 +160,43 @@ export const POSTGRES_INIT_CODE = `-- to pin the database use '-- database f/you -- $1 name1 = default arg -- $2 name2 -- $3 name3 -INSERT INTO demo VALUES (\$1::TEXT, \$2::INT, \$3::TEXT[]) RETURNING * +-- $4 name4 +INSERT INTO demo VALUES (\$1::TEXT, \$2::INT, \$3::TEXT[]) RETURNING *; +UPDATE demo SET col2 = \$4::INT WHERE col2 = \$2::INT; ` export const MYSQL_INIT_CODE = `-- to pin the database use '-- database f/your/path' -- :name1 (text) = default arg -- :name2 (int) -INSERT INTO demo VALUES (:name1, :name2) +-- :name3 (int) +INSERT INTO demo VALUES (:name1, :name2); +UPDATE demo SET col2 = :name3 WHERE col2 = :name2; ` export const BIGQUERY_INIT_CODE = `-- to pin the database use '-- database f/your/path' -- @name1 (string) = default arg -- @name2 (integer) -- @name3 (string[]) -INSERT INTO \`demodb.demo\` VALUES (@name1, @name2, @name3) +-- @name4 (integer) +INSERT INTO \`demodb.demo\` VALUES (@name1, @name2, @name3); +UPDATE \`demodb.demo\` SET col2 = @name4 WHERE col2 = @name2; ` export const SNOWFLAKE_INIT_CODE = `-- to pin the database use '-- database f/your/path' -- ? name1 (varchar) = default arg -- ? name2 (int) -INSERT INTO demo VALUES (?, ?) +INSERT INTO demo VALUES (?, ?); +-- ? name3 (int) +-- ? name2 (int) +UPDATE demo SET col2 = ? WHERE col2 = ?; ` export const MSSQL_INIT_CODE = `-- to pin the database use '-- database f/your/path' -- @p1 name1 (varchar) = default arg -- @p2 name2 (int) -INSERT INTO demo VALUES (@p1, @p2) +-- @p3 name3 (int) +INSERT INTO demo VALUES (@p1, @p2); +UPDATE demo SET col2 = @p3 WHERE col2 = @p2; ` export const GRAPHQL_INIT_CODE = `query($name4: String, $name2: Int, $name3: [String]) {