feat: Download s3 file as stream in Python and TS (#3099)

backend/windmill-api/src/job_helpers.rs

@@ -1259,7 +1259,7 @@ async fn read_object_chunk(
 ) -> error::Result<Vec<u8>> {
     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()

python-client/tests/wmill_client_test.py

@@ -1,5 +1,6 @@
 import unittest
 import wmill
+from wmill import S3Object
 import os
 
 

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,

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

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/"

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<Uint8Array|undefined> {
   !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<ReadableStream|undefined> {
+  !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
 }
 
 /**