70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
from io import BufferedReader, BytesIO
|
|
from typing import Optional, Union
|
|
|
|
import httpx
|
|
|
|
|
|
class S3BufferedReader(BufferedReader):
|
|
def __init__(self, workspace: str, windmill_client: httpx.Client, file_key: str, s3_resource_path: Optional[str], storage: Optional[str]):
|
|
params = {
|
|
"file_key": file_key,
|
|
}
|
|
if s3_resource_path is not None:
|
|
params["s3_resource_path"] = s3_resource_path
|
|
if storage is not None:
|
|
params["storage"] = storage
|
|
self._context_manager = windmill_client.stream(
|
|
"GET",
|
|
f"/w/{workspace}/job_helpers/download_s3_file",
|
|
params=params,
|
|
timeout=None,
|
|
)
|
|
|
|
def __enter__(self):
|
|
reader = self._context_manager.__enter__()
|
|
if reader.status_code >= 400:
|
|
error_bytes = reader.read()
|
|
try:
|
|
error_text = error_bytes.decode('utf-8')
|
|
except UnicodeDecodeError:
|
|
error_text = str(error_bytes)
|
|
raise httpx.HTTPStatusError(
|
|
f"Failed to load S3 file: {reader.status_code} {reader.reason_phrase} - {error_text}",
|
|
request=reader.request,
|
|
response=reader
|
|
)
|
|
self._iterator = reader.iter_bytes()
|
|
return self
|
|
|
|
def peek(self, size=0):
|
|
raise Exception("Not implemented, use read() instead")
|
|
|
|
def read(self, size=-1):
|
|
read_result = []
|
|
if size < 0:
|
|
for b in self._iterator:
|
|
read_result.append(b)
|
|
else:
|
|
for i in range(size):
|
|
try:
|
|
b = self._iterator.__next__()
|
|
except StopIteration:
|
|
break
|
|
read_result.append(b)
|
|
|
|
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: Union[BufferedReader, BytesIO]):
|
|
while True:
|
|
byte = buffered_reader.read(50 * 1024)
|
|
if not byte:
|
|
break
|
|
yield byte
|