from __future__ import annotations import atexit import datetime as dt import functools from io import BufferedReader, BytesIO import logging import os import random import time import warnings import json from json import JSONDecodeError from typing import Callable, Dict, Any, Union, Literal, Optional import re import httpx from .s3_reader import S3BufferedReader, bytes_generator from .s3_types import ( Boto3ConnectionSettings, DuckDbConnectionSettings, PolarsConnectionSettings, S3Object, ) _client: "Windmill | None" = None logger = logging.getLogger("windmill_client") JobStatus = Literal["RUNNING", "WAITING", "COMPLETED"] def _sign_s3_objects_body(s3_objects: list, expiry_secs: int | None) -> dict: # `expiry_secs` is optional but not nullable in the spec, so omit it rather than # sending an explicit null a validating gateway would reject. body: dict = {"s3_objects": s3_objects} if expiry_secs is not None: body["expiry_secs"] = expiry_secs return body class Windmill: """Windmill client for interacting with the Windmill API.""" def __init__(self, base_url=None, token=None, workspace=None, verify=True): """Initialize the Windmill client. Args: base_url: API base URL (defaults to BASE_INTERNAL_URL or WM_BASE_URL env) token: Authentication token (defaults to WM_TOKEN env) workspace: Workspace ID (defaults to WM_WORKSPACE env) verify: Whether to verify SSL certificates """ base = ( base_url or os.environ.get("BASE_INTERNAL_URL") or os.environ.get("WM_BASE_URL") ) self.base_url = f"{base}/api" self.token = token or os.environ.get("WM_TOKEN") self.headers = { "Content-Type": "application/json", "Authorization": f"Bearer {self.token}", } self.verify = verify self.client = self.get_client() self.workspace = workspace or os.environ.get("WM_WORKSPACE") self.path = os.environ.get("WM_JOB_PATH") self.mocked_api = self.get_mocked_api() assert self.workspace, ( f"workspace required as an argument or as WM_WORKSPACE environment variable" ) def worker_has_internal_server(self) -> bool: return bool( re.match(r"^https?://(localhost|127\.0\.0\.1)(:|/|$)", self.base_url or "") ) def get_mocked_api(self) -> Optional[dict]: mocked_path = os.environ.get("WM_MOCKED_API_FILE") if not mocked_path: return None logger.info("Using mocked API from %s", mocked_path) mocked_api = {"variables": {}, "resources": {}} try: with open(mocked_path, "r") as f: incoming_mocked_api = json.load(f) mocked_api = {**mocked_api, **incoming_mocked_api} except Exception as e: logger.warning( "Error parsing mocked API file at path %s Using empty mocked API.", mocked_path, ) logger.debug(e) return mocked_api def get_client(self) -> httpx.Client: """Get the HTTP client instance. Returns: Configured httpx.Client for API requests """ return httpx.Client( base_url=self.base_url, headers=self.headers, verify=self.verify, timeout=httpx.Timeout(900.0), ) def get(self, endpoint, raise_for_status=True, **kwargs) -> httpx.Response: """Make an HTTP GET request to the Windmill API. Args: endpoint: API endpoint path raise_for_status: Whether to raise an exception on HTTP errors **kwargs: Additional arguments passed to httpx.get Returns: HTTP response object """ endpoint = endpoint.lstrip("/") resp = self.client.get(f"/{endpoint}", **kwargs) if raise_for_status: try: resp.raise_for_status() except httpx.HTTPStatusError as err: error = f"{err.request.url}: {err.response.status_code}, {err.response.text}" logger.error(error) raise Exception(error) return resp def post(self, endpoint, raise_for_status=True, **kwargs) -> httpx.Response: """Make an HTTP POST request to the Windmill API. Args: endpoint: API endpoint path raise_for_status: Whether to raise an exception on HTTP errors **kwargs: Additional arguments passed to httpx.post Returns: HTTP response object """ endpoint = endpoint.lstrip("/") resp = self.client.post(f"/{endpoint}", **kwargs) if raise_for_status: try: resp.raise_for_status() except httpx.HTTPStatusError as err: error = f"{err.request.url}: {err.response.status_code}, {err.response.text}" logger.error(error) raise Exception(error) return resp def create_token(self, duration=dt.timedelta(days=1)) -> str: """Create a new authentication token. Args: duration: Token validity duration (default: 1 day) Returns: New authentication token string """ endpoint = "/users/tokens/create" payload = { "label": f"refresh {time.time()}", "expiration": (dt.datetime.now() + duration).strftime("%Y-%m-%dT%H:%M:%SZ"), } return self.post(endpoint, json=payload).text def run_script_async( self, path: str = None, hash_: str = None, args: dict = None, scheduled_in_secs: int = None, tag: str = None, ) -> str: """Create a script job and return its job id. .. deprecated:: Use run_script_by_path_async or run_script_by_hash_async instead. """ logging.warning( "run_script_async is deprecated. Use run_script_by_path_async or run_script_by_hash_async instead.", ) assert not (path and hash_), "path and hash_ are mutually exclusive" return self._run_script_async_internal(path=path, hash_=hash_, args=args, scheduled_in_secs=scheduled_in_secs, tag=tag) def _run_script_async_internal( self, path: str = None, hash_: str = None, args: dict = None, scheduled_in_secs: int = None, tag: str = None, ) -> str: """Internal helper for running scripts asynchronously.""" args = args or {} params = {"scheduled_in_secs": scheduled_in_secs} if scheduled_in_secs else {} if tag: params["tag"] = tag if os.environ.get("WM_JOB_ID"): params["parent_job"] = os.environ.get("WM_JOB_ID") if os.environ.get("WM_ROOT_FLOW_JOB_ID"): params["root_job"] = os.environ.get("WM_ROOT_FLOW_JOB_ID") if path: endpoint = f"/w/{self.workspace}/jobs/run/p/{path}" elif hash_: endpoint = f"/w/{self.workspace}/jobs/run/h/{hash_}" else: raise Exception("path or hash_ must be provided") return self.post(endpoint, json=args, params=params).text def run_script_by_path_async( self, path: str, args: dict = None, scheduled_in_secs: int = None, tag: str = None, ) -> str: """Create a script job by path and return its job id.""" return self._run_script_async_internal(path=path, args=args, scheduled_in_secs=scheduled_in_secs, tag=tag) def run_script_by_hash_async( self, hash_: str, args: dict = None, scheduled_in_secs: int = None, tag: str = None, ) -> str: """Create a script job by hash and return its job id.""" return self._run_script_async_internal(hash_=hash_, args=args, scheduled_in_secs=scheduled_in_secs, tag=tag) def run_flow_async( self, path: str, args: dict = None, scheduled_in_secs: int = None, # can only be set to false if this the job will be fully await and not concurrent with any other job # as otherwise the child flow and its own child will store their state in the parent job which will # lead to incorrectness and failures do_not_track_in_parent: bool = True, tag: str = None, ) -> str: """Create a flow job and return its job id.""" args = args or {} params = {"scheduled_in_secs": scheduled_in_secs} if scheduled_in_secs else {} if tag: params["tag"] = tag if not do_not_track_in_parent: if os.environ.get("WM_JOB_ID"): params["parent_job"] = os.environ.get("WM_JOB_ID") if os.environ.get("WM_ROOT_FLOW_JOB_ID"): params["root_job"] = os.environ.get("WM_ROOT_FLOW_JOB_ID") if path: endpoint = f"/w/{self.workspace}/jobs/run/f/{path}" else: raise Exception("path must be provided") return self.post(endpoint, json=args, params=params).text def run_script( self, path: str = None, hash_: str = None, args: dict = None, timeout: dt.timedelta | int | float | None = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = False, tag: str = None, ) -> Any: """Run script synchronously and return its result. .. deprecated:: Use run_script_by_path or run_script_by_hash instead. """ logging.warning( "run_script is deprecated. Use run_script_by_path or run_script_by_hash instead.", ) assert not (path and hash_), "path and hash_ are mutually exclusive" return self._run_script_internal( path=path, hash_=hash_, args=args, timeout=timeout, verbose=verbose, cleanup=cleanup, assert_result_is_not_none=assert_result_is_not_none, tag=tag ) def _run_script_internal( self, path: str = None, hash_: str = None, args: dict = None, timeout: dt.timedelta | int | float | None = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = False, tag: str = None, ) -> Any: """Internal helper for running scripts synchronously.""" args = args or {} if verbose: if path: logger.info(f"running `{path}` synchronously with {args = }") elif hash_: logger.info(f"running script with hash `{hash_}` synchronously with {args = }") if isinstance(timeout, dt.timedelta): timeout = timeout.total_seconds() job_id = self._run_script_async_internal(path=path, hash_=hash_, args=args, tag=tag) return self.wait_job( job_id, timeout, verbose, cleanup, assert_result_is_not_none ) def run_script_by_path( self, path: str, args: dict = None, timeout: dt.timedelta | int | float | None = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = False, tag: str = None, ) -> Any: """Run script by path synchronously and return its result.""" return self._run_script_internal( path=path, args=args, timeout=timeout, verbose=verbose, cleanup=cleanup, assert_result_is_not_none=assert_result_is_not_none, tag=tag ) def run_script_by_hash( self, hash_: str, args: dict = None, timeout: dt.timedelta | int | float | None = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = False, tag: str = None, ) -> Any: """Run script by hash synchronously and return its result.""" return self._run_script_internal( hash_=hash_, args=args, timeout=timeout, verbose=verbose, cleanup=cleanup, assert_result_is_not_none=assert_result_is_not_none, tag=tag ) def run_inline_script_preview( self, content: str, language: str, args: dict = None, ) -> Any: """Run a script on the current worker without creating a job. On agent workers (no internal server), falls back to running a normal preview job and waiting for the result. """ if self.worker_has_internal_server(): endpoint = f"/w/{self.workspace}/jobs/run_inline/preview" else: endpoint = f"/w/{self.workspace}/jobs/run_wait_result/preview" body = { "content": content, "language": language, "args": args or {}, } return self.post(endpoint, json=body).json() def wait_job( self, job_id, timeout: dt.timedelta | int | float | None = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = False, ): """Wait for a job to complete and return its result. Args: job_id: ID of the job to wait for timeout: Maximum time to wait (seconds or timedelta) verbose: Enable verbose logging cleanup: Register cleanup handler to cancel job on exit assert_result_is_not_none: Raise exception if result is None Returns: Job result when completed Raises: TimeoutError: If timeout is reached Exception: If job fails """ def cancel_job(): logger.warning(f"cancelling job: {job_id}") self.post( f"/w/{self.workspace}/jobs_u/queue/cancel/{job_id}", json={"reason": "parent script cancelled"}, ).raise_for_status() if cleanup: atexit.register(cancel_job) start_time = time.time() if isinstance(timeout, dt.timedelta): timeout = timeout.total_seconds() while True: result_res = self.get( f"/w/{self.workspace}/jobs_u/completed/get_result_maybe/{job_id}", True ).json() started = result_res["started"] completed = result_res["completed"] success = result_res["success"] if not started and verbose: logger.info(f"job {job_id} has not started yet") if cleanup and completed: atexit.unregister(cancel_job) if completed: result = result_res["result"] if success: if result is None and assert_result_is_not_none: raise Exception("Result was none") return result else: error = result["error"] raise Exception(f"Job {job_id} was not successful: {str(error)}") if timeout and ((time.time() - start_time) > timeout): msg = "reached timeout" logger.warning(msg) self.post( f"/w/{self.workspace}/jobs_u/queue/cancel/{job_id}", json={"reason": msg}, ) raise TimeoutError(msg) if verbose: logger.info(f"sleeping 0.5 seconds for {job_id = }") time.sleep(0.5) def cancel_job(self, job_id: str, reason: str = None) -> str: """Cancel a specific job by ID. Args: job_id: UUID of the job to cancel reason: Optional reason for cancellation Returns: Response message from the cancel endpoint """ logger.info(f"cancelling job: {job_id}") payload = {"reason": reason or "cancelled via cancel_job method"} response = self.post( f"/w/{self.workspace}/jobs_u/queue/cancel/{job_id}", json=payload, ) return response.text def cancel_running(self) -> dict: """Cancel currently running executions of the same script.""" logger.info("canceling running executions of this script") jobs = self.get( f"/w/{self.workspace}/jobs/list", params={ "running": "true", "script_path_exact": self.path, }, ).json() current_job_id = os.environ.get("WM_JOB_ID") logger.debug(f"{current_job_id = }") job_ids = [j["id"] for j in jobs if j["id"] != current_job_id] if job_ids: logger.info(f"cancelling the following job ids: {job_ids}") else: logger.info("no previous executions to cancel") result = {} for id_ in job_ids: result[id_] = self.post( f"/w/{self.workspace}/jobs_u/queue/cancel/{id_}", json={"reason": "killed by `cancel_running` method"}, ) return result def get_job(self, job_id: str) -> dict: """Get job details by ID. Args: job_id: UUID of the job Returns: Job details dictionary """ return self.get(f"/w/{self.workspace}/jobs_u/get/{job_id}").json() def get_root_job_id(self, job_id: str | None = None) -> dict: """Get the root job ID for a flow hierarchy. Args: job_id: Job ID (defaults to current WM_JOB_ID) Returns: Root job ID """ job_id = job_id or os.environ.get("WM_JOB_ID") return self.get(f"/w/{self.workspace}/jobs_u/get_root_job_id/{job_id}").json() def get_id_token(self, audience: str, expires_in: int | None = None) -> str: """Get an OIDC JWT token for authentication to external services. Args: audience: Token audience (e.g., "vault", "aws") expires_in: Optional expiration time in seconds Returns: JWT token string """ params = {} if expires_in is not None: params["expires_in"] = expires_in return self.post(f"/w/{self.workspace}/oidc/token/{audience}", params=params).text def get_job_status(self, job_id: str) -> JobStatus: """Get the status of a job. Args: job_id: UUID of the job Returns: Job status: "RUNNING", "WAITING", or "COMPLETED" """ job = self.get_job(job_id) job_type = job.get("type", "") assert job_type, f"{job} is not a valid job" if job_type.lower() == "completedjob": return "COMPLETED" if job.get("running"): return "RUNNING" return "WAITING" def get_result( self, job_id: str, assert_result_is_not_none: bool = True, ) -> Any: """Get the result of a completed job. Args: job_id: UUID of the completed job assert_result_is_not_none: Raise exception if result is None Returns: Job result """ result = self.get(f"/w/{self.workspace}/jobs_u/completed/get_result/{job_id}") result_text = result.text if assert_result_is_not_none and result_text is None: raise Exception(f"result is None for {job_id = }") try: return result.json() except JSONDecodeError: return result_text def get_variable(self, path: str) -> str: """Get a variable value by path. Args: path: Variable path in Windmill Returns: Variable value as string """ path = parse_variable_syntax(path) or path if self.mocked_api is not None: variables = self.mocked_api["variables"] try: result = variables[path] return result except KeyError: logger.info( f"MockedAPI present, but variable not found at {path}, falling back to real API" ) return self.get(f"/w/{self.workspace}/variables/get_value/{path}").json() def set_variable(self, path: str, value: str, is_secret: bool = False) -> None: """Set a variable value by path, creating it if it doesn't exist. Args: path: Variable path in Windmill value: Variable value to set is_secret: Whether the variable should be secret (default: False) """ path = parse_variable_syntax(path) or path if self.mocked_api is not None: self.mocked_api["variables"][path] = value return # check if variable exists r = self.get( f"/w/{self.workspace}/variables/get/{path}", raise_for_status=False ) if r.status_code == 404: # create variable self.post( f"/w/{self.workspace}/variables/create", json={ "path": path, "value": value, "is_secret": is_secret, "description": "", }, ) else: # update variable self.post( f"/w/{self.workspace}/variables/update/{path}", json={"value": value}, ) def get_resource( self, path: str, none_if_undefined: bool = False, interpolated: bool = True ) -> dict | None: """Get a resource value by path. Args: path: Resource path in Windmill none_if_undefined: Return None instead of raising if not found interpolated: if variables and resources are fully unrolled Returns: Resource value dictionary or None """ path = parse_resource_syntax(path) or path if self.mocked_api is not None: resources = self.mocked_api["resources"] try: result = resources[path] return result except KeyError: # NOTE: should mocked_api respect `none_if_undefined`? if none_if_undefined: logger.info( f"resource not found at ${path}, but none_if_undefined is True, so returning None" ) return None logger.info( f"MockedAPI present, but resource not found at ${path}, falling back to real API" ) try: if interpolated: return self.get( f"/w/{self.workspace}/resources/get_value_interpolated/{path}" ).json() else: return self.get( f"/w/{self.workspace}/resources/get_value/{path}" ).json() except Exception as e: if none_if_undefined: return None logger.error(e) raise e def set_resource( self, value: Any, path: str, resource_type: str, ): """Set a resource value by path, creating it if it doesn't exist. Args: value: Resource value to set path: Resource path in Windmill resource_type: Resource type for creation """ path = parse_resource_syntax(path) or path if self.mocked_api is not None: self.mocked_api["resources"][path] = value return # check if resource exists r = self.get( f"/w/{self.workspace}/resources/get/{path}", raise_for_status=False ) if r.status_code == 404: # create resource self.post( f"/w/{self.workspace}/resources/create", json={ "path": path, "value": value, "resource_type": resource_type, }, ) else: # update resource self.post( f"/w/{self.workspace}/resources/update_value/{path}", json={"value": value}, ) def list_resources( self, resource_type: str = None, page: int = None, per_page: int = None, ) -> list[dict]: """List resources from Windmill workspace. Args: resource_type: Optional resource type to filter by (e.g., "postgresql", "mysql", "s3") page: Optional page number for pagination per_page: Optional number of results per page Returns: List of resource dictionaries """ params = {} if resource_type is not None: params["resource_type"] = resource_type if page is not None: params["page"] = page if per_page is not None: params["per_page"] = per_page return self.get( f"/w/{self.workspace}/resources/list", params=params if params else None, ).json() def set_state(self, value: Any, path: str | None = None) -> None: """Set the workflow state. Args: value: State value to set path: Optional state resource path override. """ self.set_resource(value, path=path or self.state_path, resource_type="state") def get_state(self, path: str | None = None) -> Any: """Get the workflow state. Args: path: Optional state resource path override. Returns: State value or None if not set """ return self.get_resource(path=path or self.state_path, none_if_undefined=True, interpolated=True) def set_progress(self, value: int, job_id: Optional[str] = None): """Set job progress percentage (0-99). Args: value: Progress percentage job_id: Job ID (defaults to current WM_JOB_ID) """ workspace = get_workspace() flow_id = os.environ.get("WM_FLOW_JOB_ID") job_id = job_id or os.environ.get("WM_JOB_ID") if job_id != None: job = self.get_job(job_id) flow_id = job.get("parent_job") self.post( f"/w/{workspace}/job_metrics/set_progress/{job_id}", json={ "percent": value, "flow_job_id": flow_id or None, }, ) def get_progress(self, job_id: Optional[str] = None) -> Any: """Get job progress percentage. Args: job_id: Job ID (defaults to current WM_JOB_ID) Returns: Progress value (0-100) or None if not set """ workspace = get_workspace() job_id = job_id or os.environ.get("WM_JOB_ID") r = self.get( f"/w/{workspace}/job_metrics/get_progress/{job_id}", ) if r.status_code == 404: print(f"Job {job_id} does not exist") return None else: return r.json() def set_flow_user_state(self, key: str, value: Any) -> None: """Set the user state of a flow at a given key""" flow_id = self.get_root_job_id() r = self.post( f"/w/{self.workspace}/jobs/flow/user_states/{flow_id}/{key}", json=value, raise_for_status=False, ) if r.status_code == 404: print(f"Job {flow_id} does not exist or is not a flow") def get_flow_user_state(self, key: str) -> Any: """Get the user state of a flow at a given key""" flow_id = self.get_root_job_id() r = self.get( f"/w/{self.workspace}/jobs/flow/user_states/{flow_id}/{key}", raise_for_status=False, ) if r.status_code == 404: print(f"Job {flow_id} does not exist or is not a flow") return None else: return r.json() @property def version(self): """Get the Windmill server version. Returns: Version string """ return self.get("version").text def get_duckdb_connection_settings( self, s3_resource_path: str = "", ) -> DuckDbConnectionSettings | None: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection from DuckDB """ s3_resource_path = parse_resource_syntax(s3_resource_path) or s3_resource_path try: raw_obj = self.post( f"/w/{self.workspace}/job_helpers/v2/duckdb_connection_settings", json={} if s3_resource_path == "" else {"s3_resource_path": s3_resource_path}, ).json() return DuckDbConnectionSettings(raw_obj) except JSONDecodeError as e: raise Exception( "Could not generate DuckDB S3 connection settings from the provided resource" ) from e def get_polars_connection_settings( self, s3_resource_path: str = "", ) -> PolarsConnectionSettings: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection from Polars """ s3_resource_path = parse_resource_syntax(s3_resource_path) or s3_resource_path try: raw_obj = self.post( f"/w/{self.workspace}/job_helpers/v2/polars_connection_settings", json={} if s3_resource_path == "" else {"s3_resource_path": s3_resource_path}, ).json() return PolarsConnectionSettings(raw_obj) except JSONDecodeError as e: raise Exception( "Could not generate Polars S3 connection settings from the provided resource" ) from e def get_boto3_connection_settings( self, s3_resource_path: str = "", ) -> Boto3ConnectionSettings: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection using boto3 """ s3_resource_path = parse_resource_syntax(s3_resource_path) or s3_resource_path try: s3_resource = self.post( f"/w/{self.workspace}/job_helpers/v2/s3_resource_info", json={} if s3_resource_path == "" else {"s3_resource_path": s3_resource_path}, ).json() return self.__boto3_connection_settings(s3_resource) except JSONDecodeError as e: raise Exception( "Could not generate Boto3 S3 connection settings from the provided resource" ) from e def load_s3_file(self, s3object: S3Object | str, s3_resource_path: str | None) -> bytes: """ Load a file from the workspace s3 bucket and returns its content as bytes. '''python from wmill import S3Object s3_obj = S3Object(s3="/path/to/my_file.txt") my_obj_content = client.load_s3_file(s3_obj) file_content = my_obj_content.decode("utf-8") ''' """ s3object = parse_s3_object(s3object) with self.load_s3_file_reader(s3object, s3_resource_path) as file_reader: return file_reader.read() def load_s3_file_reader( self, s3object: S3Object | str, s3_resource_path: str | None ) -> BufferedReader: """ Load a file from the workspace s3 bucket and returns the bytes stream. '''python from wmill import S3Object s3_obj = S3Object(s3="/path/to/my_file.txt") with wmill.load_s3_file_reader(s3object, s3_resource_path) as file_reader: print(file_reader.read()) ''' """ s3object = parse_s3_object(s3object) reader = S3BufferedReader( f"{self.workspace}", self.client, s3object["s3"], s3_resource_path, s3object["storage"] if "storage" in s3object else None, ) return reader def write_s3_file( self, s3object: S3Object | str | None, file_content: BufferedReader | bytes, s3_resource_path: str | None, content_type: str | None = None, content_disposition: str | None = None, ) -> S3Object: """ Write a file to the workspace S3 bucket '''python from wmill import S3Object s3_obj = S3Object(s3="/path/to/my_file.txt") # for an in memory bytes array: file_content = b'Hello Windmill!' client.write_s3_file(s3_obj, file_content) # for a file: with open("my_file.txt", "rb") as my_file: client.write_s3_file(s3_obj, my_file) ''' """ s3object = parse_s3_object(s3object) # httpx accepts either bytes or "a bytes generator" as content. If it's a BufferedReader, we need to convert it to a generator if isinstance(file_content, BufferedReader): content_payload = bytes_generator(file_content) elif isinstance(file_content, bytes): content_payload = file_content else: raise Exception("Type of file_content not supported") query_params = {} if s3object is not None and s3object["s3"] != "": query_params["file_key"] = s3object["s3"] if s3_resource_path is not None and s3_resource_path != "": query_params["s3_resource_path"] = s3_resource_path if ( s3object is not None and "storage" in s3object and s3object["storage"] is not None ): query_params["storage"] = s3object["storage"] if content_type is not None: query_params["content_type"] = content_type if content_disposition is not None: query_params["content_disposition"] = content_disposition try: # need a vanilla client b/c content-type is not application/json here response = httpx.post( f"{self.base_url}/w/{self.workspace}/job_helpers/upload_s3_file", headers={ "Authorization": f"Bearer {self.token}", "Content-Type": "application/octet-stream", }, params=query_params, content=content_payload, verify=self.verify, timeout=None, ).json() except Exception as e: raise Exception("Could not write file to S3") from e return S3Object(s3=response["file_key"], storage=s3object.get("storage") if s3object else None) def delete_s3_object( self, s3object: S3Object | str, s3_resource_path: str | None = None, ) -> None: """ Permanently delete a file from the workspace S3 bucket. '''python from wmill import S3Object s3_obj = S3Object(s3="/path/to/my_file.txt") client.delete_s3_object(s3_obj) ''' """ s3object = parse_s3_object(s3object) query_params: Dict[str, Any] = {"file_key": s3object["s3"]} if s3_resource_path is not None and s3_resource_path != "": query_params["s3_resource_path"] = s3_resource_path if "storage" in s3object and s3object["storage"] is not None: query_params["storage"] = s3object["storage"] try: resp = self.client.delete( f"/w/{self.workspace}/job_helpers/delete_s3_file", params=query_params, ) resp.raise_for_status() except httpx.HTTPStatusError as err: error = f"{err.request.url}: {err.response.status_code}, {err.response.text}" logger.error(error) raise Exception(error) except Exception as e: raise Exception("Could not delete file from S3") from e def sign_s3_objects( self, s3_objects: list[S3Object | str], expiry_secs: int | None = None ) -> list[S3Object]: """Sign S3 objects for use by anonymous users in public apps. Args: s3_objects: List of S3 objects to sign expiry_secs: How long the signature stays valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) Returns: List of signed S3 objects """ return self.post( f"/w/{self.workspace}/apps/sign_s3_objects", json=_sign_s3_objects_body(list(map(parse_s3_object, s3_objects)), expiry_secs), ).json() def sign_s3_object(self, s3_object: S3Object | str, expiry_secs: int | None = None) -> S3Object: """Sign a single S3 object for use by anonymous users in public apps. Args: s3_object: S3 object to sign expiry_secs: How long the signature stays valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) Returns: Signed S3 object """ return self.post( f"/w/{self.workspace}/apps/sign_s3_objects", json=_sign_s3_objects_body([s3_object], expiry_secs), ).json()[0] def get_presigned_s3_public_urls( self, s3_objects: list[S3Object | str], base_url: str | None = None, expiry_secs: int | None = None, ) -> list[str]: """ Generate presigned public URLs for an array of S3 objects. If an S3 object is not signed yet, it will be signed first. Args: s3_objects: List of S3 objects to sign base_url: Optional base URL for the presigned URLs (defaults to WM_BASE_URL) expiry_secs: How long the signatures stay valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) Returns: List of signed public URLs Example: >>> s3_objs = [S3Object(s3="/path/to/file1.txt"), S3Object(s3="/path/to/file2.txt")] >>> urls = client.get_presigned_s3_public_urls(s3_objs) """ base_url = base_url or self._get_public_base_url() s3_objs = [parse_s3_object(s3_obj) for s3_obj in s3_objects] # Sign all S3 objects that need to be signed in one go s3_objs_to_sign: list[tuple[S3Object, int]] = [ (s3_obj, index) for index, s3_obj in enumerate(s3_objs) if s3_obj.get("presigned") is None ] if s3_objs_to_sign: signed_s3_objs = self.sign_s3_objects( [s3_obj for s3_obj, _ in s3_objs_to_sign], expiry_secs ) for i, (_, original_index) in enumerate(s3_objs_to_sign): s3_objs[original_index] = parse_s3_object(signed_s3_objs[i]) signed_urls: list[str] = [] for s3_obj in s3_objs: s3 = s3_obj.get("s3", "") presigned = s3_obj.get("presigned", "") storage = s3_obj.get("storage", "_default_") signed_url = f"{base_url}/api/w/{self.workspace}/s3_proxy/{storage}/{s3}?{presigned}" signed_urls.append(signed_url) return signed_urls def get_presigned_s3_public_url( self, s3_object: S3Object | str, base_url: str | None = None, expiry_secs: int | None = None, ) -> str: """ Generate a presigned public URL for an S3 object. If the S3 object is not signed yet, it will be signed first. Args: s3_object: S3 object to sign base_url: Optional base URL for the presigned URL (defaults to WM_BASE_URL) expiry_secs: How long the signature stays valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) Returns: Signed public URL Example: >>> s3_obj = S3Object(s3="/path/to/file.txt") >>> url = client.get_presigned_s3_public_url(s3_obj) """ urls = self.get_presigned_s3_public_urls([s3_object], base_url, expiry_secs) return urls[0] def _get_public_base_url(self) -> str: """Get the public base URL from environment or default to localhost""" return os.environ.get("WM_BASE_URL", "http://localhost:3000") def __boto3_connection_settings(self, s3_resource) -> Boto3ConnectionSettings: endpoint_url_prefix = "https://" if s3_resource["useSSL"] else "http://" endpoint = s3_resource["endPoint"] port = s3_resource.get("port") if port: endpoint_url = "{}{}:{}".format(endpoint_url_prefix, endpoint, port) else: endpoint_url = "{}{}".format(endpoint_url_prefix, endpoint) settings = { "endpoint_url": endpoint_url, "region_name": s3_resource["region"], "use_ssl": s3_resource["useSSL"], "aws_access_key_id": s3_resource["accessKey"], "aws_secret_access_key": s3_resource["secretKey"], # no need for path_style here as boto3 is clever enough to determine which one to use } # Include session token for OIDC/STS temporary credentials if s3_resource.get("token"): settings["aws_session_token"] = s3_resource["token"] return Boto3ConnectionSettings(settings) def whoami(self) -> dict: """Get the current user information. Returns: User details dictionary """ return self.get("/users/whoami").json() @property def user(self) -> dict: """Get the current user information (alias for whoami). Returns: User details dictionary """ return self.whoami() @property def state_path(self) -> str: """Get the state resource path from environment. Returns: State path string """ state_path = os.environ.get( "WM_STATE_PATH_NEW", os.environ.get("WM_STATE_PATH") ) if state_path is None: raise Exception("State path not found") return state_path @property def state(self) -> Any: """Get the workflow state. Returns: State value or None if not set """ return self.get_resource(path=self.state_path, none_if_undefined=True, interpolated=True) @state.setter def state(self, value: Any) -> None: """Set the workflow state.""" self.set_state(value) @staticmethod def set_shared_state_pickle(value: Any, path: str = "state.pickle") -> None: """ Set the state in the shared folder using pickle """ import pickle with open(f"/shared/{path}", "wb") as handle: pickle.dump(value, handle, protocol=pickle.HIGHEST_PROTOCOL) @staticmethod def get_shared_state_pickle(path: str = "state.pickle") -> Any: """ Get the state in the shared folder using pickle """ import pickle with open(f"/shared/{path}", "rb") as handle: return pickle.load(handle) @staticmethod def set_shared_state(value: Any, path: str = "state.json") -> None: """ Set the state in the shared folder using pickle """ import json with open(f"/shared/{path}", "w", encoding="utf-8") as f: json.dump(value, f, ensure_ascii=False, indent=4) @staticmethod def get_shared_state(path: str = "state.json") -> None: """ Get the state in the shared folder using pickle """ import json with open(f"/shared/{path}", "r", encoding="utf-8") as f: return json.load(f) def get_resume_urls(self, approver: str = None, flow_level: bool = None) -> dict: """Get URLs needed for resuming a flow after suspension. Args: approver: Optional approver name flow_level: If True, generate resume URLs for the parent flow instead of the specific step. This allows pre-approvals that can be consumed by any later suspend step in the same flow. Returns: Dictionary with approvalPage, resume, and cancel URLs """ nonce = random.randint(0, 1000000000) job_id = os.environ.get("WM_JOB_ID") or "NO_ID" params = {"approver": approver} if flow_level is not None: params["flow_level"] = flow_level return self.get( f"/w/{self.workspace}/jobs/resume_urls/{job_id}/{nonce}", params=params, ).json() def get_approval_urls(self, step_key: str = "approval", approver: str = None) -> dict: """Get the resume URLs bound to one ``wait_for_approval`` step of this workflow. Args: step_key: Checkpoint key of the approval step, as passed to ``wait_for_approval(key=...)`` approver: Optional approver name Returns: Dictionary with approvalPage, resume, and cancel URLs """ from urllib.parse import quote _assert_usable_step_key(step_key, "get_approval_urls step_key") job_id = os.environ.get("WM_JOB_ID") or "NO_ID" # Omit rather than send `approver=`: an empty value is echoed into the # returned URLs and recorded as the approver instead of "anonymous". params = {"approver": approver} if approver is not None else {} return self.get( f"/w/{self.workspace}/jobs/wac_approval_urls/{job_id}/{quote(step_key, safe='')}", params=params, ).json() def request_interactive_slack_approval( self, slack_resource_path: str, channel_id: str, message: str = None, approver: str = None, default_args_json: dict = None, dynamic_enums_json: dict = None, ) -> None: """ Sends an interactive approval request via Slack, allowing optional customization of the message, approver, and form fields. **[Enterprise Edition Only]** To include form fields in the Slack approval request, use the "Advanced -> Suspend -> Form" functionality. Learn more at: https://www.windmill.dev/docs/flows/flow_approval#form :param slack_resource_path: The path to the Slack resource in Windmill. :type slack_resource_path: str :param channel_id: The Slack channel ID where the approval request will be sent. :type channel_id: str :param message: Optional custom message to include in the Slack approval request. :type message: str, optional :param approver: Optional user ID or name of the approver for the request. :type approver: str, optional :param default_args_json: Optional dictionary defining or overriding the default arguments for form fields. :type default_args_json: dict, optional :param dynamic_enums_json: Optional dictionary overriding the enum default values of enum form fields. :type dynamic_enums_json: dict, optional :raises Exception: If the function is not called within a flow or flow preview. :raises Exception: If the required flow job or flow step environment variables are not set. :return: None **Usage Example:** >>> client.request_interactive_slack_approval( ... slack_resource_path="/u/alex/my_slack_resource", ... channel_id="admins-slack-channel", ... message="Please approve this request", ... approver="approver123", ... default_args_json={"key1": "value1", "key2": 42}, ... dynamic_enums_json={"foo": ["choice1", "choice2"], "bar": ["optionA", "optionB"]}, ... ) **Notes:** - This function must be executed within a Windmill flow or flow preview. - The function checks for required environment variables (`WM_FLOW_JOB_ID`, `WM_FLOW_STEP_ID`) to ensure it is run in the appropriate context. """ workspace = self.workspace flow_job_id = os.environ.get("WM_FLOW_JOB_ID") if not flow_job_id: raise Exception( "You can't use 'request_interactive_slack_approval' function in a standalone script or flow step preview. Please use it in a flow or a flow preview." ) # Only include non-empty parameters params = {} if message: params["message"] = message if approver: params["approver"] = approver if slack_resource_path: params["slack_resource_path"] = slack_resource_path if channel_id: params["channel_id"] = channel_id if os.environ.get("WM_FLOW_STEP_ID"): params["flow_step_id"] = os.environ.get("WM_FLOW_STEP_ID") if default_args_json: params["default_args_json"] = json.dumps(default_args_json) if dynamic_enums_json: params["dynamic_enums_json"] = json.dumps(dynamic_enums_json) self.get( f"/w/{workspace}/jobs/slack_approval/{os.environ.get('WM_JOB_ID', 'NO_JOB_ID')}", params=params, ) def username_to_email(self, username: str) -> str: """ Get email from workspace username .. deprecated:: Read the contextual variables instead: `os.environ.get("WM_END_USER_EMAIL") or os.environ.get("WM_EMAIL")`. WM_END_USER_EMAIL is the email of whoever triggered the run when it came from an app, so the fallback yields the app viewer inside an app and the executing user everywhere else - without an extra API call, and unlike this method it also resolves viewers who are not workspace members. An app viewed anonymously has no identity to report: the variable is then empty and the fallback yields the app publisher. """ return self.get(f"/w/{self.workspace}/users/username_to_email/{username}").text def send_teams_message( self, conversation_id: str, text: str, success: bool = True, card_block: dict = None, ): """ Send a message to a Microsoft Teams conversation with conversation_id, where success is used to style the message """ return self.post( f"/teams/activities", json={ "conversation_id": conversation_id, "text": text, "success": success, "card_block": card_block, }, ) def datatable(self, name: str = "main"): """Get a DataTable client for SQL queries. Args: name: Database name (default: "main") Returns: DataTableClient instance """ return DataTableClient(self, name) def ducklake(self, name: str = "main"): """Get a DuckLake client for DuckDB queries. Args: name: Database name (default: "main") Returns: DucklakeClient instance """ return DucklakeClient(self, name) def init_global_client(f): @functools.wraps(f) def wrapper(*args, **kwargs): global _client if _client is None: _client = Windmill() return f(*args, **kwargs) return wrapper def deprecate(in_favor_of: str): def decorator(f): @functools.wraps(f) def wrapper(*args, **kwargs): warnings.warn( ( f"The '{f.__name__}' method is deprecated and may be removed in the future. " f"Consider {in_favor_of}" ), DeprecationWarning, ) return f(*args, **kwargs) return wrapper return decorator @init_global_client def get_workspace() -> str: """Get the current workspace ID. Returns: Workspace ID string """ return _client.workspace @init_global_client def get_root_job_id(job_id: str | None = None) -> str: """Get the root job ID for a flow hierarchy. Args: job_id: Job ID (defaults to current WM_JOB_ID) Returns: Root job ID """ return _client.get_root_job_id(job_id) @init_global_client @deprecate("Windmill().version") def get_version() -> str: return _client.version @init_global_client def run_script_async( hash_or_path: str, args: Dict[str, Any] = None, scheduled_in_secs: int = None, tag: str = None, ) -> str: """Create a script job and return its job ID. Args: hash_or_path: Script hash or path (determined by presence of '/') args: Script arguments scheduled_in_secs: Delay before execution in seconds tag: Override the worker tag the job runs on Returns: Job ID string """ is_path = "/" in hash_or_path hash_ = None if is_path else hash_or_path path = hash_or_path if is_path else None return _client.run_script_async( hash_=hash_, path=path, args=args, scheduled_in_secs=scheduled_in_secs, tag=tag, ) @init_global_client def run_flow_async( path: str, args: Dict[str, Any] = None, scheduled_in_secs: int = None, # can only be set to false if this the job will be fully await and not concurrent with any other job # as otherwise the child flow and its own child will store their state in the parent job which will # lead to incorrectness and failures do_not_track_in_parent: bool = True, tag: str = None, ) -> str: """Create a flow job and return its job ID. Args: path: Flow path args: Flow arguments scheduled_in_secs: Delay before execution in seconds do_not_track_in_parent: Whether to track in parent job (default: True) tag: Override the worker tag the job runs on Returns: Job ID string """ return _client.run_flow_async( path=path, args=args, scheduled_in_secs=scheduled_in_secs, do_not_track_in_parent=do_not_track_in_parent, tag=tag, ) @init_global_client def run_script_sync( hash: str, args: Dict[str, Any] = None, verbose: bool = False, assert_result_is_not_none: bool = True, cleanup: bool = True, timeout: dt.timedelta = None, tag: str = None, ) -> Any: """Run a script synchronously by hash and return its result. Args: hash: Script hash args: Script arguments verbose: Enable verbose logging assert_result_is_not_none: Raise exception if result is None cleanup: Register cleanup handler to cancel job on exit timeout: Maximum time to wait tag: Override the worker tag the job runs on Returns: Script result """ return _client.run_script( hash_=hash, args=args, verbose=verbose, assert_result_is_not_none=assert_result_is_not_none, cleanup=cleanup, timeout=timeout, tag=tag, ) @init_global_client def run_script_by_path_async( path: str, args: Dict[str, Any] = None, scheduled_in_secs: Union[None, int] = None, tag: str = None, ) -> str: """Create a script job by path and return its job ID. Args: path: Script path args: Script arguments scheduled_in_secs: Delay before execution in seconds tag: Override the worker tag the job runs on Returns: Job ID string """ return _client.run_script_by_path_async( path=path, args=args, scheduled_in_secs=scheduled_in_secs, tag=tag, ) @init_global_client def run_script_by_hash_async( hash_: str, args: Dict[str, Any] = None, scheduled_in_secs: Union[None, int] = None, tag: str = None, ) -> str: """Create a script job by hash and return its job ID. Args: hash_: Script hash args: Script arguments scheduled_in_secs: Delay before execution in seconds tag: Override the worker tag the job runs on Returns: Job ID string """ return _client.run_script_by_hash_async( hash_=hash_, args=args, scheduled_in_secs=scheduled_in_secs, tag=tag, ) @init_global_client def run_script_by_path_sync( path: str, args: Dict[str, Any] = None, verbose: bool = False, assert_result_is_not_none: bool = True, cleanup: bool = True, timeout: dt.timedelta = None, tag: str = None, ) -> Any: """Run a script synchronously by path and return its result. Args: path: Script path args: Script arguments verbose: Enable verbose logging assert_result_is_not_none: Raise exception if result is None cleanup: Register cleanup handler to cancel job on exit timeout: Maximum time to wait tag: Override the worker tag the job runs on Returns: Script result """ return _client.run_script( path=path, args=args, verbose=verbose, assert_result_is_not_none=assert_result_is_not_none, cleanup=cleanup, timeout=timeout, tag=tag, ) @init_global_client def get_id_token(audience: str) -> str: """ Get a JWT token for the given audience for OIDC purposes to login into third parties like AWS, Vault, GCP, etc. """ return _client.get_id_token(audience) @init_global_client def get_job_status(job_id: str) -> JobStatus: """Get the status of a job. Args: job_id: UUID of the job Returns: Job status: "RUNNING", "WAITING", or "COMPLETED" """ return _client.get_job_status(job_id) @init_global_client def get_job(job_id: str) -> dict: """Get full job details by ID. Args: job_id: UUID of the job Returns: Job details dictionary """ return _client.get_job(job_id=job_id) @init_global_client def get_result(job_id: str, assert_result_is_not_none=True) -> Dict[str, Any]: """Get the result of a completed job. Args: job_id: UUID of the completed job assert_result_is_not_none: Raise exception if result is None Returns: Job result """ return _client.get_result( job_id=job_id, assert_result_is_not_none=assert_result_is_not_none ) @init_global_client def duckdb_connection_settings(s3_resource_path: str = "") -> DuckDbConnectionSettings: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection from DuckDB """ return _client.get_duckdb_connection_settings(s3_resource_path) @init_global_client def polars_connection_settings(s3_resource_path: str = "") -> PolarsConnectionSettings: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection from Polars """ return _client.get_polars_connection_settings(s3_resource_path) @init_global_client def boto3_connection_settings(s3_resource_path: str = "") -> Boto3ConnectionSettings: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection using boto3 """ return _client.get_boto3_connection_settings(s3_resource_path) @init_global_client def load_s3_file(s3object: S3Object | str, s3_resource_path: str | None = None) -> bytes: """ Load the entire content of a file stored in S3 as bytes """ return _client.load_s3_file( s3object, s3_resource_path if s3_resource_path != "" else None ) @init_global_client def load_s3_file_reader( s3object: S3Object | str, s3_resource_path: str | None = None ) -> BufferedReader: """ Load the content of a file stored in S3 """ return _client.load_s3_file_reader( s3object, s3_resource_path if s3_resource_path != "" else None ) @init_global_client def write_s3_file( s3object: S3Object | str | None, file_content: BufferedReader | bytes, s3_resource_path: str | None = None, content_type: str | None = None, content_disposition: str | None = None, ) -> S3Object: """ Upload a file to S3 Content type will be automatically guessed from path extension if left empty See MDN for content_disposition: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Disposition and content_type: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Type """ return _client.write_s3_file( s3object, file_content, s3_resource_path if s3_resource_path != "" else None, content_type, content_disposition, ) @init_global_client def delete_s3_object( s3object: S3Object | str, s3_resource_path: str | None = None, ) -> None: """ Permanently delete a file from the workspace S3 bucket. """ return _client.delete_s3_object( s3object, s3_resource_path if s3_resource_path != "" else None, ) @init_global_client def sign_s3_objects(s3_objects: list[S3Object | str], expiry_secs: int | None = None) -> list[S3Object]: """ Sign S3 objects to be used by anonymous users in public apps Returns a list of signed s3 tokens Args: s3_objects: List of S3 objects to sign expiry_secs: How long the signatures stay valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) """ return _client.sign_s3_objects(s3_objects, expiry_secs) @init_global_client def sign_s3_object(s3_object: S3Object| str, expiry_secs: int | None = None) -> S3Object: """ Sign S3 object to be used by anonymous users in public apps Returns a signed s3 object Args: s3_object: S3 object to sign expiry_secs: How long the signature stays valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) """ return _client.sign_s3_object(s3_object, expiry_secs) @init_global_client def get_presigned_s3_public_urls( s3_objects: list[S3Object | str], base_url: str | None = None, expiry_secs: int | None = None, ) -> list[str]: """ Generate presigned public URLs for an array of S3 objects. If an S3 object is not signed yet, it will be signed first. Args: s3_objects: List of S3 objects to sign base_url: Optional base URL for the presigned URLs (defaults to WM_BASE_URL) expiry_secs: How long the signatures stay valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) Returns: List of signed public URLs Example: >>> import wmill >>> from wmill import S3Object >>> s3_objs = [S3Object(s3="/path/to/file1.txt"), S3Object(s3="/path/to/file2.txt")] >>> urls = wmill.get_presigned_s3_public_urls(s3_objs) """ return _client.get_presigned_s3_public_urls(s3_objects, base_url, expiry_secs) @init_global_client def get_presigned_s3_public_url( s3_object: S3Object | str, base_url: str | None = None, expiry_secs: int | None = None, ) -> str: """ Generate a presigned public URL for an S3 object. If the S3 object is not signed yet, it will be signed first. Args: s3_object: S3 object to sign base_url: Optional base URL for the presigned URL (defaults to WM_BASE_URL) expiry_secs: How long the signature stays valid, in seconds (defaults to 43200 = 12h, clamped to [60, 604800]) Returns: Signed public URL Example: >>> import wmill >>> from wmill import S3Object >>> s3_obj = S3Object(s3="/path/to/file.txt") >>> url = wmill.get_presigned_s3_public_url(s3_obj) """ return _client.get_presigned_s3_public_url(s3_object, base_url, expiry_secs) @init_global_client def whoami() -> dict: """ Returns the current user """ return _client.user @init_global_client def get_state(path: str | None = None) -> Any: """ Get the state """ return _client.get_state(path=path) @init_global_client def get_resource( path: str, none_if_undefined: bool = False, interpolated: bool = True ) -> dict | None: """Get resource from Windmill""" return _client.get_resource(path, none_if_undefined, interpolated) @init_global_client def set_resource(path: str, value: Any, resource_type: str = "any") -> None: """ Set the resource at a given path as a string, creating it if it does not exist """ return _client.set_resource(value=value, path=path, resource_type=resource_type) @init_global_client def list_resources( resource_type: str = None, page: int = None, per_page: int = None, ) -> list[dict]: """List resources from Windmill workspace. Args: resource_type: Optional resource type to filter by (e.g., "postgresql", "mysql", "s3") page: Optional page number for pagination per_page: Optional number of results per page Returns: List of resource dictionaries Example: >>> # Get all resources >>> all_resources = wmill.list_resources() >>> # Get only PostgreSQL resources >>> pg_resources = wmill.list_resources(resource_type="postgresql") """ return _client.list_resources( resource_type=resource_type, page=page, per_page=per_page, ) @init_global_client def set_state(value: Any, path: str | None = None) -> None: """ Set the state """ return _client.set_state(value, path=path) @init_global_client def set_progress(value: int, job_id: Optional[str] = None) -> None: """ Set the progress """ return _client.set_progress(value, job_id) @init_global_client def get_progress(job_id: Optional[str] = None) -> Any: """ Get the progress """ return _client.get_progress(job_id) def set_shared_state_pickle(value: Any, path="state.pickle") -> None: """ Set the state in the shared folder using pickle """ return Windmill.set_shared_state_pickle(value=value, path=path) @deprecate("Windmill.get_shared_state_pickle(...)") def get_shared_state_pickle(path="state.pickle") -> Any: """ Get the state in the shared folder using pickle """ return Windmill.get_shared_state_pickle(path=path) def set_shared_state(value: Any, path="state.json") -> None: """ Set the state in the shared folder using pickle """ return Windmill.set_shared_state(value=value, path=path) def get_shared_state(path="state.json") -> None: """ Get the state in the shared folder using pickle """ return Windmill.get_shared_state(path=path) @init_global_client def get_variable(path: str) -> str: """ Returns the variable at a given path as a string """ return _client.get_variable(path) @init_global_client def set_variable(path: str, value: str, is_secret: bool = False) -> None: """ Set the variable at a given path as a string, creating it if it does not exist """ return _client.set_variable(path, value, is_secret) @init_global_client def get_flow_user_state(key: str) -> Any: """ Get the user state of a flow at a given key """ return _client.get_flow_user_state(key) @init_global_client def set_flow_user_state(key: str, value: Any) -> None: """ Set the user state of a flow at a given key """ return _client.set_flow_user_state(key, value) @init_global_client def get_state_path() -> str: """Get the state resource path from environment. Returns: State path string """ return _client.state_path @init_global_client def get_resume_urls(approver: str = None, flow_level: bool = None) -> dict: """Get URLs needed for resuming a flow after suspension. Args: approver: Optional approver name flow_level: If True, generate resume URLs for the parent flow instead of the specific step. This allows pre-approvals that can be consumed by any later suspend step in the same flow. Returns: Dictionary with approvalPage, resume, and cancel URLs """ return _client.get_resume_urls(approver, flow_level) @init_global_client def get_approval_urls(step_key: str = "approval", approver: str = None) -> dict: """Get the resume/cancel/approval-page URLs bound to one ``wait_for_approval`` step. Unlike :func:`get_resume_urls`, which signs a random nonce, these address the very ``resume_job`` record the step's built-in approval buttons use, so they are stable across replays and safe to embed in a custom notification. Args: step_key: Checkpoint key of the approval step, as passed to ``wait_for_approval(key=...)``. Keys must be unique within a workflow; reusing one raises rather than silently renaming it. The URL only resumes while that step is awaiting approval; used at any other moment it is rejected rather than banking a row a different approval would consume. Send it ahead of time — approvers just cannot act before the workflow reaches the step. ``resume`` and ``cancel`` are step-bound; ``approvalPage`` is not — it opens the job's approval page, which acts on whichever approval is pending when it is used. approver: Optional approver name Returns: Dictionary with approvalPage, resume, and cancel URLs """ return _client.get_approval_urls(step_key, approver) @init_global_client def request_interactive_slack_approval( slack_resource_path: str, channel_id: str, message: str = None, approver: str = None, default_args_json: dict = None, dynamic_enums_json: dict = None, ) -> None: return _client.request_interactive_slack_approval( slack_resource_path=slack_resource_path, channel_id=channel_id, message=message, approver=approver, default_args_json=default_args_json, dynamic_enums_json=dynamic_enums_json, ) @init_global_client def send_teams_message( conversation_id: str, text: str, success: bool, card_block: dict = None ): """Send a message to a Microsoft Teams conversation. Args: conversation_id: Teams conversation ID text: Message text success: Whether to style as success message card_block: Optional adaptive card block Returns: HTTP response from Teams """ return _client.send_teams_message(conversation_id, text, success, card_block) @init_global_client def cancel_job(job_id: str, reason: str = None) -> str: """Cancel a specific job by ID. Args: job_id: UUID of the job to cancel reason: Optional reason for cancellation Returns: Response message from the cancel endpoint """ return _client.cancel_job(job_id, reason) @init_global_client def cancel_running() -> dict: """Cancel currently running executions of the same script.""" return _client.cancel_running() @init_global_client def run_script( path: str = None, hash_: str = None, args: dict = None, timeout: dt.timedelta | int | float = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = True, tag: str = None, ) -> Any: """Run script synchronously and return its result. .. deprecated:: Use run_script_by_path or run_script_by_hash instead. """ return _client.run_script( path=path, hash_=hash_, args=args, verbose=verbose, assert_result_is_not_none=assert_result_is_not_none, cleanup=cleanup, timeout=timeout, tag=tag, ) @init_global_client def run_script_by_path( path: str, args: dict = None, timeout: dt.timedelta | int | float = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = True, tag: str = None, ) -> Any: """Run script by path synchronously and return its result.""" return _client.run_script_by_path( path=path, args=args, verbose=verbose, assert_result_is_not_none=assert_result_is_not_none, cleanup=cleanup, timeout=timeout, tag=tag, ) @init_global_client def run_script_by_hash( hash_: str, args: dict = None, timeout: dt.timedelta | int | float = None, verbose: bool = False, cleanup: bool = True, assert_result_is_not_none: bool = True, tag: str = None, ) -> Any: """Run script by hash synchronously and return its result.""" return _client.run_script_by_hash( hash_=hash_, args=args, verbose=verbose, assert_result_is_not_none=assert_result_is_not_none, cleanup=cleanup, timeout=timeout, tag=tag, ) @init_global_client def run_inline_script_preview( content: str, language: str, args: dict = None, ) -> Any: """Run a script on the current worker without creating a job""" return _client.run_inline_script_preview( content=content, language=language, args=args, ) @init_global_client def username_to_email(username: str) -> str: """ Get email from workspace username .. deprecated:: Read the contextual variables instead: `os.environ.get("WM_END_USER_EMAIL") or os.environ.get("WM_EMAIL")`. WM_END_USER_EMAIL is the email of whoever triggered the run when it came from an app, so the fallback yields the app viewer inside an app and the executing user everywhere else - without an extra API call, and unlike this function it also resolves viewers who are not workspace members. An app viewed anonymously has no identity to report: the variable is then empty and the fallback yields the app publisher. """ return _client.username_to_email(username) @init_global_client def datatable(name: str = "main") -> DataTableClient: """Get a DataTable client for SQL queries. Args: name: Database name (default: "main") Returns: DataTableClient instance """ return _client.datatable(name) @init_global_client def ducklake(name: str = "main") -> DucklakeClient: """Get a DuckLake client for DuckDB queries. Args: name: Database name (default: "main") Returns: DucklakeClient instance """ return _client.ducklake(name) def parse_resource_syntax(s: str) -> Optional[str]: """Parse resource syntax from string.""" if s is None: return None if s.startswith("$res:"): return s[5:] if s.startswith("res://"): return s[6:] return None def parse_s3_object(s3_object: S3Object | str) -> S3Object: """Parse S3 object from a `s3:///` URI string (`s3:///` for the default storage) or S3Object format. Any other string raises rather than falling back to an auto-generated key: an auto key is requested by omitting the object, and a fallback would silently misplace the upload on any typo. """ if isinstance(s3_object, str): match = re.match(r'^s3://([^/]*)/(.+)$', s3_object) if match: return S3Object(s3=match.group(2), storage=match.group(1) or None) if s3_object.startswith("s3://"): raise ValueError( f"Invalid s3 object URI {s3_object!r}: expected " "s3:/// with a non-empty key " "(s3:/// for the default storage)" ) raise ValueError( f"Invalid s3 object {s3_object!r}: expected an s3:/// " f"URI (e.g. 's3:///{s3_object}' for key {s3_object!r} in the default " "storage) or S3Object(s3=)" ) else: return s3_object def parse_variable_syntax(s: str) -> Optional[str]: """Parse variable syntax from string.""" if s.startswith("var://"): return s[6:] return None def append_to_result_stream(text: str) -> None: """Append a text to the result stream. Args: text: text to append to the result stream """ print("WM_STREAM: {}".format(text.replace(chr(10), '\\n'))) def stream_result(stream) -> None: """Stream to the result stream. Args: stream: stream to stream to the result stream """ for text in stream: append_to_result_stream(text) class DataTableClient: """Client for executing SQL queries against Windmill DataTables.""" def __init__(self, client: Windmill, name: str): """Initialize DataTableClient. Args: client: Windmill client instance name: DataTable name """ self.client = client self.name, self.schema = parse_sql_client_name(name) def query(self, sql: str, *args) -> SqlQuery: """Execute a SQL query against the DataTable. Args: sql: SQL query string with $1, $2, etc. placeholders *args: Positional arguments to bind to query placeholders Returns: SqlQuery instance for fetching results """ if self.schema is not None: sql = f'SET search_path TO "{self.schema}";\n' + sql args_dict = {} args_def = "" for i, arg in enumerate(args): args_dict[f"arg{i+1}"] = arg args_def += f"-- ${i+1} arg{i+1} ({infer_sql_type(arg)})\n" sql = args_def + sql return SqlQuery( sql, lambda sql: self.client.run_inline_script_preview( content=sql, language="postgresql", args={"database": f"datatable://{self.name}", **args_dict}, ) ) class DucklakeClient: """Client for executing DuckDB queries against Windmill DuckLake.""" def __init__(self, client: Windmill, name: str): """Initialize DucklakeClient. Args: client: Windmill client instance name: DuckLake database name """ self.client = client self.name = name def query(self, sql: str, **kwargs): """Execute a DuckDB query against the DuckLake database. Args: sql: SQL query string with $name placeholders **kwargs: Named arguments to bind to query placeholders Returns: SqlQuery instance for fetching results """ args_dict = {} args_def = "" for key, value in kwargs.items(): args_dict[key] = value args_def += f"-- ${key} ({infer_sql_type(value)})\n" attach = f"ATTACH 'ducklake://{self.name}' AS dl;USE dl;\n" sql = args_def + attach + sql return SqlQuery( sql, lambda sql: self.client.run_inline_script_preview( content=sql, language="duckdb", args=args_dict, ) ) def _qualified(self, table: str, schema: str = None) -> str: return f'dl."{schema}"."{table}"' if schema else f"dl.{table}" def _materialize_finish(self, sql, table, schema, partition, partition_col): """Return the materialize query; in a pipeline (WM_PIPELINE) append a summary read and record materialized_partition state after a successful run so SDK-materialized slices appear in the grid like `// materialize` ones. Outside a pipeline it stays a plain query (no recording).""" bind = {} if partition is None else {"_wm_partition": partition} if os.environ.get("WM_PIPELINE") != "true": return self.query(sql, **bind) t = self._qualified(table, schema) where = f" WHERE {partition_col} = $_wm_partition" if partition is not None else "" summary = ( f"\nSELECT (SELECT count(*) FROM {t}{where}) AS rows, " f"(SELECT max(snapshot_id) FROM ducklake_snapshots('dl')) AS snapshot_id;" ) q = self.query(sql + summary, **bind) # Asset path mirrors the `// materialize` engine: /. # for an explicit schema, else /
. Dropping the schema would # hide the row from the grid and collide distinct schemas under one key. asset_path = f"{self.name}/{schema}.{table}" if schema else f"{self.name}/{table}" return _RecordingSqlQuery(q, self.client, asset_path, partition or "") def upsert_partition( self, table: str, select_sql: str, partition: str = None, unique_key: str = None, partition_col: str = "_wm_partition", schema: str = None, ): """Idempotently materialize the rows of `select_sql` into ducklake `table` for one `partition` (or the whole table when `partition` is None). Client-side equivalent of the `// materialize` engine: with `unique_key` it upserts within the slice (delete-by-key + insert); without it, it replaces (whole table → CREATE OR REPLACE; partition → delete the partition + insert). Re-running the same slice is safe — the backfill / failure-recovery contract. The partition value is bound as a DuckDB arg (never string-interpolated) so it cannot inject SQL. `select_sql` is trusted (your own query). """ t = self._qualified(table, schema) # Whole-table (no partition): no partition column; replace rebuilds the # table with CREATE OR REPLACE, merge upserts the whole table by key. if partition is None: if unique_key: sql = ( f"CREATE TABLE IF NOT EXISTS {t} AS SELECT * FROM ({select_sql}) WHERE false;\n" f"BEGIN TRANSACTION;\n" f"DELETE FROM {t} WHERE {unique_key} IN (SELECT {unique_key} FROM ({select_sql}));\n" f"INSERT INTO {t} SELECT * FROM ({select_sql});\n" f"COMMIT;" ) else: sql = f"CREATE OR REPLACE TABLE {t} AS SELECT * FROM ({select_sql});" return self._materialize_finish(sql, table, schema, partition, partition_col) src = f"SELECT *, $_wm_partition AS {partition_col} FROM ({select_sql})" if unique_key: # Upsert via delete-by-key + insert (not MERGE — DuckLake's MERGE # fails writing the first rows of a fresh partition). body = ( f"DELETE FROM {t} WHERE {partition_col} = $_wm_partition " f"AND {unique_key} IN (SELECT {unique_key} FROM ({select_sql}));\n" f"INSERT INTO {t} {src};" ) else: body = ( f"DELETE FROM {t} WHERE {partition_col} = $_wm_partition;\n" f"INSERT INTO {t} {src};" ) sql = ( f"CREATE TABLE IF NOT EXISTS {t} AS " f"SELECT *, CAST(NULL AS VARCHAR) AS {partition_col} FROM ({select_sql}) WHERE false;\n" f"ALTER TABLE {t} SET PARTITIONED BY ({partition_col});\n" f"BEGIN TRANSACTION;\n{body}\nCOMMIT;" ) return self._materialize_finish(sql, table, schema, partition, partition_col) def append_partition( self, table: str, select_sql: str, partition: str = None, partition_col: str = "_wm_partition", schema: str = None, ): """INSERT-only materialization (no dedup / no replace) for an immutable event-log table — for one `partition`, or the whole table when `partition` is None. NOTE: unlike `upsert_partition`, re-running the same slice duplicates rows — use only for append-only sources.""" t = self._qualified(table, schema) # Whole-table (no partition): insert into the bare table, no partition col. if partition is None: sql = ( f"CREATE TABLE IF NOT EXISTS {t} AS SELECT * FROM ({select_sql}) WHERE false;\n" f"INSERT INTO {t} SELECT * FROM ({select_sql});" ) return self._materialize_finish(sql, table, schema, partition, partition_col) sql = ( f"CREATE TABLE IF NOT EXISTS {t} AS " f"SELECT *, CAST(NULL AS VARCHAR) AS {partition_col} FROM ({select_sql}) WHERE false;\n" f"ALTER TABLE {t} SET PARTITIONED BY ({partition_col});\n" f"INSERT INTO {t} SELECT *, $_wm_partition AS {partition_col} FROM ({select_sql});" ) return self._materialize_finish(sql, table, schema, partition, partition_col) def read( self, table: str, partition: str = None, partition_col: str = "_wm_partition", schema: str = None, ): """Read a materialized ducklake table, optionally a single partition.""" t = self._qualified(table, schema) if partition is not None: return self.query( f"SELECT * FROM {t} WHERE {partition_col} = $_wm_partition", _wm_partition=partition, ) return self.query(f"SELECT * FROM {t}") class SqlQuery: """Query result handler for DataTable and DuckLake queries.""" def __init__(self, sql: str, fetch_fn): """Initialize SqlQuery. Args: sql: SQL query string fetch_fn: Function to execute the query """ self.sql = sql self.fetch_fn = fetch_fn def fetch(self, result_collection: str | None = None): """Execute query and fetch results. Args: result_collection: Optional result collection mode Returns: Query results """ sql = self.sql if result_collection is not None: sql = f'-- result_collection={result_collection}\n{sql}' return self.fetch_fn(sql) def fetch_one(self): """Execute query and fetch first row of results. Returns: First row of query results """ return self.fetch(result_collection="last_statement_first_row") def fetch_one_scalar(self): """Execute query and fetch first row of results. Return result as a scalar value. Returns: First row of query result as a scalar value """ return self.fetch(result_collection="last_statement_first_row_scalar") def execute(self): """Execute query and don't return any results. """ self.fetch_one() class _RecordingSqlQuery: """Wraps a ducklake materialize query so that, on a successful run, the trailing summary (row count + snapshot id) is captured and the materialized_partition state is recorded (best-effort). Only used in pipeline context — outside it the helpers return a plain SqlQuery. Mirrors SqlQuery's terminal methods so `.execute()` / `.fetch_one()` behave the same.""" def __init__(self, inner, client, asset_path, partition): self._inner = inner self._client = client self._asset_path = asset_path self._partition = partition self.sql = inner.sql def execute(self): self._run() def fetch_one(self): return self._run() def fetch(self, result_collection=None): return self._run() def _run(self): try: row = self._inner.fetch_one() except Exception as e: self._record("failed", None, None, str(e)) raise snap = row.get("snapshot_id") if isinstance(row, dict) else None rows = row.get("rows") if isinstance(row, dict) else None self._record("materialized", snap, rows, None) return row def _record(self, status, snapshot_id, row_count, error): try: self._client.post( f"/w/{self._client.workspace}/assets/record_materialization", json={ "asset_kind": "ducklake", "asset_path": self._asset_path, "partition": self._partition, "status": status, "snapshot_id": snapshot_id, "row_count": row_count, "job_id": os.environ.get("WM_JOB_ID"), "error": error, }, ) except Exception: pass # best-effort; never fail the user's materialization def infer_sql_type(value) -> str: """ DuckDB executor requires explicit argument types at declaration These types exist in both DuckDB and Postgres Check that the types exist if you plan to extend this function for other SQL engines. """ if isinstance(value, bool): # Check bool before int since bool is a subclass of int in Python return "BOOLEAN" elif isinstance(value, int): return "BIGINT" elif isinstance(value, float): return "FLOAT8" elif value is None: return "TEXT" elif isinstance(value, str): return "TEXT" elif isinstance(value, dict) or isinstance(value, list): return "JSON" else: return "TEXT" def parse_sql_client_name(name: str) -> tuple[str, Optional[str]]: name = name schema = None if ":" in name: name, schema = name.split(":", 1) if not name: name = "main" return name, schema # ── Workflow-as-Code SDK ────────────────────────────────────────────── import asyncio as _asyncio import contextvars as _contextvars import sys as _sys import traceback as _traceback def _assert_usable_step_key(key: str, what: str) -> None: """A step key travels as one path segment when its URLs are minted, so it must be non-empty and free of ``/`` and dot segments — otherwise ``wait_for_approval`` would accept a key ``get_approval_urls`` can never address.""" k = key.strip() if not k or k in (".", "..") or "/" in key or "\\" in key: raise RuntimeError(f"{what} must be a non-empty step name without `/` or dot segments") class _StepSuspend(BaseException): """Raised to suspend workflow execution. Inherits from BaseException so it is not caught by bare `except Exception:` blocks.""" def __init__(self, dispatch_info: dict): self.dispatch_info = dispatch_info class _StepFailure(BaseException): """Carries the exception raised by the step a child round executes directly. That exception *is* the round's result, so a broad ``except Exception`` in the body must not be able to turn it into a successful complete — the parent would then record the caught branch's value as the step result. BaseException for the same reason ``_StepSuspend`` is; a bare ``except:`` still swallows both. """ def __init__(self, exc: BaseException): self.exc = exc class TaskError(Exception): """Raised when a WAC ``task`` or ``step`` failed. Attributes: step_key: The checkpoint key of the failed step. child_job_id: The UUID of the failed child job, or ``None`` for a ``step()``, which runs in the workflow job and has no child job. result: ``{"error": {"name", "message", "stack"?, "extra"?}}`` — the same shape whether a task or a step failed. ``name`` and ``message`` are always present; ``stack`` only when the failure had a traceback, and ``extra`` only when it carried custom fields of its own, dropped with ``extra_omitted: True`` beside it when too large to checkpoint. """ def __init__(self, message: str, *, step_key: str = "", child_job_id: Optional[str] = None, result=None): super().__init__(message) self.step_key = step_key self.child_job_id = child_job_id self.result = result def _safe_str(o) -> str: """``str()`` on the failing side's own object, which can raise in turn — a detached ORM row, a proxy over a closed connection, an ``__str__`` that itself fails. Every coercion here runs inside the ``except`` that is reporting the user's failure, so an escape would replace their error with an unrelated one and skip the checkpoint entirely.""" try: return str(o) except Exception: return f"" def _step_error_stack(exc: BaseException) -> str: """The traceback of a failed ``step()`` body, formatted the way the python executor formats a failed job's: frames only, and the frame that called into the user's code dropped. Here that first frame is ``_run_inline_step``'s own ``result = fn()``, the counterpart of the generated wrapper frame the executor strips, so a step's stack and a task's stack read alike. Taken from ``sys.exc_info()`` the way the executor takes it, falling back to the attribute: an exception overriding ``__getattribute__`` makes reading ``__traceback__`` raise, and this runs inside the ``except`` reporting the user's failure, so an escape would lose both their error and the checkpoint. """ tb = _sys.exc_info()[2] if tb is None: try: tb = exc.__traceback__ except Exception: return "" try: return "".join(_traceback.format_tb(tb)[1:]).strip() except Exception: return "" def _json_round_trip(value): """Put a value through the checkpoint's encoding without checkpointing it, so the paths that never persist anything still hand back the shape the ones that do would. ``default=str`` matches the worker wrapper's encoder.""" return json.loads(json.dumps(value, default=str)) def _step_error_marker(key: str, exc: BaseException) -> dict: """Serialize a failed ``step()`` body into the ``__wmill_error`` marker that task failures also use, so it can be stored in ``completed_steps``. The marker's final shape is decided by the backend (``wac_failure_record``), which normalizes task failures through the same function; what is built here is the raw material plus the envelope the backend recognizes.""" error = {"name": type(exc).__name__, "message": _safe_str(exc)} stack = _step_error_stack(exc) if stack: error["stack"] = stack # Custom attributes go under ``extra``, the same key the python executor uses # for a failed child job, so an exception carrying e.g. a ``code`` keeps it # whether it failed as a task or as a step. # # Coerced through ``default=str`` the way the executor writes its own error: # the fast-path POST serializes strictly, and the commonest failing step # there is — ``resp.raise_for_status()``, whose ``__dict__`` holds a request # and a response object — would otherwise fail to serialize and silently # drop every such failure onto the slow suspend-and-replay path. # Everything about the failing exception can fight back, and this runs inside # the ``except`` reporting it, so an escape replaces the user's error and # skips the checkpoint. Only a genuine ``dict`` is walked: an overridden # ``__dict__`` can raise on access, on ``.items()``, or yield non-pairs. try: _raw_extra = getattr(exc, "__dict__", None) except Exception: _raw_extra = None if type(_raw_extra) is dict and _raw_extra: safe_extra = {} for _k, _v in _raw_extra.items(): # Rebuilding the pair hashes the key again, so only the types json # can represent, and exactly those: a subclass may define __hash__. if type(_k) not in (str, int, float, bool, type(None)): continue # Per attribute so one bad value cannot take the rest, and as a pair # so an int/bool/None key arrives as the string a replay reads. # ``parse_constant`` catches what ``default`` cannot: a float is # serializable, so NaN/Infinity would go out as invalid JSON. try: safe_extra.update( json.loads( json.dumps({_k: _v}, default=_safe_str), parse_constant=lambda c: c, ) ) except (TypeError, ValueError, RecursionError): pass if safe_extra: error["extra"] = safe_extra return { "__wmill_error": True, "message": _safe_str(exc), "step_key": key, "result": {"error": error}, } def _task_error_from_marker(marker: dict, fallback_message: str) -> TaskError: """Rebuild the exception a failed task or step raises. The run that produced the failure and every later replay go through here: ``except`` is control flow, ``@workflow`` re-runs its body from the top every round, so a handler that branches on the failure it caught must be handed the same thing in every round or it dispatches different tasks on the way back.""" return TaskError( marker.get("message") or fallback_message, step_key=marker.get("step_key", ""), child_job_id=marker.get("child_job_id"), result=marker.get("result"), ) # The worker deserializes a sleep into a ``u32`` of seconds and fails the whole # job on anything wider, so a delay a multiplier has run away with has to be # capped here rather than sent. _MAX_SLEEP_SECONDS = 2**32 - 1 _RETRY_KEYS = ("attempts", "delay", "multiplier", "max_delay") # Every attempt claims its keys before the first one is dispatched, so an # unbounded ``attempts`` is a workflow that hangs allocating rather than a very # patient one. _MAX_RETRY_ATTEMPTS = 100 def _checked_retry(retry: Optional[dict]) -> Optional[dict]: """Reject a policy where it is written, rather than mid-run on a replay: the policy is a plain dict, so a misspelled key would otherwise be dropped in silence and the task would retry on a policy nobody wrote.""" if retry is None: return None unknown = sorted(k for k in retry if k not in _RETRY_KEYS) if unknown: raise ValueError( f"unknown retry option(s): {', '.join(unknown)}. Expected any of: {', '.join(_RETRY_KEYS)}" ) attempts = retry.get("attempts") if isinstance(attempts, bool) or not isinstance(attempts, int) or not 0 <= attempts <= _MAX_RETRY_ATTEMPTS: raise ValueError( f"retry attempts must be a whole number between 0 and {_MAX_RETRY_ATTEMPTS}, got {attempts!r}" ) return retry def _retry_delay_seconds(retry: dict, attempt: int) -> int: """Seconds to wait before retry number ``attempt`` (0 is the first retry).""" base = retry.get("delay") or 0 if base <= 0: return 0 # `or 1` would read an explicit `multiplier: 0` — every retry after the # first going out with no wait — as the default of 1. multiplier = retry.get("multiplier") if multiplier is None: multiplier = 1 try: grown = base * multiplier**attempt except OverflowError: # A float delay times an integer multiplier raised past ~1e308. grown = _MAX_SLEEP_SECONDS max_delay = retry.get("max_delay") if max_delay is not None: grown = min(grown, max_delay) return max(0, int(min(grown, _MAX_SLEEP_SECONDS))) _workflow_ctx: _contextvars.ContextVar["WorkflowCtx"] = _contextvars.ContextVar( "_workflow_ctx" ) class WorkflowCtx: """Internal context for workflow replay/suspension. Not user-facing — set implicitly by ``@workflow`` via contextvars. """ def __init__(self, checkpoint: dict | None = None): checkpoint = checkpoint or {} self._completed: dict = checkpoint.get("completed_steps", {}) self._counters: dict[str, int] = {} # Every key handed out by _alloc_key, so distinct names can't alias one key. self._used_keys: set[str] = set() self._pending: list = [] self._executing_key: str | None = checkpoint.get("_executing_key") # Reuse a single httpx.AsyncClient across all fast-path step() calls # in this workflow invocation. Instantiating a fresh client per call # allocates a new connection pool each time — on localhost this adds # ~15ms per step, dominating the end-to-end cost. Lazily built so no # client is created for workflows that never hit the fast path. self._inline_http_client: "httpx.AsyncClient | None" = None # Serializes fast-path POSTs across concurrent step() calls within # one workflow invocation. Wraps only the HTTP call, not fn() — so # `asyncio.gather(step("a", fn_a), step("b", fn_b))` still runs the # two fn() bodies in parallel, only the API requests are ordered. # This closes the first-write race window against `SELECT FOR UPDATE` # on a not-yet-created `v2_job_status` row: concurrent POSTs would # both see None and both overwrite each other's checkpoint because # the helper writes the whole serialized `_checkpoint` object, not # a single `completed_steps[key]`. Lazily built so the ctx can be # constructed outside an event loop (tests do this). self._inline_lock: "_asyncio.Lock | None" = None def _alloc_key(self, name: str = "step") -> str: """Name-based key: ``double`` for first call, ``double_2``, ``double_3`` for subsequent. Suffixing alone can alias — a second ``step("x")`` and a first ``step("x_2")`` both want ``x_2`` — so keep bumping past keys already handed out. Allocation order is fixed by the workflow body, so replays reproduce the same keys. """ n = self._counters.get(name, 0) + 1 key = name if n == 1 else f"{name}_{n}" while key in self._used_keys: n += 1 key = f"{name}_{n}" self._counters[name] = n self._used_keys.add(key) return key def _next_step(self, name: str, script: str, func=None, dispatch_type: str = "inline", _task_options: Optional[dict] = None, **kwargs): """Return an awaitable that either resolves from cache or suspends.""" step_name = name or script or "step" retry = (_task_options or {}).get("retry") or {} # Clamped as well as validated at decoration: a policy that reached here # another way must not spin the key loop below. max_retries = min(max(0, int(retry.get("attempts") or 0)), _MAX_RETRY_ATTEMPTS) # Claimed up front, all of them, and named off the first attempt's key: # one allocated later would shift the keys of the steps beside it, and a # ``step()`` named ``t#2`` — names are arbitrary — could alias one. # Whichever is allocated second is the one renamed, in every round alike. base_key = self._alloc_key(step_name) attempt_keys = [base_key] backoff_keys = [] for i in range(max_retries): backoff_keys.append(self._alloc_key(f"{base_key}#retry{i + 2}")) attempt_keys.append(self._alloc_key(f"{base_key}#{i + 2}")) # One pass per attempt. Every attempt the checkpoint already holds is # decided here — a failed one either retries (moving to the next key) or # is handed back to the body — so the loop always ends at the first # attempt that has yet to run. attempt = 0 while True: key = attempt_keys[attempt] if key in self._completed: val = self._completed[key] if isinstance(val, dict) and val.get("__wmill_error"): if attempt < max_retries: self._retry_backoff(backoff_keys[attempt], base_key, retry, attempt) attempt += 1 continue raise _task_error_from_marker(val, f"Task '{name}' failed") return self._resolved(val) if self._executing_key is not None: if key == self._executing_key: return self._execute_directly(func, **kwargs) else: return self._never_resolve() print(f"\n--- WAC: {key} ---") info = {"name": name or key, "script": script or key, "args": kwargs, "key": key, "dispatch_type": dispatch_type} if _task_options: for opt_key in ("timeout", "tag", "cache_ttl", "priority", "concurrent_limit", "concurrency_key", "concurrency_time_window_s", "fn_id"): if opt_key in _task_options and _task_options[opt_key] is not None: info[opt_key] = _task_options[opt_key] self._pending.append(info) return self._suspend() def _retry_backoff(self, key: str, base_key: str, retry: dict, attempt: int) -> None: """Wait out the backoff between two attempts of a retried task, as a durable sleep, and return once there is nothing to wait for — no delay configured, or the sleep already in the checkpoint. Raises where it stands rather than from a coroutine the caller has to await: a task call the body never awaits is still dispatched (the runner flushes ``_pending``), so a backoff that only fired when awaited would drop the retry and let the round report the workflow complete.""" seconds = _retry_delay_seconds(retry, attempt) if seconds < 1: return if key in self._completed: return # Child mode never raises: the parent dispatched this child only after # its own round had slept, so the loop moves on to the attempt being # executed. if self._executing_key is not None: return print(f"\n--- WAC: sleep({key}, {seconds}s) before retrying {base_key} ---") raise _StepSuspend({"mode": "sleep", "key": key, "seconds": seconds, "steps": []}) async def _resolved(self, value): return value async def _execute_directly(self, func, **kwargs): try: result = func(**kwargs) if _asyncio.iscoroutine(result): result = await result except Exception as exc: raise _StepFailure(exc) from exc raise _StepSuspend({"mode": "step_complete", "steps": [], "result": result}) async def _never_resolve(self): await _asyncio.Future() async def _suspend(self): steps = list(self._pending) self._pending.clear() raise _StepSuspend( { "mode": "parallel" if len(steps) > 1 else "sequential", "steps": steps, } ) async def _wait_for_approval( self, timeout: int = 1800, form: dict | None = None, self_approval: bool = True, key: str | None = None, skin: str | None = None, description: str | dict | None = None, ): if key is not None: _assert_usable_step_key(key, "wait_for_approval key") requested_key, key = key, self._alloc_key(key or "approval") # An explicit key is an identifier callers mint URLs against, so silently # renaming a duplicate to ``_2`` would hand them a URL for the *first* # step — which then fails with "resume request already sent" and parks the # workflow until timeout. Unnamed approvals keep auto-numbering. if requested_key and key != requested_key: raise RuntimeError( f'WAC step key "{requested_key}" is already used in this workflow. ' "Give each wait_for_approval() its own key so get_approval_urls() can address it." ) if key in self._completed: return self._completed[key] if self._executing_key is not None: await _asyncio.Future() print(f"\n--- WAC: wait_for_approval({key}) ---") raise _StepSuspend({ "mode": "approval", "key": key, "timeout": timeout, "form": form, "self_approval_disabled": not self_approval, "skin": skin, "description": description, "steps": [], }) async def _sleep(self, seconds: int): key = self._alloc_key("sleep") if key in self._completed: return if self._executing_key is not None: await _asyncio.Future() print(f"\n--- WAC: sleep({key}, {seconds}s) ---") raise _StepSuspend({ "mode": "sleep", "key": key, "seconds": max(1, int(seconds)), "steps": [], }) async def _run_inline_step(self, name: str, fn): import json as _json_mod import time as _time_mod from datetime import datetime as _dt, timezone as _tz key = self._alloc_key(name or "step") if key in self._completed: val = self._completed[key] if isinstance(val, dict) and val.get("__wmill_error"): raise _task_error_from_marker(val, f"Step '{name}' failed") return val if self._executing_key is not None: await _asyncio.Future() print(f"\n--- WAC: {key} ---") started_at = _dt.now(_tz.utc).isoformat() print(f"WM_WAC_STEP: {_json_mod.dumps({'key': key, 'started_at': started_at})}") t0 = _time_mod.monotonic() # A raised step still has to reach ``completed_steps``, or a replay with # ``_executing_key`` set finds nothing recorded and parks forever on the # ``_asyncio.Future()`` above. The control-flow signals (``_StepSuspend``, # ``_StepFailure``) and ``CancelledError`` are ``BaseException``, so they # pass through untouched. step_failed = False try: result = fn() if _asyncio.iscoroutine(result): result = await result except Exception as _exc: step_failed = True result = _step_error_marker(key, _exc) # The failure is reported as a value from here on, so nothing else # prints the traceback. Without this a step that fails and is never # caught leaves a job log whose deepest frame is inside this client. print(f"--- WAC: {key} failed ---") print(f"{type(_exc).__name__}: {_safe_str(_exc)}") _step_stack = result["result"]["error"].get("stack") if _step_stack: print(_step_stack) duration_ms = int((_time_mod.monotonic() - t0) * 1000) # Fast path: POST the delta to the new per-job API endpoint and return # the result directly, letting the workflow subprocess continue into # the next step() without unwinding. On any failure — network, auth, # timeout, source-hash mismatch, old backend without the endpoint — # fall through to raising _StepSuspend so the worker takes the legacy # suspend-and-replay path. Gated by WM_WAC_INLINE_FAST_PATH (default # on) so the old behavior stays reachable for A/B testing and rollback. _fast_path_flag = os.environ.get("WM_WAC_INLINE_FAST_PATH", "1").strip().lower() _fast_path_enabled = _fast_path_flag not in ("0", "false", "off", "no") _job_id = os.environ.get("WM_JOB_ID") _workspace = os.environ.get("WM_WORKSPACE") _base = os.environ.get("BASE_INTERNAL_URL") _token = os.environ.get("WM_TOKEN") if _fast_path_enabled and _job_id and _workspace and _base and _token: _fast_path_ok = False _stored_failure = None _replay_result = None try: # ``default=str`` is the encoder the worker wrapper uses on the # suspend path, so both arms checkpoint the same value — and a # datetime or set takes the fast path instead of silently # degrading to a suspend round. _payload = _json_mod.dumps( { "key": key, "result": result, "started_at": started_at, "duration_ms": duration_ms, }, default=str, ) _replay_result = _json_mod.loads(_payload)["result"] if self._inline_lock is None: self._inline_lock = _asyncio.Lock() # Lock wraps only the POST, not fn() above — concurrent # step() calls run fn() in parallel, then serialize on # the API request. async with self._inline_lock: if self._inline_http_client is None: self._inline_http_client = httpx.AsyncClient( timeout=httpx.Timeout(10.0), headers={ "Authorization": f"Bearer {_token}", "Content-Type": "application/json", }, ) _resp = await self._inline_http_client.post( f"{_base}/api/w/{_workspace}/jobs/wac/inline_checkpoint/{_job_id}", content=_payload, ) _resp.raise_for_status() if step_failed: # The backend normalizes the failure before storing it, # and hands back what it stored. Raising from that, not # from the marker posted above, is what makes this round # and every replay read the same record even if the two # sides ever disagree about how to build one. # # A backend predating the echo answers without a JSON # body, and the round-tripped marker below stands in. A # JSON body that will not parse is different: the record # may already be committed and its content is unknown, # so let it raise and take the suspend path instead. if "json" in _resp.headers.get("content-type", ""): _stored_failure = (_resp.json() or {}).get("failure") _fast_path_ok = True except Exception as _e: logger.info( "WAC v2 inline fast path failed for key %s, falling back to suspend: %s", key, _e, ) # fall through to the legacy suspend path if _fast_path_ok: # Raise what a replay would rebuild from the record, never the # original: a replay cannot reconstruct the original type, so # raising it here would make ``except ValueError:`` catch on this # run and miss on the next. Nothing is chained onto # ``__cause__`` for the same reason — the traceback a replay can # still show is in ``result["error"]["stack"]``. ``_stored_failure`` # is None against a backend that predates the echoed record, # which is what ``_replay_result`` below stands in for. if step_failed: # ``_replay_result``, not ``result``: the fallback has to be # what the checkpoint holds, so the round that ran the body # reads what every replay of it will. raise _task_error_from_marker( _stored_failure or _replay_result, f"Step '{name}' failed" ) # Return the round trip of what was checkpointed, never the # in-memory value: handing back the live object would let the # round that ran the body branch on a type — tuple, datetime — # that no replay of it ever sees. return _replay_result raise _StepSuspend({ "mode": "inline_checkpoint", "steps": [], "key": key, "result": result, "started_at": started_at, "duration_ms": duration_ms, }) def _fn_fingerprint(func) -> str: """A stable identity for a task's code, what its cached result is keyed on: a name is shared by any two tasks called the same, and a step key by any two tasks called at the same position, so neither can tell them apart.""" import hashlib import inspect import marshal try: src = inspect.getsource(func).encode() except (OSError, TypeError): # No source on disk: the whole code object, constants and names included. src = marshal.dumps(func.__code__) return hashlib.sha256(src).hexdigest() def task( _func=None, *, path: Optional[str] = None, tag: Optional[str] = None, timeout: Optional[int] = None, cache_ttl: Optional[int] = None, priority: Optional[int] = None, concurrency_limit: Optional[int] = None, concurrency_key: Optional[str] = None, concurrency_time_window_s: Optional[int] = None, retry: Optional[dict] = None, ): """Decorator that marks a function as a workflow task. Works in both WAC v1 (sync, HTTP-based dispatch) and WAC v2 (async, checkpoint/replay) modes: - **v2 (inside @workflow)**: dispatches as a checkpoint step. - **v1 (WM_JOB_ID set, no @workflow)**: dispatches via HTTP API. - **Standalone**: executes the function body directly. A task runs as its own job, so its result is always encoded as JSON and decoded back before the caller sees it: a ``datetime`` comes back as a string, a tuple as a list. ``retry`` re-dispatches the task after a failure, inside ``@workflow`` only. Every attempt is a step of its own (``call_api``, ``call_api#2``, ...) and the wait between two of them is a durable sleep, so a retrying task holds no worker while it backs off. Keys: ``attempts`` (retries after the first failure, a whole number from 0 to 100), ``delay`` (seconds before the first retry, sub-second delays dropped), ``multiplier`` (applied to the delay after each attempt, 1 keeps it constant), ``max_delay`` (ceiling in seconds). ``attempts`` is required, and an out-of-range or unknown key is rejected where the policy is written. A workflow sleeps once per round, so tasks backing off in the same fan-out wait one after another rather than together: the delay before a fan-out retries is the sum of every backoff pending in it, not the longest one, and it grows with both the width of the fan-out and ``attempts``. Retries with no ``delay`` all go out in a single round. ``cache_ttl`` serves a previous result of the task for that many seconds instead of running it again. The result is keyed on the task and the arguments it is called with, so anything a cached task reads from its closure must be passed in as an argument. Usage:: @task async def extract_data(url: str): ... @task(path="f/external_script", timeout=600, tag="gpu") async def run_external(x: int): ... @task(retry={"attempts": 3, "delay": 30, "multiplier": 2}) async def call_api(payload: dict): ... """ from inspect import signature as _sig _task_opts = { "timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry), } # Remove None values _task_opts = {k: v for k, v in _task_opts.items() if v is not None} or None def decorator(func) -> Callable[..., Any]: task_path = path task_name = func.__name__ _fn_opts = {**(_task_opts or {}), "fn_id": _fn_fingerprint(func)} _params_list = list(_sig(func).parameters) def _merge_args(args, kwargs): merged = dict(kwargs) for i, arg in enumerate(args): if i < len(_params_list): key = _params_list[i] if key not in merged: merged[key] = arg else: merged[f"arg{i}"] = arg return merged # Keeps the decorated function's identity: `@task` is applied to a # top-level `async def`, and a caller introspecting it should see that # function, not `wrapper`. The step key is computed from `func` above, # so this does not affect dispatch. @functools.wraps(func) def wrapper(*args, **kwargs): # WAC v2: inside a @workflow context ctx = _workflow_ctx.get(None) if ctx is not None: script = task_path if task_path else task_name merged = _merge_args(args, kwargs) return ctx._next_step(task_name, script, func, _task_options=_fn_opts, **merged) # WAC v1: running inside a Windmill job but not in a @workflow if ( os.environ.get("WM_JOB_ID") is not None and os.environ.get("MAIN_OVERRIDE") != func.__name__ ): global _client if _client is None: _client = Windmill() w_id = os.environ.get("WM_WORKSPACE") job_id = os.environ.get("WM_JOB_ID") json_args = _merge_args(args, kwargs) api_params = {} if tag is not None: api_params["tag"] = tag resp = _client.post( f"/w/{w_id}/jobs/run/workflow_as_code/{job_id}/{func.__name__}", json={"args": json_args}, params=api_params, ) child_job_id = resp.text print(f"Executing task {func.__name__} on job {child_job_id}") job_result = _client.wait_job(child_job_id) print(f"Task {func.__name__} ({child_job_id}) completed") return job_result # Standalone — execute directly, but round-trip the result: a task's # value crosses JSON in every other path, so a local run must agree. # This wrapper is sync, so an ``async def`` task hands back a # coroutine here — round-tripping that would serialize the coroutine # object itself. result = func(*args, **kwargs) if _asyncio.iscoroutine(result): async def _round_trip_awaited(): return _json_round_trip(await result) return _round_trip_awaited() return _json_round_trip(result) wrapper._is_task = True wrapper._task_path = task_path return wrapper if _func is not None: # @task without parentheses return decorator(_func) # @task() or @task(path="...", tag="...") return decorator def task_script( path: str, *, timeout: Optional[int] = None, tag: Optional[str] = None, cache_ttl: Optional[int] = None, priority: Optional[int] = None, concurrency_limit: Optional[int] = None, concurrency_key: Optional[str] = None, concurrency_time_window_s: Optional[int] = None, retry: Optional[dict] = None, ): """Create a task that dispatches to a separate Windmill script. ``retry`` takes the same policy as :func:`task`. Usage:: extract = task_script("f/data/extract", timeout=600) @workflow async def main(): data = await extract(url="https://...") """ name = path.rsplit("/", 1)[-1] _opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry)}.items() if v is not None} or None def wrapper(**kwargs): ctx = _workflow_ctx.get(None) if ctx is not None: return ctx._next_step(name, path, dispatch_type="script", _task_options=_opts, **kwargs) raise RuntimeError(f'task_script("{path}") can only be called inside a @workflow') wrapper.__name__ = name wrapper._is_task = True wrapper._task_path = path return wrapper def task_flow( path: str, *, timeout: Optional[int] = None, tag: Optional[str] = None, cache_ttl: Optional[int] = None, priority: Optional[int] = None, concurrency_limit: Optional[int] = None, concurrency_key: Optional[str] = None, concurrency_time_window_s: Optional[int] = None, retry: Optional[dict] = None, ): """Create a task that dispatches to a separate Windmill flow. ``retry`` takes the same policy as :func:`task`. Usage:: pipeline = task_flow("f/etl/pipeline", priority=10) @workflow async def main(): result = await pipeline(input=data) """ name = path.rsplit("/", 1)[-1] _opts = {k: v for k, v in {"timeout": timeout, "tag": tag, "cache_ttl": cache_ttl, "priority": priority, "concurrent_limit": concurrency_limit, "concurrency_key": concurrency_key, "concurrency_time_window_s": concurrency_time_window_s, "retry": _checked_retry(retry)}.items() if v is not None} or None def wrapper(**kwargs): ctx = _workflow_ctx.get(None) if ctx is not None: return ctx._next_step(name, path, dispatch_type="flow", _task_options=_opts, **kwargs) raise RuntimeError(f'task_flow("{path}") can only be called inside a @workflow') wrapper.__name__ = name wrapper._is_task = True wrapper._task_path = path return wrapper def workflow(func): """Decorator marking an async function as a workflow-as-code entry point. The function must be **deterministic**: given the same inputs it must call tasks in the same order on every replay. Branching on task results is fine (results are replayed from checkpoint), but branching on external state (current time, random values, external API calls) must use ``step()`` to checkpoint the value so replays see the same result. """ func._is_workflow = True return func async def step(name: str, fn): """Execute ``fn`` inline and checkpoint the result. On replay the cached value is returned without re-executing ``fn``. Use for lightweight deterministic operations (timestamps, random IDs, config reads) that should not incur the overhead of a child job. ``fn``'s result is encoded as JSON and decoded back before it is returned, so the round that runs the body sees the same types every replay sees: a ``datetime`` comes back as a string, a tuple as a list. """ ctx: WorkflowCtx | None = _workflow_ctx.get(None) if ctx is not None: return await ctx._run_inline_step(name, fn) result = fn() if _asyncio.iscoroutine(result): result = await result # Outside a workflow nothing is checkpointed, but round-trip anyway: running # the script locally must not hand back a shape a deployed run never sees. return _json_round_trip(result) async def sleep(seconds: int): """Server-side sleep — suspend the workflow for the given duration without holding a worker. Inside a @workflow, the parent job suspends and auto-resumes after ``seconds``. Outside a workflow, falls back to ``asyncio.sleep``. """ ctx: WorkflowCtx | None = _workflow_ctx.get(None) if ctx is not None: return await ctx._sleep(seconds) await _asyncio.sleep(seconds) async def wait_for_approval( timeout: int = 1800, form: dict | None = None, self_approval: bool = True, key: str | None = None, skin: Literal["detailed", "minimal"] | None = None, description: str | dict | None = None, ) -> dict: """Suspend the workflow and wait for an external approval. Pass ``key`` to name the step, then ``get_approval_urls(key)`` yields the URLs that resume exactly this approval — route them through your own channel. Without a key the steps are named ``approval``, ``approval_2``, ... Returns a dict with ``value`` (form data), ``approver``, and ``approved``. Args: timeout: Approval timeout in seconds (default 1800). form: Optional form schema for the approval page. self_approval: Whether the user who triggered the flow can approve it (default True). key: Optional checkpoint key naming this approval step. skin: ``"minimal"`` shows approvers only the request (form and approve/reject) instead of the detailed page with the workflow's details. description: Shown to approvers above the form: a string, or a rich value such as ``{"markdown": "..."}``. Example:: urls = await step("urls", lambda: get_approval_urls("manager")) await step("notify", lambda: send_email(urls["resume"], urls["cancel"])) result = await wait_for_approval(key="manager", timeout=3600) """ ctx: WorkflowCtx | None = _workflow_ctx.get(None) if ctx is not None: return await ctx._wait_for_approval( timeout=timeout, form=form, self_approval=self_approval, key=key, skin=skin, description=description, ) raise RuntimeError("wait_for_approval can only be called inside a @workflow") async def parallel(items, fn, *, concurrency: Optional[int] = None): """Process items in parallel with optional concurrency control. Each item is processed by calling ``fn(item)``, which should be a @task. Items are dispatched in batches of ``concurrency`` (default: all at once). Example:: @task async def process(item: str): ... results = await parallel(items, process, concurrency=5) """ if not items: return [] batch_size = concurrency if concurrency and concurrency > 0 else len(items) results = [] for i in range(0, len(items), batch_size): batch = items[i : i + batch_size] batch_results = await _asyncio.gather(*(fn(item) for item in batch)) results.extend(batch_results) return results async def _run_workflow_async(func, checkpoint: dict, input_args: dict): ctx = WorkflowCtx(checkpoint) token = _workflow_ctx.set(ctx) try: result = await func(**input_args) # Flush any unawaited tasks (e.g. forgotten await on last statement) if ctx._pending: steps = list(ctx._pending) ctx._pending.clear() return { "type": "dispatch", "mode": "parallel" if len(steps) > 1 else "sequential", "steps": steps, } return {"type": "complete", "result": result} except _StepFailure as e: # Re-raise the step's own exception so the child job fails with it. raise e.exc except _StepSuspend as e: info = e.dispatch_info mode = info.get("mode") if mode == "step_complete": return {"type": "complete", "result": info.get("result")} if mode == "inline_checkpoint": out = { "type": "inline_checkpoint", "key": info["key"], "result": info.get("result"), } if "started_at" in info: out["started_at"] = info["started_at"] if "duration_ms" in info: out["duration_ms"] = info["duration_ms"] return out if mode == "approval": return { "type": "approval", "key": info["key"], "timeout": info.get("timeout"), "form": info.get("form"), "skin": info.get("skin"), "description": info.get("description"), } if mode == "sleep": return { "type": "sleep", "key": info["key"], "seconds": info.get("seconds"), } return {"type": "dispatch", **info} finally: # Close the lazily-built fast-path httpx client so we don't emit # asyncio ResourceWarning('unclosed transport') on shutdown and don't # leak connection pools when this coroutine is driven from a # long-lived loop (tests, REPL, embedded callers). # # Wrapped in its own try/finally so that asyncio.CancelledError # (which is a BaseException since Python 3.8) during aclose() does # not skip the _workflow_ctx.reset(token) below. try: if ctx._inline_http_client is not None: try: await ctx._inline_http_client.aclose() except Exception: pass ctx._inline_http_client = None finally: _workflow_ctx.reset(token) def _run_workflow(func, checkpoint: dict, input_args: dict): """Synchronous wrapper that runs the workflow coroutine to completion or until it suspends.""" return _asyncio.run(_run_workflow_async(func, checkpoint, input_args)) @init_global_client def commit_kafka_offsets( trigger_path: str, topic: str, partition: int, offset: int, ) -> None: """Commit Kafka offsets for a trigger with auto_commit disabled. Args: trigger_path: Path to the Kafka trigger (from event['wm_trigger']['trigger_path']) topic: Kafka topic name (from event['topic']) partition: Partition number (from event['partition']) offset: Message offset to commit (from event['offset']) """ _client.post( f"/w/{_client.workspace}/kafka_triggers/commit_offsets/{trigger_path}", json={ "topic": topic, "partition": partition, "offset": offset, }, )