diff --git a/backend/src/parser.rs b/backend/src/parser.rs index 1f6dea9e4d..cc47d9cf53 100644 --- a/backend/src/parser.rs +++ b/backend/src/parser.rs @@ -8,21 +8,21 @@ use serde::Serialize; -#[derive(Serialize)] +#[derive(Serialize, Debug, PartialEq)] pub struct MainArgSignature { pub star_args: bool, pub star_kwargs: bool, pub args: Vec, } -#[derive(Serialize, Clone)] +#[derive(Serialize, Clone, Debug, PartialEq)] #[serde(rename_all(serialize = "lowercase"))] pub struct ObjectProperty { pub key: String, pub typ: Box, } -#[derive(Serialize, Clone)] +#[derive(Serialize, Clone, Debug, PartialEq)] #[serde(rename_all(serialize = "lowercase"))] pub enum Typ { Str(Option>), @@ -39,7 +39,7 @@ pub enum Typ { Unknown, } -#[derive(Serialize, Clone)] +#[derive(Serialize, Clone, Debug, PartialEq)] pub struct Arg { pub name: String, pub typ: Typ, diff --git a/backend/src/parser_py.rs b/backend/src/parser_py.rs index a81926437c..4ff8f7beba 100644 --- a/backend/src/parser_py.rs +++ b/backend/src/parser_py.rs @@ -23,8 +23,50 @@ use rustpython_parser::{ parser, }; +fn filter_non_main(code: &str) -> String { + const DEF_MAIN: &str = "def main("; + + let mut filtered_code = String::new(); + let mut code_iter = code.split("\n"); + let mut remaining: String = String::new(); + while let Some(line) = code_iter.next() { + if line.starts_with(DEF_MAIN) { + filtered_code += DEF_MAIN; + remaining += line.strip_prefix(DEF_MAIN).unwrap(); + remaining += &code_iter.join("\n"); + break; + } + } + if filtered_code.is_empty() { + return String::new(); + } + let mut chars = remaining.chars(); + let mut open_parens = 1; + + while let Some(c) = chars.next() { + if c == '(' { + open_parens += 1; + } else if c == ')' { + open_parens -= 1; + } + filtered_code.push(c); + if open_parens == 0 { + break; + } + } + + filtered_code.push_str(": return"); + return filtered_code; +} + pub fn parse_python_signature(code: &str) -> error::Result { - let ast = parser::parse_program(code) + let filtered_code = filter_non_main(code); + if filtered_code.is_empty() { + return Err(error::Error::BadRequest( + "No main function found".to_string(), + )); + } + let ast = parser::parse_program(&filtered_code) .map_err(|e| error::Error::ExecutionErr(format!("Error parsing code: {}", e.to_string())))? .statements; let param = ast.into_iter().find_map(|x| match x { @@ -163,11 +205,16 @@ pub fn parse_python_imports(code: &str) -> error::Result> { .collect(); Ok(lines) } else { + let code = &code + .split("\n") + .filter(|x| x.starts_with("import ") || x.starts_with("from ")) + .join("\n"); let ast = parser::parse_program(code) .map_err(|e| { error::Error::ExecutionErr(format!("Error parsing code: {}", e.to_string())) })? .statements; + let imports = ast .into_iter() .filter_map(|x| match x { @@ -191,6 +238,7 @@ pub fn parse_python_imports(code: &str) -> error::Result> { .filter(|x| !STDIMPORTS.contains(&x.as_str())) .unique() .collect(); + Ok(imports) } } @@ -198,12 +246,10 @@ pub fn parse_python_imports(code: &str) -> error::Result> { #[cfg(test)] mod tests { - // Note this useful idiom: importing names from outer (for mod tests) scope. use super::*; #[test] fn test_parse_python_sig() -> anyhow::Result<()> { - //let code = "print(2 + 3, fd=sys.stderr)"; let code = " import os @@ -215,7 +261,124 @@ def main(test1: str, name: datetime.datetime = datetime.now(), byte: bytes = byt return {\"len\": len(name), \"splitted\": name.split() } "; - println!("{}", serde_json::to_string(&parse_python_signature(code)?)?); + //println!("{}", serde_json::to_string()?); + assert_eq!( + parse_python_signature(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + name: "test1".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false + }, + Arg { + name: "name".to_string(), + typ: Typ::Unknown, + default: Some(json!("")), + has_default: true + }, + Arg { + name: "byte".to_string(), + typ: Typ::Bytes, + default: Some(json!("")), + has_default: true + } + ] + } + ); + + Ok(()) + } + + #[test] + fn test_parse_python_sig_2() -> anyhow::Result<()> { + let code = " + +import os + +def main(test1: str, + name: datetime.datetime = datetime.now(), + byte: bytes = bytes(1)): + + print(f\"Hello World and a warm welcome especially to {name}\") + print(\"The env variable at `all/pretty_secret`: \", os.environ.get(\"ALL_PRETTY_SECRET\")) + return {\"len\": len(name), \"splitted\": name.split() } + +"; + //println!("{}", serde_json::to_string()?); + assert_eq!( + parse_python_signature(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + name: "test1".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false + }, + Arg { + name: "name".to_string(), + typ: Typ::Unknown, + default: Some(json!("")), + has_default: true + }, + Arg { + name: "byte".to_string(), + typ: Typ::Bytes, + default: Some(json!("")), + has_default: true + } + ] + } + ); + + Ok(()) + } + + #[test] + fn test_parse_python_sig_3() -> anyhow::Result<()> { + let code = " + +import os + +def main(test1: str, + name: datetime.datetime = datetime.now(), + byte: bytes = bytes(1)): return + +"; + //println!("{}", serde_json::to_string()?); + assert_eq!( + parse_python_signature(code)?, + MainArgSignature { + star_args: false, + star_kwargs: false, + args: vec![ + Arg { + name: "test1".to_string(), + typ: Typ::Str(None), + default: None, + has_default: false + }, + Arg { + name: "name".to_string(), + typ: Typ::Unknown, + default: Some(json!("")), + has_default: true + }, + Arg { + name: "byte".to_string(), + typ: Typ::Bytes, + default: Some(json!("")), + has_default: true + } + ] + } + ); Ok(()) } diff --git a/backend/src/worker.rs b/backend/src/worker.rs index f80aa7e32c..6f6c38a279 100644 --- a/backend/src/worker.rs +++ b/backend/src/worker.rs @@ -576,6 +576,7 @@ async fn handle_nondep_job( workspace_id = %job.workspace_id, "started setup python dependencies" ); + let child = if !disable_nsjail { Command::new(nsjail_path) .current_dir(job_dir) @@ -745,6 +746,7 @@ print(res_json) .stderr(Stdio::piped()) .spawn()? }; + *status = handle_child(job, db, logs, last_line, timeout, child).await; tracing::info!( worker_name = %worker_name, @@ -1379,6 +1381,69 @@ mod tests { } } + #[sqlx::test(fixtures("base"))] + async fn test_python_flow_2(db: DB) { + initialize_tracing().await; + + let flow: FlowValue = serde_json::from_value(serde_json::json!({ + "modules": [ + { + "value": { + "type": "rawscript", + "content": "import os\nimport wmill\nfrom datetime import datetime\n\nimport re\nimport requests\nfrom datetime import datetime\nfrom html import unescape\n\n\ndef extract(html):\n r = r\"]*alt=\\\"(?P[^\\\"]+)\\\")|([^>]*src=\\\"(?P[^\\\"]+)\\\"))+[^>]*>\"\n match = re.search(r, html)\n if match:\n return {'url': match.group(\"url\"), 'desc': unescape(match.group(\"alt\"))}\n return None\n\n\ndef fetch(date=None):\n if date is None:\n date = datetime.now()\n url = f\"https://dilbert.com/strip/{date:%Y-%m-%d}\"\n return requests.get(url, allow_redirects=False).text\n\n\ndef get_today_comic(date=None):\n return extract(fetch(date))\n\n\n# Our webeditor includes a syntax, type checker through a language server running pyright\n# and the autoformatter Black in our servers. Use Cmd/Ctrl + S to autoformat the code.\n# Beware that the code is only saved when you click Save and not across reload.\n# You can however navigate to any steps safely.\n\"\"\"\nThe client is used to interact with windmill itself through its standard API.\nOne can explore the methods available through autocompletion of `client.XXX`.\nOnly the most common methods are included for ease of use. Request more as\nfeedback if you feel you are missing important ones.\n\"\"\"\n\n\ndef main(\n date: str = None\n):\n dateToFetch = datetime.strptime(date, '%Y-%m-%d') if (date is not None and len(date)>0) else datetime.now()\n print(f\"Fetchind Dilber from {dateToFetch}\")\n # retrieve variables, including secrets by querying the windmill platform.\n # secret fetching is audited by windmill.\n # secret = wmill.get_variable(\"g/all/pretty_secret\")\n return get_today_comic(dateToFetch)\n", + "language": "python3" + }, + "input_transform": { + "date": { + "expr": "undefined", + "type": "javascript" + } + } + }, + { + "value": { + "type": "rawscript", + "content": "\nimport requests\n\n\ndef main(bot_token: str, chat_id: str, url: str, caption: str = \"\"):\n return requests.get(\n f\"https://api.telegram.org/bot{bot_token}/sendPhoto\",\n params={\"chat_id\": chat_id, \"photo\": url, \"caption\": caption},\n ).url\n", + "language": "python3" + }, + "input_transform": { + "url": { + "type": "static", + "value": null + }, + "caption": { + "type": "static", + "value": "" + }, + "chat_id": { + "type": "static", + "value": null + }, + "bot_token": { + "type": "static", + "value": null + } + } + } + ] +})).unwrap(); + + for i in 0..10 { + println!("python flow iteration: {}", i); + let result = run_job_in_new_worker_until_complete( + &db, + JobPayload::RawFlow { value: flow.clone(), path: None }, + ) + .await; + + assert_eq!( + result, + serde_json::json!("https://api.telegram.org/botNone/sendPhoto?caption="), + "iteration: {i}" + ); + } + } + #[sqlx::test(fixtures("base"))] async fn test_python_job(db: DB) { initialize_tracing().await;