345 lines
9.3 KiB
Python
345 lines
9.3 KiB
Python
from typing import Any, Union, Dict
|
|
from typing import Generic, TypeVar, TypeAlias
|
|
|
|
import os
|
|
|
|
from time import sleep
|
|
|
|
from windmill_api.client import AuthenticatedClient
|
|
|
|
from enum import Enum
|
|
|
|
from windmill_api.types import Unset
|
|
|
|
S = TypeVar("S")
|
|
|
|
|
|
class Resource(Generic[S]):
|
|
pass
|
|
|
|
|
|
postgresql = TypeAlias
|
|
mysql = TypeAlias
|
|
bigquery = TypeAlias
|
|
|
|
|
|
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, 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: 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
|
|
"""
|
|
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) -> 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 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 | None = None, none_if_undefined: bool = False) -> Any:
|
|
"""
|
|
Returns the resource at a given path
|
|
"""
|
|
from windmill_api.api.resource import get_resource as get_resource_api
|
|
|
|
path = path or get_state_path()
|
|
parsed = get_resource_api.sync_detailed(
|
|
workspace=get_workspace(), path=path, client=create_client()
|
|
).parsed
|
|
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"
|
|
)
|
|
|
|
if isinstance(parsed.value, Unset):
|
|
return None
|
|
|
|
raw = parsed.value
|
|
return _transform_leaf(raw)
|
|
|
|
|
|
def get_state() -> Any:
|
|
"""
|
|
Get the state
|
|
"""
|
|
return get_resource(None, True)
|
|
|
|
|
|
def set_resource(
|
|
value: Any, path: str | None = None, resource_type: str = "state"
|
|
) -> 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.api.resource import (
|
|
exists_resource,
|
|
update_resource,
|
|
create_resource,
|
|
)
|
|
from windmill_api.models.update_resource_json_body import UpdateResourceJsonBody
|
|
|
|
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.sync_detailed(
|
|
workspace=get_workspace(),
|
|
client=client,
|
|
path=path,
|
|
json_body=UpdateResourceJsonBody(value=value),
|
|
)
|
|
|
|
|
|
def set_state(value: Any) -> None:
|
|
"""
|
|
Set the state
|
|
"""
|
|
set_resource(value, None)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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_variable(path: str) -> str:
|
|
"""
|
|
Returns the variable at a given path as a string
|
|
"""
|
|
from windmill_api.api.variable import get_variable as get_variable_api
|
|
|
|
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 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),
|
|
)
|
|
|
|
|
|
def get_state_path() -> str:
|
|
state_path = os.environ.get("WM_STATE_PATH")
|
|
if state_path is None:
|
|
raise Exception("State path not found")
|
|
return state_path
|
|
|
|
|
|
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 _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
|