From 16912b484d7459cdf8cb112cac0800d3fd61512c Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Mon, 1 Sep 2025 14:55:08 +0000 Subject: [PATCH] feat: cache relative imports (#6504) * all * all * update --- backend/windmill-api/src/resources.rs | 1 + backend/windmill-api/src/scripts.rs | 73 +++++++++++++++++-- backend/windmill-api/src/variables.rs | 2 + backend/windmill-common/src/variables.rs | 10 ++- backend/windmill-worker/loader.py | 38 +++++++--- backend/windmill-worker/src/bun_executor.rs | 1 + backend/windmill-worker/src/common.rs | 1 + backend/windmill-worker/src/deno_executor.rs | 1 + .../windmill-worker/src/python_executor.rs | 2 + 9 files changed, 112 insertions(+), 17 deletions(-) diff --git a/backend/windmill-api/src/resources.rs b/backend/windmill-api/src/resources.rs index 571ecdfac9..b474e710db 100644 --- a/backend/windmill-api/src/resources.rs +++ b/backend/windmill-api/src/resources.rs @@ -592,6 +592,7 @@ pub async fn transform_json_value<'c>( job.flow_step_id.clone(), job.root_job.map(|x| x.to_string()), Some(job.scheduled_for.clone()), + None, ) .await; diff --git a/backend/windmill-api/src/scripts.rs b/backend/windmill-api/src/scripts.rs index 49765660ab..ef687bf418 100644 --- a/backend/windmill-api/src/scripts.rs +++ b/backend/windmill-api/src/scripts.rs @@ -28,6 +28,7 @@ use axum::{ }; use hyper::StatusCode; use itertools::Itertools; +use quick_cache::sync::Cache; use serde::{Deserialize, Serialize}; use serde_json::json; use serde_json::value::RawValue; @@ -1361,8 +1362,9 @@ async fn toggle_workspace_error_handler( async fn get_tokened_raw_script_by_path( Extension(user_db): Extension, Extension(db): Extension, - Path((w_id, token, path)): Path<(String, String, StripPath)>, Extension(cache): Extension>, + Path((w_id, token, path)): Path<(String, String, StripPath)>, + Query(query): Query, ) -> Result { let authed = cache .get_authed(Some(w_id.clone()), &token) @@ -1373,6 +1375,7 @@ async fn get_tokened_raw_script_by_path( Extension(user_db), Extension(db), Path((w_id, path)), + Query(query), ) .await; } @@ -1381,13 +1384,21 @@ async fn get_empty_ts_script_by_path() -> String { return String::new(); } +#[derive(Deserialize)] +struct RawScriptByPathQuery { + // used to make cache immutable with respect to importer + cache_key: Option, + // used specifically for python to cache folders on import success to avoid extra db calls on package fetch + cache_folders: Option, +} async fn raw_script_by_path( authed: ApiAuthed, Extension(user_db): Extension, Extension(db): Extension, Path((w_id, path)): Path<(String, StripPath)>, + Query(query): Query, ) -> Result { - raw_script_by_path_internal(path, user_db, db, authed, w_id, false).await + raw_script_by_path_internal(path, user_db, db, authed, w_id, false, query).await } async fn raw_script_by_path_unpinned( @@ -1395,8 +1406,9 @@ async fn raw_script_by_path_unpinned( Extension(user_db): Extension, Extension(db): Extension, Path((w_id, path)): Path<(String, StripPath)>, + Query(query): Query, ) -> Result { - raw_script_by_path_internal(path, user_db, db, authed, w_id, true).await + raw_script_by_path_internal(path, user_db, db, authed, w_id, true, query).await } lazy_static::lazy_static! { @@ -1404,6 +1416,12 @@ lazy_static::lazy_static! { std::env::var("DEBUG_RAW_SCRIPT_ENDPOINTS").is_ok(); } +lazy_static::lazy_static! { + pub static ref RAW_SCRIPT_CACHE: Cache = Cache::new(1000); + pub static ref CACHE_FOLDERS_PATH: Cache = Cache::new(1000); + +} + async fn raw_script_by_path_internal( path: StripPath, user_db: UserDB, @@ -1411,9 +1429,18 @@ async fn raw_script_by_path_internal( authed: ApiAuthed, w_id: String, unpin: bool, + query: RawScriptByPathQuery, ) -> Result { let path = path.to_path(); check_scopes(&authed, || format!("scripts:read:{}", path))?; + let cache_path = query.cache_key.map(|x| format!("{w_id}:{path}:{x}")); + if let Some(cache_path) = cache_path.clone() { + let cached_content = RAW_SCRIPT_CACHE.get(&cache_path); + if let Some(cached_content) = cached_content { + return Ok(cached_content); + } + } + if !path.ends_with(".py") && !path.ends_with(".ts") && !path.ends_with(".go") @@ -1431,6 +1458,27 @@ async fn raw_script_by_path_internal( .trim_end_matches(".ts") .trim_end_matches(".go") .trim_end_matches(".sh"); + + // folder cache is only useful for python given it needs to recuse over all intermediate folders to find the package. + // When a script exists in a folder, we can cache the fact that the folder exists to avoid extra db calls. + let mut split_path = path.split("/").collect::>(); + let folder_path = if query.cache_folders.is_some() && split_path.len() > 2 { + Some(format!("{w_id}:{path}/")) + } else { + None + }; + + let has_folder_cache = folder_path.is_some(); + if let Some(cache_folders) = folder_path { + let cached_content = CACHE_FOLDERS_PATH.get(&cache_folders); + if let Some(cached_ts) = cached_content { + if cached_ts >= chrono::Utc::now().timestamp() - 300 { + // 5 minutes + return Ok("WINDMILL_IS_FOLDER".to_string()); + } + } + } + let mut tx = user_db.begin(&authed).await?; let content_o = sqlx::query_scalar!( @@ -1484,11 +1532,24 @@ async fn raw_script_by_path_internal( let content = not_found_if_none(content_o, "Script", path)?; - if unpin { - return Ok(remove_pinned_imports(&content)?); + let content = if unpin { + remove_pinned_imports(&content)? } else { - return Ok(content); + content + }; + + if has_folder_cache { + while split_path.len() >= 2 { + split_path.pop(); + let npath = split_path.join("/"); + CACHE_FOLDERS_PATH.insert(format!("{w_id}:{npath}/"), chrono::Utc::now().timestamp()); + } } + + if let Some(cache_path) = cache_path { + RAW_SCRIPT_CACHE.insert(cache_path, content.clone()); + } + Ok(content) } async fn exists_script_by_path( diff --git a/backend/windmill-api/src/variables.rs b/backend/windmill-api/src/variables.rs index 12791753d0..f7476d9f26 100644 --- a/backend/windmill-api/src/variables.rs +++ b/backend/windmill-api/src/variables.rs @@ -26,6 +26,7 @@ use windmill_audit::ActionKind; use windmill_common::{ db::UserDB, error::{Error, JsonResult, Result}, + scripts::ScriptHash, utils::{not_found_if_none, paginate, Pagination, StripPath, WarnAfterExt}, variables::{ build_crypt, get_reserved_variables, ContextualVariable, CreateVariable, ListableVariable, @@ -77,6 +78,7 @@ async fn list_contextual_variables( Some("c".to_string()), Some("017e0ad5-f499-73b6-5488-92a61c5196dd".to_string()), Some(chrono::offset::Utc::now()), + Some(ScriptHash(1234567890)), ) .await .to_vec(), diff --git a/backend/windmill-common/src/variables.rs b/backend/windmill-common/src/variables.rs index a1ab8e45cd..f43292755e 100644 --- a/backend/windmill-common/src/variables.rs +++ b/backend/windmill-common/src/variables.rs @@ -7,6 +7,7 @@ */ use crate::error; +use crate::scripts::ScriptHash; use crate::utils::WarnAfterExt; use crate::worker::Connection; use crate::{worker::WORKER_GROUP, BASE_URL, DB}; @@ -211,6 +212,7 @@ pub async fn get_reserved_variables( step_id: Option, root_flow_id: Option, scheduled_for: Option>, + runnable_id: Option, ) -> Vec { let state_path = { let trigger = if schedule_path.is_some() { @@ -366,7 +368,13 @@ pub async fn get_reserved_variables( ContextualVariable { name: "WM_WORKER_GROUP".to_string(), value: WORKER_GROUP.clone(), - description: "name of the worker group the job is running on".to_string(), + description: "Name of the worker group the job is running on".to_string(), + is_custom: false, + }, + ContextualVariable { + name: "WM_RUNNABLE_ID".to_string(), + value: runnable_id.map(|x| x.to_string()).unwrap_or_else(|| "".to_string()), + description: "Hash of the script. Useful as cache key for cache that should be runnable specific.".to_string(), is_custom: false, }, ].into_iter().chain(custom_envs.into_iter().map(|(name, value)| ContextualVariable { diff --git a/backend/windmill-worker/loader.py b/backend/windmill-worker/loader.py index 3d40c93b42..646ec64029 100644 --- a/backend/windmill-worker/loader.py +++ b/backend/windmill-worker/loader.py @@ -2,6 +2,7 @@ import sys import os from importlib.abc import MetaPathFinder, Loader from importlib.machinery import ModuleSpec, SourceFileLoader +import urllib.response class WindmillLoader(Loader): @@ -27,7 +28,15 @@ class WindmillFinder(MetaPathFinder): if l <= 2: return ModuleSpec(name, WindmillLoader(name)) elif l > 2: + script_path = "/".join(splitted) + folder = os.getcwd() + "/tmp/" + "/".join(splitted[:-1]) + fullpath = folder + "/" + splitted[-1] + ".py" + + if os.path.exists(fullpath): + return ModuleSpec(name, SourceFileLoader(name, fullpath)) + + import urllib.parse import urllib.request @@ -35,20 +44,29 @@ class WindmillFinder(MetaPathFinder): "Authorization": f"Bearer {os.environ.get('WM_TOKEN')}", "User-Agent": "windmill/beta" } - url = f"{os.environ.get('BASE_INTERNAL_URL')}/api/w/{os.environ.get('WM_WORKSPACE')}/scripts/raw/p/{script_path}.py" + + query_params = "?cache_folders=true" + runnable_id = os.environ.get('WM_RUNNABLE_ID') + if runnable_id: + query_params += f"&cache_key={runnable_id}" + url = f"{os.environ.get('BASE_INTERNAL_URL')}/api/w/{os.environ.get('WM_WORKSPACE')}/scripts/raw/p/{script_path}.py{query_params}" req = urllib.request.Request(url, None, headers) try: with urllib.request.urlopen(req) as response: - r = response.read().decode("utf-8") - 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) - return ModuleSpec(name, SourceFileLoader(name, fullpath)) - except: - # raise ImportError(f"Script {script_path} not found") + os.makedirs(folder, exist_ok=True) + r = response.read().decode("utf-8") + if r == "WINDMILL_IS_FOLDER": + return ModuleSpec(name, WindmillLoader(name)) + with open(fullpath, "w+") as f: + f.write(r) + return ModuleSpec(name, SourceFileLoader(name, fullpath)) + except urllib.error.HTTPError as e: + if e.code != 404: + print(f"Error fetching script {script_path}: HTTP {e.code} - {e.reason}") + return ModuleSpec(name, WindmillLoader(name)) + except Exception as e: + print(f"Error fetching script {script_path}: {e}") return ModuleSpec(name, WindmillLoader(name)) diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 9cad2ebba7..d0ad991bf2 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -1589,6 +1589,7 @@ pub async fn start_worker( None, None, None, + None, ) .await; let context_envs = build_envs_map(context.to_vec()).await; diff --git a/backend/windmill-worker/src/common.rs b/backend/windmill-worker/src/common.rs index 11718f3cf5..1df44c34df 100644 --- a/backend/windmill-worker/src/common.rs +++ b/backend/windmill-worker/src/common.rs @@ -464,6 +464,7 @@ pub async fn get_reserved_variables( job.flow_step_id.clone(), job.flow_innermost_root_job.clone().map(|x| x.to_string()), Some(job.scheduled_for.clone()), + job.runnable_id, ) .await .to_vec(); diff --git a/backend/windmill-worker/src/deno_executor.rs b/backend/windmill-worker/src/deno_executor.rs index 4cc6649525..607e6375d2 100644 --- a/backend/windmill-worker/src/deno_executor.rs +++ b/backend/windmill-worker/src/deno_executor.rs @@ -544,6 +544,7 @@ pub async fn start_worker( None, None, None, + None, ) .await; let context_envs = build_envs_map(context.to_vec()).await; diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index eb6a6e7f30..70f980052b 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -2145,6 +2145,7 @@ pub async fn start_worker( None, None, None, + None, ) .await .to_vec(); @@ -2264,6 +2265,7 @@ for line in sys.stdin: None, None, None, + None, ) .await;