diff --git a/python-client/tests/wmill_client_test.py b/python-client/tests/wmill_client_test.py index 88ea15294d..1acbcb0bf7 100644 --- a/python-client/tests/wmill_client_test.py +++ b/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() diff --git a/python-client/wmill/wmill/client.py b/python-client/wmill/wmill/client.py index fe09747794..a193037f23 100644 --- a/python-client/wmill/wmill/client.py +++ b/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() diff --git a/python-client/wmill/wmill/s3_reader.py b/python-client/wmill/wmill/s3_reader.py index 69b0a8a8b6..9afcebcdbf 100644 --- a/python-client/wmill/wmill/s3_reader.py +++ b/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