feat: migrate s3 client to object_store (#3116)

* feat: migrate s3 client to object_store

* remove multipart from open API

* Remove multipart in favor of a single stream

* progress report

* add progress reader on server side

* small nit fix

* fix read chunk

* Fix TS and python SDK

* Fix download button

* Fix download button object viewer

* fix list

* Better errors

* export loadS3FileContent

* revert changes SDK

* fix browser

* small file list unavailable fix

* Old endpoints throws informative error messages

* Typescript SDK uses raw fetch

* update python SDK

* Error if uploaded file > 50Mb

* revert python SDL changes

* Update python SDK method docs

---------

Co-authored-by: Ruben Fiszel <ruben@windmill.dev>
Co-authored-by: Ruben Fiszel <ruben@rubenfiszel.com>

python-client/tests/wmill_client_test.py

@@ -15,6 +15,7 @@ class TestStringMethods(unittest.TestCase):
         os.environ["WM_TOKEN"] = self._token
         os.environ["BASE_INTERNAL_URL"] = self._host
 
+    @unittest.skip("skipping")
     def test_duckdb_connection_settings(self):
         settings = wmill.duckdb_connection_settings(self._resource_path)
         self.assertIsNotNone(settings)
@@ -35,6 +36,7 @@ SET s3_secret_access_key='80yMndIMcyXwEujxVNINQbf0tBlIzRaLPyM2m1n4';
         settings = wmill.polars_connection_settings(self._resource_path)
         print(settings)
 
+    @unittest.skip("skipping")
     def test_polars_connection_settings(self):
         settings = wmill.polars_connection_settings(self._resource_path)
         s3fs_args_expected = {
@@ -59,6 +61,7 @@ SET s3_secret_access_key='80yMndIMcyXwEujxVNINQbf0tBlIzRaLPyM2m1n4';
         )
         self.assertEqual(settings.polars_cloud_options, polars_cloud_options_expected)
 
+    @unittest.skip("skipping")
     def test_boto3_connection_settings(self):
         settings = wmill.boto3_connection_settings(self._resource_path)
         expected_settings = {
@@ -72,6 +75,24 @@ SET s3_secret_access_key='80yMndIMcyXwEujxVNINQbf0tBlIzRaLPyM2m1n4';
         self.assertEqual(settings["endpoint_url"], "http://localhost:9000")
         self.assertEqual(settings.endpoint_url, "http://localhost:9000")
 
+    @unittest.skip("skipping")
+    def test_download_s3_file(self):
+        with wmill.load_s3_file_reader(S3Object(s3="region.csv")) as file_content, open(
+            "region.csv", "wb"
+        ) as output_file:
+            output_file.write(file_content.read())
+
+    @unittest.skip("skipping")
+    def test_download_s3_file_content(self):
+        file_content = wmill.load_s3_file(S3Object(s3="region.csv"))
+        print(file_content)
+
+    @unittest.skip("skipping")
+    def test_upload_s3_file(self):
+        with open("region.csv", "rb") as file_content:
+            file_key = wmill.write_s3_file(S3Object(s3="region.csv"), file_content)
+        print(file_key)
+
 
 if __name__ == "__main__":
     unittest.main()

python-client/wmill/wmill/client.py

@@ -386,7 +386,8 @@ class Windmill:
         file_content = my_obj_content.decode("utf-8")
         '''
         """
-        return self.load_s3_file_reader(s3object, s3_resource_path).read()
+        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) -> BufferedReader:
         """
@@ -396,24 +397,17 @@ class Windmill:
         from wmill import S3Object
 
         s3_obj = S3Object(s3="/path/to/my_file.txt")
-        my_obj_content_reader = client.load_s3_file_reader(s3_obj)
-        file_content = my_obj_content_reader.read().decode("utf-8")
+        with wmill.load_s3_file(s3object, s3_resource_path) as file_reader:
+            print(file_reader.read())
         '''
         """
-
-        result = S3BufferedReader(
-            workspace=f"{self.workspace}",
-            windmill_client=self.client,
-            file_key=s3object["s3"],
-            s3_resource_path=s3_resource_path,
-        )
-        return result
+        reader = S3BufferedReader(f"{self.workspace}", self.client, s3object["s3"], s3_resource_path)
+        return reader
 
     def write_s3_file(
         self,
         s3object: S3Object | None,
         file_content: BufferedReader | bytes,
-        file_expiration: dt.datetime | None,
         s3_resource_path: str | None,
     ) -> S3Object:
         """
@@ -441,38 +435,25 @@ class Windmill:
         else:
             raise Exception("Type of file_content not supported")
 
-        file_key = s3object["s3"] if s3object is not None else None
-        parts = []
-        upload_id = None
-        chunk = content_reader.read(5 * 1024 * 1024)
-        if len(chunk) == 0:
-            raise Exception("File content is empty, nothing to upload")
-        while True:
-            chunk_2 = content_reader.read(5 * 1024 * 1024)
-            reader_done = len(chunk_2) == 0
-            try:
-                response = self.post(
-                    f"/w/{self.workspace}/job_helpers/multipart_upload_s3_file",
-                    json={
-                        "file_key": file_key,
-                        "part_content": [b for b in chunk],
-                        "upload_id": upload_id,
-                        "parts": parts,
-                        "is_final": reader_done,
-                        "cancel_upload": False,
-                        "s3_resource_path": s3_resource_path,
-                        "file_expiration": file_expiration.isoformat() if file_expiration else None,
-                    },
-                ).json()
-            except Exception as e:
-                raise Exception("Could not write file to S3") from e
-            parts = response["parts"]
-            upload_id = response["upload_id"]
-            file_key = response["file_key"]
-            if response["is_done"]:
-                break
-            chunk = chunk_2
-        return S3Object(s3=file_key)
+        query_params = {}
+        if s3object is not None and s3object["s3"] != "":
+            query_params["file_key"] = s3object["s3"]
+        if s3_resource_path is not None and s3_resource_path != "":
+            query_params["s3_resource_path"] = s3_resource_path
+
+        try:
+            # need a vanilla client b/c content-type is not application/json here
+            response = httpx.post(
+                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,
+                verify=self.verify,
+                timeout=None,
+            ).json()
+        except Exception as e:
+            raise Exception("Could not write file to S3") from e
+        return S3Object(s3=response["file_key"])
 
     def __boto3_connection_settings(self, s3_resource) -> Boto3ConnectionSettings:
         endpoint_url_prefix = "https://" if s3_resource["useSSL"] else "http://"
@@ -727,7 +708,7 @@ 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 entire content of a file stored in S3
+    Load the entire content of a file stored in S3 as bytes
     """
     return _client.load_s3_file(s3object, s3_resource_path if s3_resource_path != "" else None)
 
@@ -735,7 +716,7 @@ def load_s3_file(s3object: S3Object, s3_resource_path: str = "") -> bytes:
 @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
+    Load the content of a file stored in S3
     """
     return _client.load_s3_file_reader(s3object, s3_resource_path if s3_resource_path != "" else None)
 
@@ -744,15 +725,12 @@ def load_s3_file_reader(s3object: S3Object, s3_resource_path: str = "") -> Buffe
 def write_s3_file(
     s3object: S3Object | None,
     file_content: BufferedReader | bytes,
-    file_expiration: dt.datetime | None = None,
     s3_resource_path: str = "",
 ) -> S3Object:
     """
     Upload a file to S3
     """
-    return _client.write_s3_file(
-        s3object, file_content, file_expiration, s3_resource_path if s3_resource_path != "" else None
-    )
+    return _client.write_s3_file(s3object, file_content, s3_resource_path if s3_resource_path != "" else None)
 
 
 @init_global_client

python-client/wmill/wmill/s3_reader.py

@@ -6,144 +6,43 @@ 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
+        params = {
+            "file_key": file_key,
+        }
+        if s3_resource_path is not None:
+            params["s3_resource_path"] = s3_resource_path
+        self._context_manager = windmill_client.stream(
+            "GET",
+            f"/w/{workspace}/job_helpers/download_s3_file",
+            params=params,
+            timeout=None,
+        )
 
-        self._part_number: int | None = 0
-        self._current_chunk: list[int] = []
-        self._position_in_chunk = 0
+    def __enter__(self):
+        reader = self._context_manager.__enter__()
+        self._iterator = reader.iter_bytes()
+        return self
 
     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
+        raise Exception("Not implemented, use read() instead")
 
     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
+        if size < 0:
+            for b in self._iterator:
+                read_result += b
+        else:
+            for i in range(size):
+                try:
+                    b = self._iterator.__next__()
+                except StopIteration:
                     break
+                read_result += b
 
-                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 = []
+        return self.read(size)
 
-        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
+    def __exit__(self, *args):
+        self._context_manager.__exit__(*args)