diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index cc0d7fb628..290ddd5376 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/backend/parsers/windmill-parser-py/src/lib.rs @@ -187,6 +187,9 @@ static PYTHON_IMPORTS_REPLACEMENT: phf::Map<&'static str, &'static str> = phf_ma }; fn replace_import(x: String) -> String { + if x.starts_with('.') { + return "requests".to_string(); + } PYTHON_IMPORTS_REPLACEMENT .get(&x) .map(|x| x.to_owned()) @@ -224,12 +227,23 @@ pub fn parse_python_imports(code: &str) -> error::Result> { StmtKind::Import { names } => Some( names .into_iter() - .map(|x| x.node.name.split('.').next().unwrap_or("").to_string()) + .map(|x| { + let name = x.node.name; + if name.starts_with('.') { + ".".to_string() + } else { + name.split('.').next().unwrap_or("").to_string() + } + }) .map(replace_import) .collect::>(), ), StmtKind::ImportFrom { level: _, module: Some(mod_), names: _ } => { - let imprt = mod_.split('.').next().unwrap_or("").replace("_", "-"); + let imprt = if mod_.starts_with('.') { + mod_.split('.').next().unwrap_or("").replace("_", "-") + } else { + ".".to_string() + }; Some(vec![replace_import(imprt)]) } diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 527d38624e..b8cc247aac 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -658,7 +658,7 @@ impl QueuedJob { self.script_path .as_ref() .map(String::as_str) - .unwrap_or("NO_FLOW_PATH") + .unwrap_or("tmp/main") } } diff --git a/backend/windmill-worker/nsjail/run.python3.config.proto b/backend/windmill-worker/nsjail/run.python3.config.proto index ef4e9a77d8..33919ae1fc 100644 --- a/backend/windmill-worker/nsjail/run.python3.config.proto +++ b/backend/windmill-worker/nsjail/run.python3.config.proto @@ -66,8 +66,8 @@ mount { } mount { - src: "{JOB_DIR}/inner.py" - dst: "/tmp/inner.py" + src: "{JOB_DIR}/{MAIN}.py" + dst: "/tmp/{MAIN}.py" is_bind: true } @@ -79,8 +79,8 @@ mount { } mount { - src: "{JOB_DIR}/main.py" - dst: "/tmp/main.py" + src: "{JOB_DIR}/wrapper.py" + dst: "/tmp/wrapper.py" is_bind: true } diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 623dbd6447..9f6d8e823c 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/backend/windmill-worker/src/worker.rs @@ -1706,7 +1706,7 @@ async fn create_args_and_out_file( } lazy_static! { - static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(u|f)\."#).unwrap(); + static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(((u|f)\.)|\.)"#).unwrap(); } #[tracing::instrument(level = "trace", skip_all)] @@ -1777,9 +1777,14 @@ async fn handle_python_job( let relative_imports = RELATIVE_IMPORT_REGEX.is_match(&inner_content); - let _ = write_file(job_dir, "inner.py", inner_content).await?; + let script_path_splitted = &job.script_path().split("/"); + let dirs = script_path_splitted.clone().take(script_path_splitted.clone().count() - 1).join("/"); + let last = script_path_splitted.clone().last().unwrap(); + let module_dir = format!("{}/{}", job_dir, dirs ); + tokio::fs::create_dir_all(format!("{module_dir}/")).await?; + let _ = write_file(&module_dir, &format!("{last}.py"), inner_content).await?; if relative_imports { - let _ = write_file(job_dir, "loader.py", RELATIVE_PYTHON_LOADER).await?; + let _ = write_file(&job_dir, "loader.py", RELATIVE_PYTHON_LOADER).await?; } let sig = windmill_parser_py::parse_python_signature(inner_content)?; @@ -1840,6 +1845,7 @@ async fn handle_python_job( .join("\n") }; + let module_dir_dot = dirs.replace("/", "."); let wrapper_content: String = format!( r#" import json @@ -1848,8 +1854,8 @@ import json {import_datetime} import traceback import sys +from {module_dir_dot} import {last} as inner_script -inner_script = __import__("inner") with open("args.json") as f: kwargs = json.load(f, strict=False) @@ -1876,7 +1882,7 @@ except Exception as e: sys.exit(1) "#, ); - write_file(job_dir, "main.py", &wrapper_content).await?; + write_file(job_dir, "wrapper.py", &wrapper_content).await?; let mut reserved_variables = get_reserved_variables(job, &token, db).await?; let additional_python_paths_folders = additional_python_paths.iter().join(":"); @@ -1904,6 +1910,7 @@ mount {{ .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()) .replace("{SHARED_MOUNT}", shared_mount) .replace("{SHARED_DEPENDENCIES}", shared_deps.as_str()) + .replace("{MAIN}", format!("{dirs}/{last}").as_str()) .replace( "{ADDITIONAL_PYTHON_PATHS}", additional_python_paths_folders.as_str(), @@ -1935,7 +1942,8 @@ mount {{ "--", PYTHON_PATH.as_str(), "-u", - "/tmp/main.py", + "-m", + "wrapper", ]) .stdout(Stdio::piped()) .stderr(Stdio::piped()) @@ -1947,7 +1955,7 @@ mount {{ .envs(reserved_variables) .env("PATH", PATH_ENV.as_str()) .env("BASE_INTERNAL_URL", base_internal_url) - .args(vec!["-u", "main.py"]) + .args(vec!["-u", "-m", "wrapper"]) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .spawn()?