From 616088979378712eab5b3abbb646d3e688d2cece Mon Sep 17 00:00:00 2001 From: Guillaume Bouvignies Date: Mon, 29 Jan 2024 14:49:08 +0100 Subject: [PATCH] feat: Download s3 file as stream in Python and TS (#3099) --- backend/windmill-api/src/job_helpers.rs | 2 +- python-client/tests/wmill_client_test.py | 1 + python-client/wmill/wmill/client.py | 60 +++++---- python-client/wmill/wmill/s3_reader.py | 149 +++++++++++++++++++++++ typescript-client/build.sh | 2 +- typescript-client/client.ts | 80 ++++++++---- 6 files changed, 241 insertions(+), 53 deletions(-) create mode 100644 python-client/wmill/wmill/s3_reader.py diff --git a/backend/windmill-api/src/job_helpers.rs b/backend/windmill-api/src/job_helpers.rs index c31f8f4e68..2dceea4d97 100644 --- a/backend/windmill-api/src/job_helpers.rs +++ b/backend/windmill-api/src/job_helpers.rs @@ -1259,7 +1259,7 @@ async fn read_object_chunk( ) -> error::Result> { let s3_object = s3_client .get_object() - .range(format!("bytes={}-{}", from_byte, from_byte + length).to_string()) + .range(format!("bytes={}-{}", from_byte, from_byte + length - 1).to_string()) .bucket(s3_bucket) .key(file_key) .send() diff --git a/python-client/tests/wmill_client_test.py b/python-client/tests/wmill_client_test.py index 1c7806abd0..4b4efc5711 100644 --- a/python-client/tests/wmill_client_test.py +++ b/python-client/tests/wmill_client_test.py @@ -1,5 +1,6 @@ import unittest import wmill +from wmill import S3Object import os diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index c212af21b6..d57176f658 100644 --- a/python-client/wmill/wmill/client.py +++ b/python-client/wmill/wmill/client.py @@ -14,6 +14,7 @@ from typing import Dict, Any, Union, Literal import httpx +from .s3_reader import S3BufferedReader from .s3_types import Boto3ConnectionSettings, DuckDbConnectionSettings, PolarsConnectionSettings, S3Object _client: "Windmill | None" = None @@ -374,6 +375,20 @@ class Windmill: raise Exception("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: + """ + 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") + ''' + """ + return self.load_s3_file_reader(s3object, s3_resource_path).read() + + def load_s3_file_reader(self, s3object: S3Object, s3_resource_path: str | None) -> BufferedReader: """ Load a file from the workspace s3 bucket and returns the bytes stream. @@ -381,35 +396,18 @@ class Windmill: from wmill import S3Object s3_obj = S3Object(s3="/path/to/my_file.txt") - my_obj = client.load_s3_file(s3_obj) - file_content = my_obj["Body"].read().decode("utf-8") + my_obj_content_reader = client.load_s3_file_reader(s3_obj) + file_content = my_obj_content_reader.read().decode("utf-8") ''' """ - part_number = 0 - file_total_size = None - file_content: list[int] = [] - while True: - if part_number is None: - break - try: - part_response = self.post( - f"/w/{self.workspace}/job_helpers/multipart_download_s3_file", - json={ - "file_key": s3object["s3"], - "part_number": part_number, - "file_size": file_total_size, - "s3_resource_path": s3_resource_path, - }, - ).json() - except JSONDecodeError as e: - raise Exception("Could not generate download S3 file part") from e - - if len(part_response["part_content"]) > 0: - file_content = file_content + part_response["part_content"] - part_number = part_response["next_part_number"] - file_total_size = part_response["file_size"] - return bytes(file_content) + result = S3BufferedReader( + workspace=f"{self.workspace}", + windmill_client=self.client, + file_key=s3object["s3"], + s3_resource_path=s3_resource_path, + ) + return result def write_s3_file( self, @@ -729,11 +727,19 @@ def boto3_connection_settings(s3_resource_path: str = "") -> Boto3ConnectionSett @init_global_client def load_s3_file(s3object: S3Object, s3_resource_path: str = "") -> bytes: """ - Load the content of a file stored in S3 + Load the entire content of a file stored in S3 """ 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, s3_resource_path: str = "") -> BufferedReader: + """ + Load the content of a file stored in S3 as a buffered reader + """ + 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 | None, diff --git a/python-client/wmill/wmill/s3_reader.py b/python-client/wmill/wmill/s3_reader.py new file mode 100644 index 0000000000..d14357bea5 --- /dev/null +++ b/python-client/wmill/wmill/s3_reader.py @@ -0,0 +1,149 @@ +from io import BufferedReader +from json import JSONDecodeError + +import httpx + + +class S3BufferedReader(BufferedReader): + def __init__(self, workspace: str, windmill_client: httpx.Client, file_key: str, s3_resource_path: str | None): + self._workspace = workspace + self._client = windmill_client + self._file_key = file_key + self._s3_resource_path = s3_resource_path + self._file_size: int | None = None + + self._part_number: int | None = 0 + self._current_chunk: list[int] = [] + self._position_in_chunk = 0 + + def peek(self, size=0): + read_result = [] + + if size > 0 or ( + len(self._current_chunk) > self._position_in_chunk + and len(self._current_chunk) > self._position_in_chunk + size + ): + payload_to_return = self._current_chunk[self._position_in_chunk : (self._position_in_chunk + size)] + read_result += payload_to_return + return bytes(read_result) + + if self._position_in_chunk < len(self._current_chunk): + payload_to_return = self._current_chunk[self._position_in_chunk :] + read_result += bytes(payload_to_return) + + previous_chunk = self._current_chunk + previous_part_number = self._part_number + previous_position_in_chunk = self._position_in_chunk + try: + while len(read_result) < size or self._part_number is not None: + self._download_new_chunk() + if size > 0 and size - len(read_result) < len(self._current_chunk): + payload_to_return = self._current_chunk[: (size - len(read_result))] + self._position_in_chunk = size - len(read_result) + read_result += bytes(payload_to_return) + break + + read_result += bytes(self._current_chunk) + if self._part_number is None: + break + finally: + # always roll back the changes to the stream state + self._current_chunk = previous_chunk + self._part_number = previous_part_number + self._position_in_chunk = previous_position_in_chunk + return read_result + + def read(self, size=-1): + read_result = [] + + if size > 0 and ( + len(self._current_chunk) > self._position_in_chunk + and len(self._current_chunk) > self._position_in_chunk + size + ): + payload_to_return = self._current_chunk[self._position_in_chunk : (self._position_in_chunk + size)] + self._position_in_chunk += size + read_result += payload_to_return + return bytes(read_result) + + if self._position_in_chunk < len(self._current_chunk): + payload_to_return = self._current_chunk[self._position_in_chunk :] + self._position_in_chunk = len(self._current_chunk) + read_result += payload_to_return + + previous_chunk = self._current_chunk + previous_part_number = self._part_number + previous_position_in_chunk = self._position_in_chunk + try: + while len(read_result) < size or self._part_number is not None: + self._download_new_chunk() + if size > 0 and size - len(read_result) < len(self._current_chunk): + payload_to_return = self._current_chunk[: (size - len(read_result))] + self._position_in_chunk = size - len(read_result) + read_result += payload_to_return + break + + read_result += self._current_chunk + if self._part_number is None: + break + except Exception as e: + # roll back the changes to the stream state + self._current_chunk = previous_chunk + self._part_number = previous_part_number + self._position_in_chunk = previous_position_in_chunk + raise e + return bytes(read_result) + + def read1(self, size=-1): + read_result = [] + + if size < 0: + payload_to_return = self._current_chunk[self._position_in_chunk :] + self._position_in_chunk = len(self._current_chunk) + read_result += payload_to_return + return bytes(read_result) + + if size > 0 and len(self._current_chunk) > self._position_in_chunk: + end_byte = min(self._position_in_chunk + size, len(self._current_chunk)) + payload_to_return = self._current_chunk[self._position_in_chunk : end_byte] + self._position_in_chunk = end_byte + read_result += payload_to_return + return bytes(read_result) + + # no bytes in current buffer, load a new chunk + self._download_new_chunk() + end_byte = min(size, len(self._current_chunk)) + payload_to_return = self._current_chunk[:end_byte] + self._position_in_chunk = end_byte + read_result += payload_to_return + return bytes(read_result) + + def close(self): + self._part_number = 0 + self._current_chunk = [] + self._position_in_chunk = 0 + + def _download_new_chunk( + self, + ): + try: + raw_response = self._client.post( + f"/w/{self._workspace}/job_helpers/multipart_download_s3_file", + json={ + "file_key": self._file_key, + "part_number": self._part_number, + "file_size": self._file_size, + "s3_resource_path": self._s3_resource_path, + }, + ) + try: + raw_response.raise_for_status() + except httpx.HTTPStatusError as err: + raise Exception(f"{err.request.url}: {err.response.status_code}, {err.response.text}") + response = raw_response.json() + except JSONDecodeError as e: + raise Exception("Could not generate download S3 file part") from e + + self._current_chunk = response["part_content"] + self._part_number = response["next_part_number"] + self._file_size = response["file_size"] + self._position_in_chunk = 0 diff --git a/typescript-client/build.sh b/typescript-client/build.sh index 3fda57dd15..45cf304c98 100755 --- a/typescript-client/build.sh +++ b/typescript-client/build.sh @@ -6,7 +6,7 @@ rm -rf "${script_dirpath}/src" npx --yes openapi-typescript-codegen --input "${script_dirpath}/../backend/windmill-api/openapi.yaml" \ --output "${script_dirpath}/src" --useOptions \ -# && sed -i '213 i \\ request.referrerPolicy = \"no-referrer\"\n' src/core/request.ts + && sed -i '213 i \\ request.referrerPolicy = \"no-referrer\"\n' src/core/request.ts cp "${script_dirpath}/client.ts" "${script_dirpath}/src/" cp "${script_dirpath}/s3Types.ts" "${script_dirpath}/src/" diff --git a/typescript-client/client.ts b/typescript-client/client.ts index c98839d5a9..f441767794 100644 --- a/typescript-client/client.ts +++ b/typescript-client/client.ts @@ -1,4 +1,3 @@ -import { Readable } from "stream"; import { ResourceService, VariableService, @@ -298,7 +297,7 @@ export async function denoS3LightClientSettings( * Load the content of a file stored in S3. If the s3ResourcePath is undefined, it will default to the workspace S3 resource. * * ```typescript - * let fileContentStream = await wmill.loadS3File(inputFile) + * let fileContent = await wmill.loadS3File(inputFile) * // if the file is a raw text file, it can be decoded and printed directly: * const text = new TextDecoder().decode(fileContentStream) * console.log(text); @@ -309,6 +308,60 @@ export async function loadS3File( s3ResourcePath: string | undefined ): Promise { !clientSet && setClient(); + const fileContentStream = await loadS3FileStream(s3object, s3ResourcePath) + if (fileContentStream === undefined) { + return undefined + } + + // we read the stream until completion and put the content in an Uint8Array + const reader = fileContentStream.getReader() + const chunks: Uint8Array[] = []; + while (true) { + const {value: chunk, done} = await reader.read(); + if (done) { + break; + } + chunks.push(chunk); + } + let fileContentLength = 0; + chunks.forEach(item => { + fileContentLength += item.length; + }); + let fileContent = new Uint8Array(fileContentLength); + let offset = 0; + chunks.forEach(chunk => { + fileContent.set(chunk, offset); + offset += chunk.length; + }); + return fileContent +} + +/** + * Load the content of a file stored in S3 as a stream. If the s3ResourcePath is undefined, it will default to the workspace S3 resource. + * + * ```typescript + * let fileContentStream = await wmill.loadS3FileStream(inputFile) + * // Use a read to read the stream: + * const chunks: Uint8Array[] = []; + * const reader = fileContentStream.getReader() + * while (true) { + * const {value, done} = await reader.read(); + * if (done) { + * break; + * } + * chunks.push(chunk); + * } + * // then the chunks can be concatenated into a single Uint8Array and if the + * // file is a raw text file, it can be decoded and printed directly: + * const text = new TextDecoder().decode(fileContent) + * console.log(text); + * ``` + */ +export async function loadS3FileStream( + s3object: S3Object, + s3ResourcePath: string | undefined +): Promise { + !clientSet && setClient(); let part_number: number | undefined = 0 let file_total_size: number|undefined = undefined @@ -347,28 +400,7 @@ export async function loadS3File( await fetch(controller) } }) - - // For now we read all the stream in here. In the future return the stream and let the users consume it as they wish - const reader = fileContentStream.getReader() - const chunks: Uint8Array[] = []; - while (true) { - const {value: chunk, done} = await reader.read(); - if (done) { - break; - } - chunks.push(chunk); - } - let fileContentLength = 0; - chunks.forEach(item => { - fileContentLength += item.length; - }); - let fileContent = new Uint8Array(fileContentLength); - let offset = 0; - chunks.forEach(chunk => { - fileContent.set(chunk, offset); - offset += chunk.length; - }); - return fileContent + return fileContentStream } /**