From 37ffdaed60fb750f1466b440353e1d8409eaea90 Mon Sep 17 00:00:00 2001 From: Stephan Fitzpatrick Date: Wed, 22 Nov 2023 00:49:50 -0800 Subject: [PATCH] feat(python): Refactor Windmill Python client (remove `windmill-api`) (#2665) * "feat(python): Refactor Windmill Python client for better encapsulation and maintainability" This PR obviates the need for the `windmill-api` library. I believe this makes the client package is easier to understand, debug, build, (and test) without it. Additional updates were made to improve logging and add more robust error handling. Here's what Jetbrains' AI assistant came up with to describe the changes based on the diff--I think it did a decent job: > Simplified the Windmill Python client by refactoring out repeated code into more compact, reusable methods. Transitioned the client functions into a Windmill client class, enabling a better encapsulation of the client's state. Updated the README example to reflect this change. This improves code maintainability by making the code easier to understand and update, and improves user experience by providing a more intuitive client interface. * "Refactor post method in Windmill Python client" Removed the hard-coded param 'refresh_client' from the post method in wmill/client.py. since it's no longer used. * "Update build script for Python client" Updated the build script for the Python client for the backend to now include scaffolding code for generating the OpenAPI client, making changes to the generated client, and building the client. . * "Add raise_for_status option in http methods" Enhanced 'get' and 'post' methods in the client class to include a new optional parameter 'raise_for_status'. This allows for better error handling by raising exceptions for 4XX and 5XX responses, if requested. This way, non critical API calls can continue execution even if they receive a client or server error. * "Removed refresh_client condition in post method" * Replace `create_job` with `start_execution` in wmill client This commit changes the method `create_job` to `start_execution` in the wmill client. The change was made to better reflect the function's purpose and make the code even more self-explanatory. Additionally, references to this changed method in README.md and various portions of client.py are also updated. The change will enhance readability and make it easier for newcomers to understand the code. * Fix type hints and default arguments in wmill client This commit adjusts typing hints for several methods from Dict[str, Any] to Any in wmill client since the result of a script isn't always a dictionary. Null arguments are also adjusted from {} to None, ensuring better Python standard practices and less unpredicted behaviors. These changes are aimed to enhance maintainability and make the functions more resilient. * Update client.py * Update pyproject.toml * Update client.py --------- Co-authored-by: Ruben Fiszel --- python-client/wmill/README.md | 21 +- python-client/wmill/pyproject.toml | 2 +- python-client/wmill/wmill/client.py | 946 +++++++++++++++------------- 3 files changed, 531 insertions(+), 438 deletions(-) diff --git a/python-client/wmill/README.md b/python-client/wmill/README.md index f51f1cfdd7..032fffacd2 100644 --- a/python-client/wmill/README.md +++ b/python-client/wmill/README.md @@ -2,9 +2,6 @@ The core client for the [Windmill](https://windmill.dev) platform. -It is a convenient wrapper around the exhaustive, automatically generated from -OpenApi but less user-friendly -[windmill-api](https://pypi.org/project/windmill-api/). ## Quickstart @@ -13,15 +10,15 @@ import wmill def main(): - #os.environ.set("WM_TOKEN", "") OPTIONAL to set token used by the wmill client - version = wmill.get_version() - resource = wmill.get_resource("u/user/resource_path") + client = wmill.Windmill( + # token=... <- this is optional. otherwise the client will look for the WM_TOKEN env var + ) - # run synchronously, will return the result - res = wmill.run_script_sync(hash="000000000000002a", args={}) - print(res) + print(client.version) + print(client.get("u/user/resource_path")) - for _ in range(3): - # run asynchrnously, will return immediately. Can be scheduled - wmill.run_script_async(hash="000000000000002a", args={}, scheduled_in_secs=10) + job_id = client.start_execution(path="path/to/script") + print(job_id) + + return client.run_script(path="path/to/script", args={"arg1": "value1"}) ``` diff --git a/python-client/wmill/pyproject.toml b/python-client/wmill/pyproject.toml index bd36e0e533..aff6e25af6 100644 --- a/python-client/wmill/pyproject.toml +++ b/python-client/wmill/pyproject.toml @@ -16,7 +16,7 @@ include = ["wmill/py.typed"] [tool.poetry.dependencies] python = "^3.7" -windmill-api = "^1.213.0" +httpx = "^0.25" [build-system] requires = ["poetry>=1.0.2", "poetry-dynamic-versioning"] diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index 127cc163c1..081726173a 100644 --- a/python-client/wmill/wmill/client.py +++ b/python-client/wmill/wmill/client.py @@ -1,527 +1,623 @@ -from typing import Any, Union, Dict -from typing import Generic, TypeVar, Optional +from __future__ import annotations -import os -import json -from datetime import timedelta -import logging import atexit +import datetime as dt +import functools +import logging +import os +import random import time +import warnings +from json import JSONDecodeError +from typing import Dict, Any, Union, Literal -from time import sleep -from windmill_api.models.whoami_response_200 import WhoamiResponse200 +import httpx -from windmill_api.client import AuthenticatedClient +_client: "Windmill | None" = None -from enum import Enum +logger = logging.getLogger("windmill_client") -from windmill_api.types import Unset - -S = TypeVar("S") +JobStatus = Literal["RUNNING", "WAITING", "COMPLETED"] -class Resource(Generic[S]): - pass +class Windmill: + def __init__(self, base_url=None, token=None): + base = base_url or os.environ.get("BASE_INTERNAL_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.client = self.get_client() + self.workspace = os.environ.get("WM_WORKSPACE") + self.path = os.environ.get("WM_JOB_PATH") -class JobStatus(Enum): - WAITING = 1 - RUNNING = 2 - COMPLETED = 3 - - -_client: "AuthenticatedClient | None" = None - -logger = logging.getLogger("wmill_client") - - -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, verify_ssl=False) - 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 - """ - from windmill_api.api.settings import backend_version - - 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, int] = None, -) -> str: - """ - Launch the run of a script and return immediately its job id - """ - from windmill_api.api.job import run_script_by_hash - - from windmill_api.models.run_script_by_hash_json_body import RunScriptByHashJsonBody - - 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: Optional[Dict[str, Any]] = None, - verbose: bool = False, - assert_result_is_not_none: bool = True, - cleanup: bool = True, - timeout: Optional[timedelta] = None, -) -> Dict[str, Any]: - """ - Run a script, wait for it to complete and return the result of the launched script - """ - args = args or {} - job_id = run_script_async(hash, args, None) - - def cancel_job(): - from windmill_api.api.job.cancel_queued_job import ( - sync_detailed, - CancelQueuedJobJsonBody, - ) - logger.warning(f"cancelling job {job_id}") - return sync_detailed( - workspace=get_workspace(), - id=job_id, - client=create_client(), - json_body=CancelQueuedJobJsonBody(reason="killed by exit handler"), + def get_client(self) -> httpx.Client: + return httpx.Client( + base_url=self.base_url, + headers=self.headers, ) - if cleanup: - atexit.register(cancel_job) + def get(self, endpoint, raise_for_status=True, **kwargs) -> httpx.Response: + 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 - nb_iter = 0 + def post(self, endpoint, raise_for_status=True, **kwargs) -> httpx.Response: + endpoint = endpoint.lstrip("/") + resp = self.client.post(f"/{endpoint}", **kwargs) + if raise_for_status: + resp.raise_for_status() + return resp - start_time = time.time() - timeout_seconds = timeout.total_seconds() if timeout else None + def create_token(self, duration=dt.timedelta(days=1)) -> str: + 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, + ) -> str: + """Create a job and return its job id.""" + assert not (path and hash_), "path and hash_ are mutually exclusive" + args = args or {} + params = {"scheduled_in_secs": scheduled_in_secs} if scheduled_in_secs else {} + 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( + self, + 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, + ) -> Any: + """Run script synchronously and return its result.""" + args = args or {} - while get_job_status(job_id) != JobStatus.COMPLETED: - if timeout_seconds is not None: - elapsed_time = time.time() - start_time - if elapsed_time > timeout_seconds: - msg = f"Script execution timed out after {timeout_seconds} seconds" - logger.warning(msg) - raise TimeoutError(msg) if verbose: - print(f"Waiting for {job_id} to complete...") - if nb_iter < 10: - sleep(2.0) - else: - sleep(5.0) - nb_iter += 1 - - result = get_result( - job_id, - assert_result_is_not_none=assert_result_is_not_none, - ) + logger.info(f"running `{path}` synchronously with {args = }") - # the job finished--we don't need to cancel it anymore - if cleanup: - atexit.unregister(cancel_job) + if isinstance(timeout, dt.timedelta): + timeout = timeout.total_seconds - error = isinstance(result, dict) and result.get("error") + start_time = time.time() - assert not error, error + job_id = self.start_execution(path=path, hash_=hash_, args=args) - return result + 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) -def run_script_by_path_async( - path: str, - args: Dict[str, Any] = {}, - scheduled_in_secs: Union[None, int] = None, -) -> str: - """ - Launch the run of a script and return immediately its job id - """ - from windmill_api.api.job import run_script_by_path + while True: + job = self.get(f"/w/{self.workspace}/jobs_u/get/{job_id}").json() - from windmill_api.models.run_script_by_path_json_body import RunScriptByPathJsonBody - - return run_script_by_path.sync_detailed( - client=create_client(), - workspace=get_workspace(), - path=path, - json_body=RunScriptByPathJsonBody.from_dict(args), - scheduled_in_secs=scheduled_in_secs, - parent_job=os.environ.get("DT_JOB_ID"), - ).content.decode("us-ascii") - - -def run_script_by_path_sync( - path: str, - args: Dict[str, Any] = {}, - verbose: bool = False, - assert_result_is_not_none: bool = True, - cleanup: bool = True, - timeout: Optional[timedelta] = None, -) -> Dict[str, Any]: - """ - Run a script, wait for it to complete and return the result of the launched script - """ - args = args or {} - job_id = run_script_by_path_async(path, args, None) - - def cancel_job(): - from windmill_api.api.job.cancel_queued_job import ( - sync_detailed, - CancelQueuedJobJsonBody, - ) - logger.warning(f"cancelling job {job_id}") - return sync_detailed( - workspace=get_workspace(), - id=job_id, - client=create_client(), - json_body=CancelQueuedJobJsonBody(reason="killed by exit handler"), - ) - - if cleanup: - atexit.register(cancel_job) - - nb_iter = 0 - - start_time = time.time() - timeout_seconds = timeout.total_seconds() if timeout else None - - while get_job_status(job_id) != JobStatus.COMPLETED: - if timeout_seconds is not None: - elapsed_time = time.time() - start_time - if elapsed_time > timeout_seconds: - msg = f"Script execution timed out after {timeout_seconds} seconds" + 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: - print(f"Waiting for {job_id} to complete...") - if nb_iter < 10: - sleep(2.0) + + result = job.get("result") + canceled, canceled_reason = job.get("canceled"), job.get("canceled_reason") + success = job.get("success") + + if cleanup and (canceled or success): + atexit.unregister(cancel_job) + + if canceled: + raise Exception(f"job canceled: {canceled_reason}") + + if success: + if assert_result_is_not_none and result is None: + raise Exception(f"result is None for {job_id = }") + return result + + if verbose: + logger.info(f"sleeping 0.5 seconds for {job_id = }") + + time.sleep(0.5) + + 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: - sleep(5.0) - nb_iter += 1 - - result = get_result( - job_id, - assert_result_is_not_none=assert_result_is_not_none, - ) + logger.info("no previous executions to cancel") - # the job finished--we don't need to cancel it anymore - if cleanup: - atexit.unregister(cancel_job) + result = {} - error = isinstance(result, dict) and result.get("error") - - assert not error, error - - return result - - -def get_job_status(job_id: str) -> JobStatus: - """ - Returns the status of a queued or completed job - """ - from windmill_api.models.get_job_response_200_type import GetJobResponse200Type - from windmill_api.api.job import get_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, assert_result_is_not_none: bool = True) -> Dict[str, Any]: - """ - Returns the result of a completed job - """ - from windmill_api.api.job import get_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 assert_result_is_not_none and res.result is None: - raise Exception(f"result was null for completed job {job_id}") - else: - return res.result - - -def get_resource(path: Union[str, None] = None, none_if_undefined: bool = False) -> Any: - """ - Returns the resource at a given path - """ - from windmill_api.api.resource import ( - get_resource_value_interpolated as get_resource_api, - ) - - path = path or get_state_path() - parsed = get_resource_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client()) - try: - content = parsed.content.decode("utf-8") - parsed = json.loads(content) - except: - parsed = None - - if parsed is None: - if none_if_undefined: - return None - else: - raise Exception( - f"Resource at path {path} does not exist or you do not have read permissions on it: {content}" + 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 parsed + return result + + def get_job_status(self, job_id: str) -> JobStatus: + resp = self.get(f"/w/{self.workspace}/jobs_u/get/{job_id}", raise_for_status=False) + assert not resp.status_code == 404, f"{job_id} not found" + resp_json = resp.json() + job_type = resp_json.get("type", "") + assert job_type, f"{resp_json} is not a valid job" + if job_type.lower() == "completedjob": + return "COMPLETED" + additional_properties = resp_json.get("additional_properties", {}) + if "running" not in additional_properties: + raise Exception(f"{job_id} is not running") + if additional_properties.get("running"): + return "RUNNING" + return "WAITING" + + def get_result( + self, + job_id: str, + assert_result_is_not_none: bool = True, + ) -> Any: + 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 variable from Windmill""" + return self.get(f"/w/{self.workspace}/variables/get_value/{path}").json() + + def set_variable(self, path: str, value: str) -> None: + """Set variable from Windmill""" + # 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": False, + "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, + ) -> str | None: + """Get resource from Windmill""" + try: + return self.get(f"/w/{self.workspace}/resources/get_value_interpolated/{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, + ): + # 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 set_state(self, value: Any): + self.set_resource(value, path=self.state_path, resource_type="state") + + @property + def version(self): + return self.get("version").text + + def get_duckdb_connection_settings( + self, + s3_resource: Any, + none_if_undefined: bool = False, + ) -> Union[str, None]: + """ + Convenient helpers that takes an S3 resource as input and returns the settings necessary to + initiate an S3 connection from DuckDB + """ + try: + return self.post( + f"/w/{self.workspace}/job_helpers/duckdb_connection_settings", + json={"s3_resource": s3_resource}, + ).json() + except JSONDecodeError as e: + if none_if_undefined: + return None + raise Exception( + "Could not generate DuckDB S3 connection settings from the provided resource" + ) from e + + def get_polars_connection_settings( + self, + s3_resource: Any, + none_if_undefined: bool = False, + ) -> Any: + """ + Convenient helpers that takes an S3 resource as input and returns the settings necessary to + initiate an S3 connection from Polars + """ + try: + return self.post( + f"/w/{self.workspace}/job_helpers/polars_connection_settings", + json={"s3_resource": s3_resource}, + ).json() + except JSONDecodeError as e: + if none_if_undefined: + return None + raise Exception( + "Could not generate Polars S3 connection settings from the provided resource" + ) from e + + def whoami(self) -> dict: + return self.get("/users/whoami").json() + + @property + def user(self) -> dict: + return self.whoami() + + @property + def state_path(self) -> str: + 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: + return self.get_resource(path=self.state_path, none_if_undefined=True) + + @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: + """ + Set 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) -> dict: + nonce = random.randint(0, 1000000000) + return self.get( + f"/w/{self.workspace}/jobs_u/get_resume_urls/{self.path}/{nonce}", + params={"approver": approver}, + ).json() -def duckdb_connection_settings(s3_resource: Any, none_if_undefined: bool = False) -> Union[str, None]: +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: + return _client.workspace + + +@init_global_client +@deprecate("Windmill().version") +def get_version() -> str: + return _client.version + + +@init_global_client +def run_script_async( + hash: str, + args: Dict[str, Any] = None, + scheduled_in_secs: int = None, +) -> str: + return _client.start_execution( + hash_=hash, + args=args, + scheduled_in_secs=scheduled_in_secs, + ) + + +@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, +) -> Any: + return _client.run_script( + hash_=hash, + args=args, + verbose=verbose, + assert_result_is_not_none=assert_result_is_not_none, + cleanup=cleanup, + timeout=timeout, + ) + + +@init_global_client +def run_script_by_path_async( + path: str, + args: Dict[str, Any] = None, + scheduled_in_secs: Union[None, int] = None, +) -> str: + return _client.start_execution( + path=path, + args=args, + scheduled_in_secs=scheduled_in_secs, + ) + + +@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, +) -> Any: + return _client.run_script( + path=path, + args=args, + verbose=verbose, + assert_result_is_not_none=assert_result_is_not_none, + cleanup=cleanup, + timeout=timeout, + ) + + +@init_global_client +def get_job_status(job_id: str) -> JobStatus: + return _client.get_job_status(job_id) + + +@init_global_client +def get_result(*args, **kwargs) -> Dict[str, Any]: + return _client.get_result(*args, **kwargs) + + +@init_global_client +def duckdb_connection_settings(*args, **kwargs) -> Union[str, None]: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection from DuckDB """ - from windmill_api.api.helpers import duckdb_connection_settings - from windmill_api.models.s3_resource import S3Resource - - parsed = duckdb_connection_settings.sync_detailed( - workspace=get_workspace(), - client=create_client(), - json_body={ - "s3_resource": s3_resource, - }, - ) - try: - content = parsed.content.decode("utf-8") - parsed = json.loads(content) - except: - parsed = None - - if parsed is None: - if none_if_undefined: - return None - else: - raise Exception(f"Could not generate DuckDB S3 connection settings from the provided resource") - return parsed + return _client.get_duckdb_connection_settings(*args, **kwargs) -def polars_connection_settings(s3_resource: Any, none_if_undefined: bool = False) -> Any: +@init_global_client +def polars_connection_settings(*args, **kwargs) -> Any: """ Convenient helpers that takes an S3 resource as input and returns the settings necessary to initiate an S3 connection from Polars """ - from windmill_api.api.helpers import polars_connection_settings - - parsed = polars_connection_settings.sync_detailed( - workspace=get_workspace(), - client=create_client(), - json_body={ - "s3_resource": s3_resource, - }, - ) - try: - content = parsed.content.decode("utf-8") - parsed = json.loads(content) - except: - parsed = None - - if parsed is None: - if none_if_undefined: - return None - else: - raise Exception(f"Could not generate Polars S3 connection settings from the provided resource") - return parsed + return _client.get_polars_connection_settings(*args, **kwargs) -def whoami() -> Union[WhoamiResponse200, None]: +@init_global_client +def whoami() -> dict: """ Returns the current user """ - from windmill_api.api.user import whoami - - return whoami.sync(client=create_client(), workspace=get_workspace()) + return _client.user +@init_global_client +@deprecate("Windmill().state") def get_state() -> Any: """ Get the state """ - return get_resource(None, True) + return _client.state -def set_resource(value: Any, path: Union[str, None] = None, resource_type: str = "state") -> None: +@init_global_client +def get_resource( + path: str, + none_if_undefined: bool = False, +) -> str | None: + """Get resource from Windmill""" + return _client.get_resource(path, none_if_undefined) + +@init_global_client +def set_resource(**kwargs) -> None: """ Set the resource at a given path as a string, creating it if it does not exist """ - from windmill_api.models.create_resource_json_body import CreateResourceJsonBody - from windmill_api.models.update_resource_value_json_body import ( - UpdateResourceValueJsonBody, - ) - from windmill_api.api.resource import ( - exists_resource, - update_resource_value, - create_resource, - ) - - path = path or get_state_path() - workspace = get_workspace() - client = create_client() - if not exists_resource.sync_detailed(workspace=workspace, path=path, client=client).parsed: - create_resource.sync_detailed( - workspace=workspace, - client=client, - json_body=CreateResourceJsonBody(path=path, value=value, resource_type=resource_type), - ) - else: - update_resource_value.sync_detailed( - workspace=get_workspace(), - client=client, - path=path, - json_body=UpdateResourceValueJsonBody(value=value), - ) + return _client.set_resource(**kwargs) +@init_global_client def set_state(value: Any) -> None: """ Set the state """ - set_resource(value, None) + return _client.set_state(value) -def set_shared_state_pickle(value: Any, path: str = "state.pickle") -> None: +def set_shared_state_pickle(**kwargs) -> 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) + return Windmill.set_shared_state_pickle(**kwargs) -def get_shared_state_pickle(path: str = "state.pickle") -> Any: +@deprecate("Windmill.get_shared_state_pickle(...)") +def get_shared_state_pickle(**kwargs) -> Any: """ Get the state in the shared folder using pickle """ - import pickle - - with open(f"/shared/{path}", "rb") as handle: - return pickle.load(handle) + return Windmill.get_shared_state_pickle(**kwargs) -def set_shared_state(value: Any, path: str = "state.json") -> None: +def set_shared_state(**kwargs) -> 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) + return Windmill.set_shared_state(**kwargs) -def get_shared_state(path: str = "state.json") -> None: +def get_shared_state(**kwargs) -> None: """ Set 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) + return Windmill.get_shared_state(**kwargs) +@init_global_client def get_variable(path: str) -> str: """ Returns the variable at a given path as a string """ - from windmill_api.api.variable import get_variable_value as get_variable_api - - res = get_variable_api.sync_detailed(workspace=get_workspace(), path=path, client=create_client()) - parsed = res.parsed - if parsed is None: - raise Exception( - f"Variable at path {path} does not exist or you do not have read permissions on it: {res.content.decode('utf-8')}" - ) - - return parsed + return _client.get_variable(path) +@init_global_client def set_variable(path: str, value: str) -> None: """ Set the variable at a given path as a string, creating it if it does not exist """ - from windmill_api.api.variable import ( - exists_variable, - update_variable, - create_variable, - ) - - from windmill_api.models.update_variable_json_body import UpdateVariableJsonBody - from windmill_api.models.create_variable_json_body import CreateVariableJsonBody - - workspace = get_workspace() - client = create_client() - if not exists_variable.sync_detailed(workspace=workspace, path=path, client=client).parsed: - create_variable.sync_detailed( - workspace=workspace, - client=client, - json_body=CreateVariableJsonBody(path=path, value=value, is_secret=False, description=""), - ) - else: - update_variable.sync_detailed( - workspace=get_workspace(), - path=path, - client=client, - json_body=UpdateVariableJsonBody(value=value), - ) + return _client.set_variable(path, value) +@init_global_client def get_state_path() -> str: - state_path = os.environ.get("WM_STATE_PATH_NEW") or os.environ.get("WM_STATE_PATH") - if state_path is None: - raise Exception("State path not found") - return state_path + return _client.state_path -def get_resume_urls(approver: Union[str, None] = None) -> Dict: - from windmill_api.api.job import get_resume_urls as get_resume_urls_api - - workspace = get_workspace() - client = create_client() - job_id = os.environ.get("WM_JOB_ID") or "NO_ID" - import random - - nonce = random.randint(0, 1000000000) - res = get_resume_urls_api.sync_detailed(workspace, job_id, nonce, client=client, approver=approver) - if res.parsed is not None: - return res.parsed.to_dict() - else: - raise Exception("Failed to get resume urls") +@init_global_client +def get_resume_urls(approver: str = None) -> dict: + return _client.get_resume_urls(approver)