feat: add support for contextual vars in SQL (#6791)

Author: hugocasa

backend/windmill-worker/src/bigquery_executor.rs

@@ -19,7 +19,8 @@ use serde::Deserialize;
 
 use crate::common::{build_args_values, resolve_job_timeout};
 use crate::common::{
-    build_http_client, s3_mode_args_to_worker_data, OccupancyMetrics, S3ModeWorkerData,
+    build_http_client, get_reserved_variables, s3_mode_args_to_worker_data, OccupancyMetrics,
+    S3ModeWorkerData,
 };
 use crate::handle_child::run_future_with_polling_update_job_poller;
 use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
@@ -313,6 +314,7 @@ pub async fn do_bigquery(
     worker_name: &str,
     column_order: &mut Option<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
+    parent_runnable_path: Option<String>,
 ) -> windmill_common::error::Result<Box<RawValue>> {
     let bigquery_args = build_args_values(job, client, conn).await?;
 
@@ -364,8 +366,15 @@ pub async fn do_bigquery(
         .map_err(|x| Error::ExecutionErr(x.to_string()))?
         .args;
 
-    let (query, args_to_skip) =
-        &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &bigquery_args)?;
+    let reserved_variables =
+        get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
+
+    let (query, args_to_skip) = &sanitize_and_interpolate_unsafe_sql_args(
+        query,
+        &sig,
+        &bigquery_args,
+        &reserved_variables,
+    )?;
 
     let queries = parse_sql_blocks(query);
 

backend/windmill-worker/src/duckdb_executor.rs

@@ -18,7 +18,7 @@ use windmill_parser_sql::{parse_duckdb_sig, parse_sql_blocks};
 use windmill_queue::{CanceledBy, MiniPulledJob};
 
 use crate::agent_workers::get_ducklake_from_agent_http;
-use crate::common::{build_args_values, OccupancyMetrics};
+use crate::common::{build_args_values, get_reserved_variables, OccupancyMetrics};
 use crate::handle_child::run_future_with_polling_update_job_poller;
 #[cfg(feature = "mysql")]
 use crate::mysql_executor::MysqlDatabase;
@@ -37,6 +37,7 @@ pub async fn do_duckdb(
     // TODO
     #[allow(unused_variables)] column_order_ref: &mut Option<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
+    parent_runnable_path: Option<String>,
 ) -> Result<Box<RawValue>> {
     let token = client.token.clone();
     let hidden_passwords = Arc::new(Mutex::new(Vec::<String>::new()));
@@ -48,7 +49,11 @@ pub async fn do_duckdb(
         let sig = parse_duckdb_sig(query)?.args;
         let mut job_args = build_args_values(job, client, conn).await?;
 
-        let (query, _) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args)?;
+        let reserved_variables =
+            get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
+
+        let (query, _) =
+            &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args, &reserved_variables)?;
         let query = transform_s3_uris(query).await?;
 
         let job_args = {

backend/windmill-worker/src/mssql_executor.rs

@@ -21,7 +21,9 @@ use windmill_parser_sql::{parse_db_resource, parse_mssql_sig, parse_s3_mode};
 use windmill_queue::MiniPulledJob;
 use windmill_queue::{append_logs, CanceledBy};
 
-use crate::common::{build_args_values, s3_mode_args_to_worker_data, OccupancyMetrics};
+use crate::common::{
+    build_args_values, get_reserved_variables, s3_mode_args_to_worker_data, OccupancyMetrics,
+};
 use crate::handle_child::run_future_with_polling_update_job_poller;
 use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
 use windmill_common::client::AuthedClient;
@@ -56,23 +58,25 @@ lazy_static::lazy_static! {
 
 pub async fn do_mssql(
     job: &MiniPulledJob,
-    client: &AuthedClient,
+    authed_client: &AuthedClient,
     query: &str,
     conn: &Connection,
     mem_peak: &mut i32,
     canceled_by: &mut Option<CanceledBy>,
     worker_name: &str,
     occupancy_metrics: &mut OccupancyMetrics,
     job_dir: &str,
+    parent_runnable_path: Option<String>,
 ) -> error::Result<Box<RawValue>> {
-    let mssql_args = build_args_values(job, client, conn).await?;
+    let mssql_args = build_args_values(job, authed_client, conn).await?;
 
     let inline_db_res_path = parse_db_resource(&query);
-    let s3 = parse_s3_mode(&query)?.map(|s3| s3_mode_args_to_worker_data(s3, client.clone(), job));
+    let s3 = parse_s3_mode(&query)?
+        .map(|s3| s3_mode_args_to_worker_data(s3, authed_client.clone(), job));
 
     let db_arg = if let Some(inline_db_res_path) = inline_db_res_path {
         Some(
-            client
+            authed_client
                 .get_resource_value_interpolated::<serde_json::Value>(
                     &inline_db_res_path,
                     Some(job.id.to_string()),
@@ -189,8 +193,11 @@ pub async fn do_mssql(
         .map_err(|x| Error::ExecutionErr(x.to_string()))?
         .args;
 
+    let reserved_variables =
+        get_reserved_variables(job, &authed_client.token, conn, parent_runnable_path).await?;
+
     let (query, args_to_skip) =
-        &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &mssql_args)?;
+        &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &mssql_args, &reserved_variables)?;
 
     let mut prepared_query = Query::new(query.to_owned());
     for arg in &sig {

backend/windmill-worker/src/mysql_executor.rs

@@ -26,7 +26,10 @@ use windmill_queue::CanceledBy;
 use windmill_queue::MiniPulledJob;
 
 use crate::{
-    common::{build_args_values, s3_mode_args_to_worker_data, OccupancyMetrics, S3ModeWorkerData},
+    common::{
+        build_args_values, get_reserved_variables, s3_mode_args_to_worker_data, OccupancyMetrics,
+        S3ModeWorkerData,
+    },
     handle_child::run_future_with_polling_update_job_poller,
     sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args,
 };
@@ -149,6 +152,7 @@ pub async fn do_mysql(
     worker_name: &str,
     column_order: &mut Option<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
+    parent_runnable_path: Option<String>,
 ) -> windmill_common::error::Result<Box<RawValue>> {
     let job_args = build_args_values(job, client, conn).await?;
 
@@ -198,7 +202,11 @@ pub async fn do_mysql(
         .map_err(|x| Error::ExecutionErr(x.to_string()))?
         .args;
 
-    let (query, args_to_skip) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args)?;
+    let reserved_variables =
+        get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
+
+    let (query, args_to_skip) =
+        &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args, &reserved_variables)?;
 
     let using_named_params = RE_ARG_MYSQL_NAMED.captures_iter(query).count() > 0;
 

backend/windmill-worker/src/oracledb_executor.rs

@@ -23,11 +23,11 @@ use windmill_queue::CanceledBy;
 
 use crate::{
     common::{
-        build_args_values, check_executor_binary_exists, s3_mode_args_to_worker_data,
-        OccupancyMetrics, S3ModeWorkerData,
+        build_args_values, check_executor_binary_exists, get_reserved_variables,
+        s3_mode_args_to_worker_data, OccupancyMetrics, S3ModeWorkerData,
     },
     handle_child::run_future_with_polling_update_job_poller,
-    sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args
+    sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args,
 };
 use windmill_common::client::AuthedClient;
 
@@ -343,6 +343,7 @@ pub async fn do_oracledb(
     worker_name: &str,
     column_order: &mut Option<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
+    parent_runnable_path: Option<String>,
 ) -> windmill_common::error::Result<Box<RawValue>> {
     check_executor_binary_exists(
         "the Oracle client lib",
@@ -381,7 +382,11 @@ pub async fn do_oracledb(
         .map_err(|x| Error::ExecutionErr(x.to_string()))?
         .args;
 
-    let (query, args_to_skip) = sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args)?;
+    let reserved_variables =
+        get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
+
+    let (query, args_to_skip) =
+        sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args, &reserved_variables)?;
 
     let (statement_values, errors) = get_statement_values(sig.clone(), &job_args, &args_to_skip);
 

backend/windmill-worker/src/pg_executor.rs

@@ -37,7 +37,8 @@ use windmill_parser_sql::{
 use windmill_queue::{CanceledBy, MiniPulledJob};
 
 use crate::common::{
-    build_args_values, s3_mode_args_to_worker_data, sizeof_val, OccupancyMetrics, S3ModeWorkerData,
+    build_args_values, get_reserved_variables, s3_mode_args_to_worker_data, sizeof_val,
+    OccupancyMetrics, S3ModeWorkerData,
 };
 use crate::handle_child::run_future_with_polling_update_job_poller;
 use crate::sanitized_sql_params::sanitize_and_interpolate_unsafe_sql_args;
@@ -186,6 +187,7 @@ pub async fn do_postgresql(
     worker_name: &str,
     column_order: &mut Option<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
+    parent_runnable_path: Option<String>,
 ) -> error::Result<Box<RawValue>> {
     let pg_args = build_args_values(job, client, conn).await?;
 
@@ -312,7 +314,11 @@ pub async fn do_postgresql(
 
     let sig = parse_pgsql_sig(&query).map_err(|x| Error::ExecutionErr(x.to_string()))?;
 
-    let (query, _) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig.args, &pg_args)?;
+    let reserved_variables =
+        get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
+
+    let (query, _) =
+        &sanitize_and_interpolate_unsafe_sql_args(query, &sig.args, &pg_args, &reserved_variables)?;
 
     let queries = parse_sql_blocks(query);
 

backend/windmill-worker/src/sanitized_sql_params.rs

@@ -1,11 +1,16 @@
 use anyhow::anyhow;
-use std::collections::HashMap;
+use regex::Regex;
+use std::collections::{HashMap, HashSet};
 
 use serde_json::Value;
 use windmill_common::error;
 use windmill_parser::Arg;
 use windmill_parser_sql::{SANITIZED_ENUM_STR, SANITIZED_RAW_STRING_STR};
 
+lazy_static::lazy_static! {
+    static ref RE_SQL_CONTEXTUAL_VAR: Regex = Regex::new(r"%%WM_[A-Z_]+%%").unwrap();
+}
+
 /// Identifier must be a continuous ASCII alphanumeric word, not starting with
 /// a number, that can contain underscores
 fn sanitize_identifier(arg: &Arg, input: &str) -> Result<(), error::Error> {
@@ -28,14 +33,39 @@ fn sanitize_identifier(arg: &Arg, input: &str) -> Result<(), error::Error> {
     }
 }
 
+fn replace_contextual_variables(
+    code: &mut String,
+    contextual_variables: &HashMap<String, String>,
+) -> () {
+    let vars = RE_SQL_CONTEXTUAL_VAR
+        .find_iter(&code)
+        .map(|m| m.as_str().to_string())
+        .collect::<HashSet<_>>();
+
+    for var_pattern in vars {
+        let var_name = var_pattern
+            .strip_prefix("%%")
+            .unwrap()
+            .strip_suffix("%%")
+            .unwrap();
+        let var_value = contextual_variables.get(var_name);
+        if let Some(var_value) = var_value {
+            *code = code.replace(&var_pattern, var_value);
+        }
+    }
+}
+
 pub fn sanitize_and_interpolate_unsafe_sql_args(
     code: &str,
     args: &Vec<Arg>,
     args_map: &HashMap<String, Value>,
+    contextual_variables: &HashMap<String, String>,
 ) -> Result<(String, Vec<String>), error::Error> {
     let mut ret = code.to_string();
     let mut args_to_skip = vec![];
 
+    replace_contextual_variables(&mut ret, contextual_variables);
+
     for arg in args {
         if let Some(typ) = &arg.otyp {
             let pattern = format!("%%{}%%", arg.name);

backend/windmill-worker/src/snowflake_executor.rs

@@ -20,7 +20,7 @@ use windmill_queue::{CanceledBy, MiniPulledJob, HTTP_CLIENT};
 
 use serde::{Deserialize, Serialize};
 
-use crate::common::build_args_values;
+use crate::common::{build_args_values, get_reserved_variables};
 use crate::common::{
     build_http_client, resolve_job_timeout, s3_mode_args_to_worker_data, OccupancyMetrics,
     S3ModeWorkerData,
@@ -132,12 +132,14 @@ fn do_snowflake_inner<'a>(
     skip_collect: bool,
     http_client: &'a Client,
     s3: Option<S3ModeWorkerData>,
+    reserved_variables: &HashMap<String, String>,
 ) -> windmill_common::error::Result<BoxFuture<'a, windmill_common::error::Result<Box<RawValue>>>> {
     let sig = parse_snowflake_sig(&query)
         .map_err(|x| Error::ExecutionErr(x.to_string()))?
         .args;
 
-    let (query, args_to_skip) = &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args)?;
+    let (query, args_to_skip) =
+        &sanitize_and_interpolate_unsafe_sql_args(query, &sig, &job_args, reserved_variables)?;
 
     body.insert("statement".to_string(), json!(query));
 
@@ -282,6 +284,7 @@ pub async fn do_snowflake(
     worker_name: &str,
     column_order: &mut Option<Vec<String>>,
     occupancy_metrics: &mut OccupancyMetrics,
+    parent_runnable_path: Option<String>,
 ) -> windmill_common::error::Result<Box<RawValue>> {
     let snowflake_args = build_args_values(job, client, conn).await?;
 
@@ -403,6 +406,9 @@ pub async fn do_snowflake(
 
     let http_client = build_http_client(timeout_duration)?;
 
+    let reserved_variables =
+        get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
+
     let result_f = if queries.len() > 1 {
         let futures = queries
             .iter()
@@ -419,6 +425,7 @@ pub async fn do_snowflake(
                     annotations.return_last_result && i < queries.len() - 1,
                     &http_client,
                     s3.clone(),
+                    &reserved_variables,
                 )
             })
             .collect::<windmill_common::error::Result<Vec<_>>>()?;
@@ -449,6 +456,7 @@ pub async fn do_snowflake(
             false,
             &http_client,
             s3.clone(),
+            &reserved_variables,
         )?
     };
     let r = run_future_with_polling_update_job_poller(

backend/windmill-worker/src/worker.rs

@@ -3016,6 +3016,7 @@ async fn handle_code_execution_job(
             worker_name,
             column_order,
             occupancy_metrics,
+            parent_runnable_path,
         )
         .await;
     } else if language == Some(ScriptLang::Mysql) {
@@ -3035,6 +3036,7 @@ async fn handle_code_execution_job(
             worker_name,
             column_order,
             occupancy_metrics,
+            parent_runnable_path,
         )
         .await;
     } else if language == Some(ScriptLang::Bigquery) {
@@ -3065,6 +3067,7 @@ async fn handle_code_execution_job(
                 worker_name,
                 column_order,
                 occupancy_metrics,
+                parent_runnable_path,
             )
             .await;
         }
@@ -3088,6 +3091,7 @@ async fn handle_code_execution_job(
                 worker_name,
                 column_order,
                 occupancy_metrics,
+                parent_runnable_path,
             )
             .await;
         }
@@ -3119,6 +3123,7 @@ async fn handle_code_execution_job(
                 worker_name,
                 occupancy_metrics,
                 job_dir,
+                parent_runnable_path,
             )
             .await;
         }
@@ -3150,6 +3155,7 @@ async fn handle_code_execution_job(
                 worker_name,
                 column_order,
                 occupancy_metrics,
+                parent_runnable_path,
             )
             .await;
         }
@@ -3174,6 +3180,7 @@ async fn handle_code_execution_job(
                 worker_name,
                 column_order,
                 occupancy_metrics,
+                parent_runnable_path,
             )
             .await;
         }