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

from typing import Any, Union, Dict

import os

from time import sleep

from windmill_api.client import AuthenticatedClient
from windmill_api.api.settings import backend_version
from windmill_api.api.job import run_script_by_hash, get_job, get_completed_job

from windmill_api.api.resource import get_resource as get_resource_api
from windmill_api.api.variable import get_variable as get_variable_api
from windmill_api.models.get_job_response_200_type import GetJobResponse200Type

from windmill_api.models.run_script_by_hash_json_body import RunScriptByHashJsonBody

from enum import Enum

from windmill_api.types import Unset


VAR_RESOURCE_PREFIX = "$var:"


class JobStatus(Enum):
    WAITING = 1
    RUNNING = 2
    COMPLETED = 3

_client: "AuthenticatedClient | None" = None

def create_client(base_url: "str | None" = None, token: "str | None" = None) -> AuthenticatedClient:
    env_base_url = os.environ.get("BASE_INTERNAL_URL")

    if env_base_url is not None:
        env_base_url = env_base_url + "/api"

    base_url_: str = base_url or env_base_url or "http://localhost:8000/api"
    token_ : str = token or os.environ.get("WM_TOKEN") or ""
    global _client
    if _client is None:
        _client = AuthenticatedClient(base_url=base_url_, token=token_, timeout=30)
    return _client

def get_workspace() -> str:
    from_env = os.environ.get("WM_WORKSPACE")
    if from_env is None:
        raise Exception("Workspace not passed as WM_WORKSPACE")
    return from_env

def get_version() -> str:
    """
    Returns the current version of the backend
    """
    return backend_version.sync_detailed(client=create_client()).content.decode("us-ascii")

def run_script_async(
    hash: str,
    args: Dict[str, Any] = {},
    scheduled_in_secs: Union[None, float] = None,
) -> str:
    """
    Launch the run of a script and return immediately its job id
    """
    return run_script_by_hash.sync_detailed(
        client=create_client(),
        workspace=get_workspace(),
        hash_=hash,
        json_body=RunScriptByHashJsonBody.from_dict(args),
        scheduled_in_secs=scheduled_in_secs,
        parent_job=os.environ.get("DT_JOB_ID"),
    ).content.decode("us-ascii")

def run_script_sync(hash: str, args: Dict[str, Any] = {}, verbose: bool = False) -> Dict[str, Any]:
    """
    Run a script, wait for it to complete and return the result of the launched script
    """
    job_id = run_script_async(hash, args, None)
    nb_iter = 0
    while get_job_status(job_id) != JobStatus.COMPLETED:
        if verbose:
            print(f"Waiting for {job_id} to complete...")
        if nb_iter < 10:
            sleep(2.0)
        else:
            sleep(5.0)
        nb_iter += 1
    return get_result(job_id)

def get_job_status(job_id: str) -> JobStatus:
    """
    Returns the status of a queued or completed job
    """
    res = get_job.sync_detailed(client=create_client(), workspace=get_workspace(), id=job_id).parsed
    if not res:
        raise Exception(f"Job {job_id} not found")
    elif not res.type:
        raise Exception(f"Unexpected type not found for job {job_id}")
    elif res.type == GetJobResponse200Type.COMPLETEDJOB:
        return JobStatus.COMPLETED
    else:
        if not "running" in res.additional_properties:
            raise Exception(f"Unexpected running not found for completed job {job_id}")
        elif bool(res.additional_properties["running"]):
            return JobStatus.RUNNING
        else:
            return JobStatus.WAITING

def get_result(job_id: str) -> Dict[str, Any]:
    """
    Returns the result of a completed job
    """
    res = get_completed_job.sync_detailed(client=create_client(), workspace=get_workspace(), id=job_id).parsed
    if not res:
        raise Exception(f"Job {job_id} not found")
    if not res.result:
        raise Exception(f"Unexpected result not found for completed job {job_id}")
    else:
        return res.result.to_dict()  # type: ignore

def get_resource(path: str) -> Dict[str, Any]:
    """
    Returns the resource at a given path as a python dict.
    """
    parsed = get_resource_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client()).parsed
    if parsed is None:
        raise Exception(f"Resource at path {path} does not exist or you do not have read permissions on it")

    if isinstance(parsed.value, Unset):
        return {}

    raw_dict = parsed.value
    res = _transform_leaves(raw_dict)

    return res

def get_variable(path: str) -> str:
    """
    Returns the variable at a given path as a string
    """
    res = get_variable_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client()).parsed
    if res is None:
        raise Exception(f"Variable at path {path} does not exist or you do not have read permissions on it")

    return res.value # type: ignore

def _transform_leaves(d: Dict[str, Any]) -> Dict[str, Any]:
    return {k: _transform_leaf(v) for k, v in d.items()}

def _transform_leaf(v: Any) -> Any:
    if isinstance(v, dict):
        return Client._transform_leaves(v)  # type: ignore
    elif isinstance(v, str):
        if v.startswith(VAR_RESOURCE_PREFIX):
            var_name = v[len(VAR_RESOURCE_PREFIX):]
            return get_variable(var_name)
        else:
            return v
    else:
        return v