feat: cache relative imports (#6504)

* all

* all

* update

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;
 

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<UserDB>,
     Extension(db): Extension<DB>,
-    Path((w_id, token, path)): Path<(String, String, StripPath)>,
     Extension(cache): Extension<Arc<AuthCache>>,
+    Path((w_id, token, path)): Path<(String, String, StripPath)>,
+    Query(query): Query<RawScriptByPathQuery>,
 ) -> Result<String> {
     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<String>,
+    // used specifically for python to cache folders on import success to avoid extra db calls on package fetch
+    cache_folders: Option<bool>,
+}
 async fn raw_script_by_path(
     authed: ApiAuthed,
     Extension(user_db): Extension<UserDB>,
     Extension(db): Extension<DB>,
     Path((w_id, path)): Path<(String, StripPath)>,
+    Query(query): Query<RawScriptByPathQuery>,
 ) -> Result<String> {
-    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<UserDB>,
     Extension(db): Extension<DB>,
     Path((w_id, path)): Path<(String, StripPath)>,
+    Query(query): Query<RawScriptByPathQuery>,
 ) -> Result<String> {
-    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<String, String> = Cache::new(1000);
+    pub static ref CACHE_FOLDERS_PATH: Cache<String, i64> = 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<String> {
     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::<Vec<&str>>();
+    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(

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

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<String>,
     root_flow_id: Option<String>,
     scheduled_for: Option<chrono::DateTime<Utc>>,
+    runnable_id: Option<ScriptHash>,
 ) -> Vec<ContextualVariable> {
     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 {

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

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;

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

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;

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;