diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index ef4cfa0f46..ea50393018 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -516,6 +516,7 @@ fn parse_file(path: &str) -> Option { pub struct PythonAnnotations { pub no_cache: bool, pub no_postinstall: bool, + pub skip_result_postprocessing: bool, pub py310: bool, pub py311: bool, pub py312: bool, diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs index a776bbe7c7..7721a2fc2b 100644 --- a/backend/windmill-worker/src/python_executor.rs +++ b/backend/windmill-worker/src/python_executor.rs @@ -842,6 +842,8 @@ pub async fn handle_python_job( ) -> windmill_common::error::Result> { let script_path = crate::common::use_flow_root_path(job.runnable_path()); + let annotations = PythonAnnotations::parse(inner_content); + let (py_version, mut additional_python_paths) = handle_python_deps( job_dir, requirements_o, @@ -856,10 +858,10 @@ pub async fn handle_python_job( canceled_by, &mut Some(occupancy_metrics), precomputed_agent_info, + annotations, ) .await?; - let PythonAnnotations { no_postinstall, .. } = PythonAnnotations::parse(inner_content); tracing::debug!("Finished handling python dependencies"); let python_path = get_python_path( py_version, @@ -872,7 +874,7 @@ pub async fn handle_python_job( ) .await?; - if !no_postinstall { + if !annotations.no_postinstall { if let Err(e) = postinstall(&mut additional_python_paths, job_dir, job, conn).await { tracing::error!("Postinstall stage has failed. Reason: {e}"); } @@ -938,6 +940,8 @@ pub async fn handle_python_job( "".to_string() }; + let postprocessor = get_result_postprocessor(annotations.skip_result_postprocessing); + let os_main_override = if let Some(main_override) = main_name.as_ref() { format!("os.environ[\"MAIN_OVERRIDE\"] = \"{main_override}\"\n") } else { @@ -984,7 +988,8 @@ def res_to_json(res): for k, v in res.items(): if type(v).__name__ == 'bytes': res[k] = to_b_64(v) - return re.sub(replace_invalid_fields, ' null ', json.dumps(res, separators=(',', ':'), default=str).replace('\n', '')) + unprocessed = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '') + return {postprocessor} try: {preprocessor} @@ -1428,6 +1433,7 @@ async fn handle_python_deps( canceled_by: &mut Option, occupancy_metrics: &mut Option<&mut OccupancyMetrics>, precomputed_agent_info: Option, + annotations: PythonAnnotations, ) -> error::Result<(PyVersion, Vec)> { create_dependencies_dir(job_dir).await; @@ -1445,7 +1451,6 @@ async fn handle_python_deps( let mut annotated_pyv_numeric = None; let is_deployed = requirements_o.is_some(); let instance_pyv = PyVersion::from_instance_version(job_id, w_id, conn).await; - let annotations = windmill_common::worker::PythonAnnotations::parse(inner_content); let requirements = match requirements_o { Some(r) => r, None => { @@ -2326,6 +2331,15 @@ fn get_pyv_from_requirements_lines(requirements_lines: &[&str]) -> PyVersion { } } +// Returns code snippet that needs to be injected into wrapper to post-process results or leave unprocessed +fn get_result_postprocessor<'a>(skip: bool) -> &'a str { + if skip { + "unprocessed" + } else { + "re.sub(replace_invalid_fields, ' null ', unprocessed)" + } +} + #[cfg(feature = "enterprise")] use crate::JobCompletedSender; #[cfg(feature = "enterprise")] @@ -2373,6 +2387,7 @@ pub async fn start_worker( .await .to_vec(); + let annotations = PythonAnnotations::parse(inner_content); let context_envs = build_envs_map(context).await; let (_, additional_python_paths) = handle_python_deps( job_dir, @@ -2388,6 +2403,7 @@ pub async fn start_worker( &mut canceled_by, &mut None, None, + annotations, ) .await?; @@ -2405,6 +2421,7 @@ pub async fn start_worker( ) = prepare_wrapper(job_dir, false, None, None, inner_content, script_path).await?; { + let postprocessor = get_result_postprocessor(annotations.skip_result_postprocessing); let indented_transforms = transforms .lines() .map(|x| format!(" {}", x)) @@ -2456,7 +2473,8 @@ for line in sys.stdin: for k, v in res.items(): if type(v).__name__ == 'bytes': res[k] = to_b_64(v) - res_json = re.sub(replace_invalid_fields, ' null ', json.dumps(res, separators=(',', ':'), default=str).replace('\n', '')) + unprocessed = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '') + res_json = {postprocessor} sys.stdout.write("wm_res[success]:" + res_json + "\n") except BaseException as e: exc_type, exc_value, exc_traceback = sys.exc_info()