From 104c50ab8ae4f61eff572aabc7d5bdd724624f9e Mon Sep 17 00:00:00 2001 From: Diego Imbert Date: Tue, 8 Jul 2025 14:31:49 +0200 Subject: [PATCH] support new syntax in python client --- .../components/copilot/chat/script/core.ts | 8 +-- python-client/wmill/wmill/client.py | 63 +++++++++++++++---- 2 files changed, 55 insertions(+), 16 deletions(-) diff --git a/frontend/src/lib/components/copilot/chat/script/core.ts b/frontend/src/lib/components/copilot/chat/script/core.ts index 4dd7a66694..1bab4df243 100644 --- a/frontend/src/lib/components/copilot/chat/script/core.ts +++ b/frontend/src/lib/components/copilot/chat/script/core.ts @@ -83,8 +83,8 @@ wmill.runScriptAsync(path: string, args?: Record): Promise wmill.waitJob(jobId: string): Promise // Wait for job completion and get result // S3 file operations (if S3 is configured) -wmill.loadS3File(s3object: S3Object): Promise // Load file content from S3 -wmill.writeS3File(s3object: S3Object, content: string | Blob): Promise // Write file to S3 +wmill.loadS3File(s3object: S3Object | string): Promise // Load file content from S3 +wmill.writeS3File(s3object: S3Object | string, content: string | Blob): Promise // Write file to S3 // Flow operations wmill.setFlowUserState(key: string, value: any): Promise // Set flow user state @@ -119,8 +119,8 @@ wmill.run_script_async(path: str, args: dict = None, scheduled_in_secs: int = No wmill.wait_job(job_id: str, timeout = None) -> Any // Wait for job completion and get result // S3 file operations (if S3 is configured) -wmill.load_s3_file(s3object: S3Object, s3_resource_path: str = None) -> bytes // Load file content from S3 -wmill.write_s3_file(s3object: S3Object, file_content: bytes, s3_resource_path: str = None) -> S3Object // Write file to S3 +wmill.load_s3_file(s3object: S3Object | str, s3_resource_path: str = None) -> bytes // Load file content from S3 +wmill.write_s3_file(s3object: S3Object | str, file_content: bytes, s3_resource_path: str = None) -> S3Object // Write file to S3 // Flow operations wmill.run_flow_async(path: str, args: dict = None) -> str // Run flow asynchronously diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index a5f0f72f95..8ab4b420a3 100644 --- a/python-client/wmill/wmill/client.py +++ b/python-client/wmill/wmill/client.py @@ -11,7 +11,8 @@ import time import warnings import json from json import JSONDecodeError -from typing import Dict, Any, Union, Literal +from typing import Dict, Any, Union, Literal, Optional +import re import httpx @@ -312,6 +313,7 @@ class Windmill: return result_text def get_variable(self, path: str) -> str: + path = parse_variable_syntax(path) or path if self.mocked_api is not None: variables = self.mocked_api["variables"] try: @@ -326,6 +328,7 @@ class Windmill: 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: + path = parse_variable_syntax(path) or path if self.mocked_api is not None: self.mocked_api["variables"][path] = value return @@ -358,6 +361,7 @@ class Windmill: path: str, none_if_undefined: bool = False, ) -> dict | None: + path = parse_resource_syntax(path) or path if self.mocked_api is not None: resources = self.mocked_api["resources"] try: @@ -391,6 +395,7 @@ class Windmill: path: str, resource_type: str, ): + path = parse_resource_syntax(path) or path if self.mocked_api is not None: self.mocked_api["resources"][path] = value return @@ -485,6 +490,7 @@ class Windmill: 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", @@ -506,6 +512,7 @@ class Windmill: 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", @@ -527,6 +534,7 @@ class Windmill: 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", @@ -540,7 +548,7 @@ class Windmill: "Could not generate Boto3 S3 connection settings from the provided resource" ) from e - def load_s3_file(self, s3object: S3Object, s3_resource_path: str | None) -> bytes: + 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. @@ -552,11 +560,12 @@ class Windmill: 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, s3_resource_path: str | None + self, s3object: S3Object | str, s3_resource_path: str | None ) -> BufferedReader: """ Load a file from the workspace s3 bucket and returns the bytes stream. @@ -569,6 +578,7 @@ class Windmill: print(file_reader.read()) ''' """ + s3object = parse_s3_object(s3object) reader = S3BufferedReader( f"{self.workspace}", self.client, @@ -580,7 +590,7 @@ class Windmill: def write_s3_file( self, - s3object: S3Object | None, + s3object: S3Object | str | None, file_content: BufferedReader | bytes, s3_resource_path: str | None, content_type: str | None = None, @@ -603,6 +613,7 @@ class Windmill: 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) @@ -644,12 +655,12 @@ class Windmill: raise Exception("Could not write file to S3") from e return S3Object(s3=response["file_key"]) - def sign_s3_objects(self, s3_objects: list[S3Object]) -> list[S3Object]: + def sign_s3_objects(self, s3_objects: list[S3Object | str]) -> list[S3Object]: return self.post( - f"/w/{self.workspace}/apps/sign_s3_objects", json={"s3_objects": s3_objects} + f"/w/{self.workspace}/apps/sign_s3_objects", json={"s3_objects": list(map(parse_s3_object, s3_objects))} ).json() - def sign_s3_object(self, s3_object: S3Object) -> S3Object: + def sign_s3_object(self, s3_object: S3Object | str) -> S3Object: return self.post( f"/w/{self.workspace}/apps/sign_s3_objects", json={"s3_objects": [s3_object]}, @@ -1027,7 +1038,7 @@ def boto3_connection_settings(s3_resource_path: str = "") -> Boto3ConnectionSett @init_global_client -def load_s3_file(s3object: S3Object, s3_resource_path: str | None = None) -> bytes: +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 """ @@ -1038,7 +1049,7 @@ def load_s3_file(s3object: S3Object, s3_resource_path: str | None = None) -> byt @init_global_client def load_s3_file_reader( - s3object: S3Object, s3_resource_path: str | None = None + s3object: S3Object | str, s3_resource_path: str | None = None ) -> BufferedReader: """ Load the content of a file stored in S3 @@ -1050,7 +1061,7 @@ def load_s3_file_reader( @init_global_client def write_s3_file( - s3object: S3Object | None, + s3object: S3Object | str | None, file_content: BufferedReader | bytes, s3_resource_path: str | None = None, content_type: str | None = None, @@ -1075,7 +1086,7 @@ def write_s3_file( @init_global_client -def sign_s3_objects(s3_objects: list[S3Object]) -> list[S3Object]: +def sign_s3_objects(s3_objects: list[S3Object | str]) -> list[S3Object]: """ Sign S3 objects to be used by anonymous users in public apps Returns a list of signed s3 tokens @@ -1084,7 +1095,7 @@ def sign_s3_objects(s3_objects: list[S3Object]) -> list[S3Object]: @init_global_client -def sign_s3_object(s3_object: S3Object) -> S3Object: +def sign_s3_object(s3_object: S3Object| str) -> S3Object: """ Sign S3 object to be used by anonymous users in public apps Returns a signed s3 object @@ -1336,3 +1347,31 @@ def task(*args, **kwargs): return f(args[0], None) else: return lambda x: f(x, kwargs.get("tag")) + +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 string or S3Object format.""" + if isinstance(a, str): + match = re.match(r'^s3://([^/]*)/(.*)$', s3_object) + if match: + return S3Object(s3=match.group(2) or "", storage=match.group(1) or None) + return S3Object(s3="") + else: + return s3_object + + + +def parse_variable_syntax(s: str) -> Optional[str]: + """Parse variable syntax from string.""" + if s.startswith("var://"): + return s[6:] + return None \ No newline at end of file