feat(backend): allow relative imports for python

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<Vec<String>> {
                     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::<Vec<String>>(),
                     ),
                     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)])
                     }

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")
     }
 }
 

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
 }
 

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()?