feat: add relative imports for python scripts (#1075)

* all

* all

backend/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.analysis.typeCheckingMode": "basic"
+}

backend/parsers/windmill-parser-py/src/lib.rs

@@ -175,7 +175,9 @@ fn constant_to_value(c: &Constant) -> serde_json::Value {
 static PYTHON_IMPORTS_REPLACEMENT: phf::Map<&'static str, &'static str> = phf_map! {
     "psycopg2" => "psycopg2-binary",
     "yaml" => "pyyaml",
-    "git" => "GitPython"
+    "git" => "GitPython",
+    "u" => "requests",
+    "f" => "requests"
 };
 
 fn replace_import(x: String) -> String {

backend/windmill-api/src/scripts.rs

@@ -469,10 +469,7 @@ async fn raw_script_by_path(
     Extension(user_db): Extension<UserDB>,
     Path((w_id, path)): Path<(String, StripPath)>,
 ) -> Result<String> {
-    let path = path
-        .to_path()
-        .strip_suffix(".ts")
-        .ok_or_else(|| Error::BadRequest("Raw script path must end with .ts".to_string()))?;
+    let path = path.to_path().split(".").next().unwrap_or_default();
     let mut tx = user_db.begin(&authed).await?;
 
     let content_o = sqlx::query_scalar!(

backend/windmill-worker/loader.py

@@ -0,0 +1,55 @@
+import sys
+import os
+from importlib.abc import MetaPathFinder, Loader
+from importlib.machinery import ModuleSpec, SourceFileLoader
+
+
+class WindmillLoader(Loader):
+    def __init__(self, path):
+        self.path = path
+
+    def create_module(self, spec):
+        return None
+
+    def exec_module(self, module):
+        module.__path__ = self.path
+        return None
+
+
+class WindmillFinder(MetaPathFinder):
+    @classmethod
+    def find_spec(cls, name, path, target=None):
+        splitted = name.split(".")
+
+        if splitted[0] != "f" and splitted[0] != "u":
+            return None
+        l = len(splitted)
+        if l <= 2:
+            return ModuleSpec(name, WindmillLoader(name))
+        elif l == 3:
+            script_path = "/".join(splitted)
+            import requests
+
+            url = f"{os.environ.get('BASE_INTERNAL_URL')}/api/w/{os.environ.get('WM_WORKSPACE')}/scripts/raw/p/{script_path}"
+
+            r = requests.get(
+                url, headers={"Authorization": f"Bearer {os.environ.get('WM_TOKEN')}"}
+            )
+
+            if r.status_code == 200:
+                folder = os.getcwd() + "/tmp/" + "/".join(splitted[:-1])
+                fullpath = folder + "/" + splitted[-1] + ".py"
+                os.makedirs(folder, exist_ok=True)
+                with open(fullpath, "w+") as f:
+                    f.write(r.text)
+                return ModuleSpec(name, SourceFileLoader(name, fullpath))
+            else:
+                print(r.text, r.status_code)
+                raise ImportError(f"Script {script_path} not found")
+        else:
+            raise ImportError(
+                "Import can only be done at the top level of a folder or user space"
+            )
+
+
+sys.meta_path.append(WindmillFinder)

backend/windmill-worker/nsjail/run.python3.config.proto

@@ -72,6 +72,13 @@ mount {
     is_bind: true
 }
 
+mount {
+    src: "{JOB_DIR}/loader.py"
+    dst: "/tmp/loader.py"
+    is_bind: true
+    mandatory: false
+}
+
 mount {
     src: "{JOB_DIR}/main.py"
     dst: "/tmp/main.py"

backend/windmill-worker/src/worker.rs

@@ -9,6 +9,8 @@
 use const_format::concatcp;
 use git_version::git_version;
 use itertools::Itertools;
+use lazy_static::lazy_static;
+use regex::Regex;
 use sqlx::{Pool, Postgres, Transaction};
 use std::{borrow::Borrow, collections::HashMap, io, panic, process::Stdio, time::Duration};
 use tracing::{trace_span, Instrument};
@@ -271,8 +273,10 @@ const NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT: &str = include_str!("../nsjail/download
 const NSJAIL_CONFIG_RUN_PYTHON3_CONTENT: &str = include_str!("../nsjail/run.python3.config.proto");
 const NSJAIL_CONFIG_RUN_GO_CONTENT: &str = include_str!("../nsjail/run.go.config.proto");
 const NSJAIL_CONFIG_RUN_BASH_CONTENT: &str = include_str!("../nsjail/run.bash.config.proto");
-
 const NSJAIL_CONFIG_RUN_DENO_CONTENT: &str = include_str!("../nsjail/run.deno.config.proto");
+
+const RELATIVE_PYTHON_LOADER: &str = include_str!("../loader.py");
+
 const MAX_LOG_SIZE: u32 = 200000;
 const GO_REQ_SPLITTER: &str = "//go.sum";
 const GIT_VERSION: &str = git_version!(args = ["--tag", "--always"], fallback = "unknown-version");
@@ -1538,6 +1542,10 @@ async fn create_args_and_out_file(
     Ok(())
 }
 
+lazy_static! {
+    static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(u|f)\."#).unwrap();
+}
+
 #[tracing::instrument(level = "trace", skip_all)]
 async fn handle_python_job(
     WorkerConfig { base_internal_url, base_url, disable_nuser, disable_nsjail, .. }: &WorkerConfig,
@@ -1608,12 +1616,17 @@ async fn handle_python_job(
 
     set_logs(logs, &job.id, db).await;
 
+    let relative_imports = RELATIVE_IMPORT_REGEX.is_match(&inner_content);
+
     let _ = write_file(job_dir, "inner.py", inner_content).await?;
+    if relative_imports {
+        let _ = write_file(job_dir, "loader.py", RELATIVE_PYTHON_LOADER).await?;
+    }
 
     let sig = windmill_parser_py::parse_python_signature(inner_content)?;
     let transforms = sig
         .args
-        .into_iter()
+        .iter()
         .map(|x| match x.typ {
             windmill_parser::Typ::Bytes => {
                 format!(
@@ -1636,11 +1649,36 @@ async fn handle_python_job(
         .join("");
     create_args_and_out_file(client, job, job_dir).await?;
 
+    let import_loader = if relative_imports {
+        "import loader"
+    } else {
+        ""
+    };
+    let import_base64 = if sig
+        .args
+        .iter()
+        .any(|x| x.typ == windmill_parser::Typ::Bytes)
+    {
+        "import base64"
+    } else {
+        ""
+    };
+    let import_datetime = if sig
+        .args
+        .iter()
+        .any(|x| x.typ == windmill_parser::Typ::Datetime)
+    {
+        "from datetime import datetime"
+    } else {
+        ""
+    };
     let wrapper_content: String = format!(
         r#"
 import json
 import base64
-from datetime import datetime
+{import_loader}
+{import_base64}
+{import_datetime}
 
 inner_script = __import__("inner")