From 5347cd46a996b4cf48a96fbb873e4d029ca4f75f Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 9 Jan 2023 20:49:30 +0800 Subject: [PATCH] feat: add relative imports for python scripts (#1075) * all * all --- backend/.vscode/settings.json | 3 + backend/parsers/windmill-parser-py/src/lib.rs | 4 +- backend/windmill-api/src/scripts.rs | 5 +- backend/windmill-worker/loader.py | 55 +++++++++++++++++++ .../nsjail/run.python3.config.proto | 7 +++ backend/windmill-worker/src/worker.rs | 44 ++++++++++++++- 6 files changed, 110 insertions(+), 8 deletions(-) create mode 100644 backend/.vscode/settings.json create mode 100644 backend/windmill-worker/loader.py diff --git a/backend/.vscode/settings.json b/backend/.vscode/settings.json new file mode 100644 index 0000000000..dc3f727cad --- /dev/null +++ b/backend/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.analysis.typeCheckingMode": "basic" +} diff --git a/backend/parsers/windmill-parser-py/src/lib.rs b/backend/parsers/windmill-parser-py/src/lib.rs index 36ac100e06..fac1c78810 100644 --- a/backend/parsers/windmill-parser-py/src/lib.rs +++ b/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 { diff --git a/backend/windmill-api/src/scripts.rs b/backend/windmill-api/src/scripts.rs index c1f43fb85e..4c5d8be3e7 100644 --- a/backend/windmill-api/src/scripts.rs +++ b/backend/windmill-api/src/scripts.rs @@ -469,10 +469,7 @@ async fn raw_script_by_path( Extension(user_db): Extension, Path((w_id, path)): Path<(String, StripPath)>, ) -> Result { - 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!( diff --git a/backend/windmill-worker/loader.py b/backend/windmill-worker/loader.py new file mode 100644 index 0000000000..dda3776b98 --- /dev/null +++ b/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) diff --git a/backend/windmill-worker/nsjail/run.python3.config.proto b/backend/windmill-worker/nsjail/run.python3.config.proto index 0c05a08cc3..5747d3cb3f 100644 --- a/backend/windmill-worker/nsjail/run.python3.config.proto +++ b/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" diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs index 8316e85eee..5ae8d674e4 100644 --- a/backend/windmill-worker/src/worker.rs +++ b/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")