windmill/python-client/wmill/wmill/client.py

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", *, role: Optional[str] = None):
        """Get a DataTable client for SQL queries.

        Args:
            name: Database name (default: "main")
            role: Connect as this data table role instead of the data table's default one.
                Only meaningful on a data table under roles, and only for a role you are a
                tenant of.

        Returns:
            DataTableClient instance
        """
        return DataTableClient(self, name, role=role)

    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", *, role: Optional[str] = None) -> DataTableClient:
    """Get a DataTable client for SQL queries.

    Args:
        name: Database name (default: "main")
        role: Connect as this data table role instead of the data table's default one.

    Returns:
        DataTableClient instance
    """
    return _client.datatable(name, role=role)

@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://<storage>/<key>` URI string (`s3:///<key>`
    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://<storage>/<key> with a non-empty key "
                "(s3:///<key> for the default storage)"
            )
        raise ValueError(
            f"Invalid s3 object {s3_object!r}: expected an s3://<storage>/<key> "
            f"URI (e.g. 's3:///{s3_object}' for key {s3_object!r} in the default "
            "storage) or S3Object(s3=<key>)"
        )
    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)

# Interpolated into a `-- role <name>` line, so a value carrying a newline could append
# statements of its own. Mirrors the server's own role-name rule.
_ROLE_NAME_RE = re.compile(r"^[A-Za-z0-9_-]{1,63}$")


class DataTableClient:
    """Client for executing SQL queries against Windmill DataTables."""

    def __init__(self, client: Windmill, name: str, role: Optional[str] = None):
        """Initialize DataTableClient.

        Args:
            client: Windmill client instance
            name: DataTable name
            role: Data table role to connect as, or None for the data table's default
        """
        if role is not None and not _ROLE_NAME_RE.match(role):
            raise ValueError(
                f"Invalid data table role '{role}': only letters, digits, '_' and '-' are allowed"
            )
        self.client = client
        self.role = role
        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
        # Must lead: the executor's annotation parser stops at the first non-comment line.
        if self.role is not None:
            sql = f"-- role {self.role}\n" + 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: <lake>/<schema>.<table>
        # for an explicit schema, else <lake>/<table>. 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"<unrepresentable {type(o).__name__}>"


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"):
                    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 ``<key>_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 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. A task is keyed on its step key (its name and
    call order) and the workflow's input, not on the arguments it is called
    with, so cache one only when whether it runs, and what it receives, follow
    from the workflow's input alone. A ``task_script`` target is keyed on the
    arguments it is called with. It has no effect on a ``task_flow`` target,
    which keeps its flow's own cache policy.

    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__

        _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=_task_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,
        },
    )