fix: Python buffered reader (#3136)

* fix: Python buffered reader

* use bytes generator only for buferred reader

python-client/tests/wmill_client_test.py

@@ -93,6 +93,21 @@ SET s3_secret_access_key='80yMndIMcyXwEujxVNINQbf0tBlIzRaLPyM2m1n4';
             file_key = wmill.write_s3_file(S3Object(s3="region.csv"), file_content)
         print(file_key)
 
+    @unittest.skip("skipping")
+    def test_upload_s3_raw_bytes(self):
+        file_key = wmill.write_s3_file(
+            S3Object(s3="hello-world.txt"), b"Hello Windmill!"
+        )
+        print(file_key)
+
+    @unittest.skip("skipping")
+    def test_download_upload_s3_file(self):
+        with wmill.load_s3_file_reader(S3Object(s3="customer.csv")) as file_content:
+            file_key = wmill.write_s3_file(
+                S3Object(s3="customer_test.csv"), file_content
+            )
+        print(file_key)
+
 
 if __name__ == "__main__":
     unittest.main()

python-client/wmill/wmill/client.py

@@ -14,7 +14,7 @@ from typing import Dict, Any, Union, Literal
 
 import httpx
 
-from .s3_reader import S3BufferedReader
+from .s3_reader import S3BufferedReader, bytes_generator
 from .s3_types import Boto3ConnectionSettings, DuckDbConnectionSettings, PolarsConnectionSettings, S3Object
 
 _client: "Windmill | None" = None
@@ -427,11 +427,11 @@ class Windmill:
             client.write_s3_file(s3_obj, my_file)
         '''
         """
-        content_reader: BufferedReader | BytesIO
+        # 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_reader = file_content
+            content_payload = bytes_generator(file_content)
         elif isinstance(file_content, bytes):
-            content_reader = BytesIO(file_content)
+            content_payload = file_content
         else:
             raise Exception("Type of file_content not supported")
 
@@ -447,7 +447,7 @@ class Windmill:
                 f"{self.base_url}/w/{self.workspace}/job_helpers/upload_s3_file",
                 headers={"Authorization": f"Bearer {self.token}", "Content-Type": "application/octet-stream"},
                 params=query_params,
-                content=content_reader,
+                content=content_payload,
                 verify=self.verify,
                 timeout=None,
             ).json()

python-client/wmill/wmill/s3_reader.py

@@ -1,4 +1,4 @@
-from io import BufferedReader
+from io import BufferedReader, BytesIO
 from json import JSONDecodeError
 
 import httpx
@@ -39,10 +39,18 @@ class S3BufferedReader(BufferedReader):
                     break
                 read_result.append(b)
 
-        return b''.join(read_result)
+        return b"".join(read_result)
 
     def read1(self, size=-1):
         return self.read(size)
 
     def __exit__(self, *args):
         self._context_manager.__exit__(*args)
+
+
+def bytes_generator(buffered_reader: BufferedReader | BytesIO):
+    while True:
+        byte = buffered_reader.read(1)
+        if not byte:
+            break
+        yield byte