"""Amazon S3 filesystem abstraction layer (PRIVATE).""" from __future__ import annotations import concurrent.futures import io import itertools import logging import math import socket from contextlib import contextmanager from errno import ESPIPE from types import TracebackType from typing import TYPE_CHECKING, Any, BinaryIO, Iterator, cast import boto3 from botocore.exceptions import ReadTimeoutError from botocore.loaders import Loader from botocore.model import ServiceModel from awswrangler import _utils, exceptions from awswrangler._config import apply_configs from awswrangler._distributed import engine from awswrangler.s3._describe import _describe_object if TYPE_CHECKING: from mypy_boto3_s3 import S3Client _logger: logging.Logger = logging.getLogger(__name__) _S3_RETRYABLE_ERRORS: tuple[Any, Any, Any] = (socket.timeout, ConnectionError, ReadTimeoutError) _MIN_WRITE_BLOCK: int = 5_242_880 # 5 MB (5 * 2**20) _MIN_PARALLEL_READ_BLOCK: int = 5_242_880 # 5 MB (5 * 2**20) _BOTOCORE_LOADER = Loader() _S3_JSON_MODEL = _BOTOCORE_LOADER.load_service_model(service_name="s3", type_name="service-2") _S3_SERVICE_MODEL = ServiceModel(_S3_JSON_MODEL, service_name="s3") def _snake_to_camel_case(s: str) -> str: return "".join(c.title() for c in s.split("_")) def get_botocore_valid_kwargs(function_name: str, s3_additional_kwargs: dict[str, Any]) -> dict[str, Any]: """Filter and keep only the valid botocore key arguments.""" s3_operation_model = _S3_SERVICE_MODEL.operation_model(_snake_to_camel_case(function_name)) allowed_kwargs = s3_operation_model.input_shape.members.keys() # type: ignore[union-attr] return {k: v for k, v in s3_additional_kwargs.items() if k in allowed_kwargs} def _fetch_range( range_values: tuple[int, int], bucket: str, key: str, s3_client: "S3Client", boto3_kwargs: dict[str, Any], version_id: str | None = None, ) -> tuple[int, bytes]: start, end = range_values _logger.debug("Fetching: s3://%s/%s - VersionId: %s - Range: %s-%s", bucket, key, version_id, start, end) if version_id: boto3_kwargs["VersionId"] = version_id resp = _utils.try_it( f=s3_client.get_object, ex=_S3_RETRYABLE_ERRORS, base=0.5, max_num_tries=6, Bucket=bucket, Key=key, Range=f"bytes={start}-{end - 1}", **boto3_kwargs, ) return start, resp["Body"].read() class _UploadProxy: def __init__(self, use_threads: bool | int): self.closed = False self._exec: concurrent.futures.ThreadPoolExecutor | None self._results: list[dict[str, str | int]] = [] self._cpus: int = _utils.ensure_cpu_count(use_threads=use_threads) if self._cpus > 1: self._exec = concurrent.futures.ThreadPoolExecutor(max_workers=self._cpus) self._futures: list[Any] = [] else: self._exec = None @staticmethod def _sort_by_part_number(parts: list[dict[str, str | int]]) -> list[dict[str, str | int]]: return sorted(parts, key=lambda k: k["PartNumber"]) @staticmethod def _caller( bucket: str, key: str, part: int, upload_id: str, data: bytes, s3_client: "S3Client", boto3_kwargs: dict[str, Any], ) -> dict[str, str | int]: _logger.debug("Upload part %s started.", part) resp = _utils.try_it( f=s3_client.upload_part, ex=_S3_RETRYABLE_ERRORS, base=0.5, max_num_tries=6, Bucket=bucket, Key=key, Body=data, PartNumber=part, UploadId=upload_id, **boto3_kwargs, ) _logger.debug("Upload part %s done.", part) return {"PartNumber": part, "ETag": resp["ETag"]} def upload( self, bucket: str, key: str, part: int, upload_id: str, data: bytes, s3_client: "S3Client", boto3_kwargs: dict[str, Any], ) -> None: """Upload Part.""" if self._exec is not None: _utils.block_waiting_available_thread(seq=self._futures, max_workers=self._cpus) future = self._exec.submit( _UploadProxy._caller, bucket=bucket, key=key, part=part, upload_id=upload_id, data=data, s3_client=s3_client, boto3_kwargs=boto3_kwargs, ) self._futures.append(future) else: self._results.append( self._caller( bucket=bucket, key=key, part=part, upload_id=upload_id, data=data, s3_client=s3_client, boto3_kwargs=boto3_kwargs, ) ) def close(self) -> list[dict[str, str | int]]: """Close the proxy.""" if self.closed is True: return [] if self._exec is not None: try: for future in concurrent.futures.as_completed(self._futures): self._results.append(future.result()) finally: self._exec.shutdown(wait=True) self.closed = True return self._sort_by_part_number(parts=self._results) class _S3ObjectBase(io.RawIOBase): """Class to abstract S3 objects as ordinary files.""" def __init__( self, path: str, s3_block_size: int, mode: str, use_threads: bool | int, s3_client: "S3Client", s3_additional_kwargs: dict[str, str] | None, newline: str | None, encoding: str | None, version_id: str | None = None, ) -> None: super().__init__() self._use_threads = use_threads self._newline: str = "\n" if newline is None else newline self._encoding: str = "utf-8" if encoding is None else encoding self._bucket, self._key = _utils.parse_path(path=path) self._version_id = version_id if mode not in {"rb", "wb", "r", "w"}: raise NotImplementedError(f"File mode must be {'rb', 'wb', 'r', 'w'}, not {mode}") self._mode: str = "rb" if mode is None else mode self._one_shot_download: bool = False if 0 < s3_block_size < 3: raise exceptions.InvalidArgumentValue( "s3_block_size MUST > 2 to define a valid size or " "< 1 to avoid blocks and always execute one shot downloads." ) if s3_block_size <= 0: _logger.debug("s3_block_size of %d, enabling one_shot_download.", s3_block_size) self._one_shot_download = True self._s3_block_size: int = s3_block_size self._s3_half_block_size: int = s3_block_size // 2 self._s3_additional_kwargs: dict[str, str] = {} if s3_additional_kwargs is None else s3_additional_kwargs self._client: "S3Client" = s3_client self._loc: int = 0 if self.readable() is True: self._cache: bytes = b"" self._start: int = 0 self._end: int = 0 func = engine.dispatch_func(_describe_object, "python") _, desc = func( s3_client=self._client, path=path, s3_additional_kwargs=self._s3_additional_kwargs, version_id=version_id, ) size: int | None = desc.get("ContentLength", None) if size is None: raise exceptions.InvalidArgumentValue(f"S3 object w/o defined size: {path}") self._size: int = size _logger.debug("self._size: %s", self._size) _logger.debug("self._s3_block_size: %s", self._s3_block_size) elif self.writable() is True: self._mpu: dict[str, Any] = {} self._buffer: io.BytesIO = io.BytesIO() self._parts_count: int = 0 self._size = 0 self._upload_proxy: _UploadProxy = _UploadProxy(use_threads=self._use_threads) else: raise RuntimeError(f"Invalid mode: {self._mode}") def __enter__(self) -> "_S3ObjectBase": return self def __exit__( self, exception_type: type[BaseException] | None, exception_value: BaseException | None, traceback: TracebackType | None, ) -> None: """Close the context.""" _logger.debug("exc_type: %s", exception_type) _logger.debug("exc_value: %s", exception_value) _logger.debug("exc_traceback: %s", traceback) self.close() def __del__(self) -> None: """Delete object tear down.""" self.close() def __next__(self) -> bytes: """Next line.""" out: bytes | None = self.readline() if not out: raise StopIteration return out next = __next__ def __iter__(self) -> "_S3ObjectBase": """Iterate over lines.""" return self @staticmethod def _merge_range(ranges: list[tuple[int, bytes]]) -> bytes: return b"".join(data for start, data in sorted(ranges, key=lambda r: r[0])) def _fetch_range_proxy(self, start: int, end: int) -> bytes: _logger.debug("Fetching: s3://%s/%s - Range: %s-%s", self._bucket, self._key, start, end) boto3_kwargs: dict[str, Any] = get_botocore_valid_kwargs( function_name="get_object", s3_additional_kwargs=self._s3_additional_kwargs ) cpus: int = _utils.ensure_cpu_count(use_threads=self._use_threads) range_size: int = end - start if cpus < 2 or range_size < (2 * _MIN_PARALLEL_READ_BLOCK): return _fetch_range( range_values=(start, end), bucket=self._bucket, key=self._key, s3_client=self._client, boto3_kwargs=boto3_kwargs, version_id=self._version_id, )[1] sizes: tuple[int, ...] = _utils.get_even_chunks_sizes( total_size=range_size, chunk_size=_MIN_PARALLEL_READ_BLOCK, upper_bound=False ) ranges: list[tuple[int, int]] = [] chunk_start: int = start for size in sizes: ranges.append((chunk_start, chunk_start + size)) chunk_start += size with concurrent.futures.ThreadPoolExecutor(max_workers=cpus) as executor: return self._merge_range( ranges=list( executor.map( _fetch_range, ranges, itertools.repeat(self._bucket), itertools.repeat(self._key), itertools.repeat(self._client), itertools.repeat(boto3_kwargs), itertools.repeat(self._version_id), ) ), ) def _fetch(self, start: int, end: int) -> None: if end > self._size: raise ValueError(f"Trying to fetch byte (at position {end - 1}) beyond file size ({self._size})") if start < 0: raise ValueError(f"Trying to fetch byte (at position {start}) beyond file range ({self._size})") if start >= self._start and end <= self._end: return None # Does not require download if self._one_shot_download: self._start = 0 self._end = self._size self._cache = self._fetch_range_proxy(self._start, self._end) return None if end - start >= self._s3_block_size: # Fetching length greater than cache length self._cache = self._fetch_range_proxy(start, end) self._start = start self._end = end return None # Calculating block START and END positions _logger.debug("Downloading: %s (start) / %s (end)", start, end) mid: int = int(math.ceil((start + (end - 1)) / 2)) new_block_start: int = mid - self._s3_half_block_size new_block_start = new_block_start + 1 if self._s3_block_size % 2 == 0 else new_block_start new_block_end: int = mid + self._s3_half_block_size + 1 _logger.debug("new_block_start: %s / new_block_end: %s / mid: %s", new_block_start, new_block_end, mid) if new_block_start < 0 and new_block_end > self._size: # both ends overflowing new_block_start = 0 new_block_end = self._size elif new_block_end > self._size: # right overflow new_block_start = new_block_start - (new_block_end - self._size) new_block_start = 0 if new_block_start < 0 else new_block_start new_block_end = self._size elif new_block_start < 0: # left overflow new_block_end = new_block_end - new_block_start new_block_end = self._size if new_block_end > self._size else new_block_end new_block_start = 0 _logger.debug( "new_block_start: %s / new_block_end: %s/ self._start: %s / self._end: %s", new_block_start, new_block_end, self._start, self._end, ) # Calculating missing bytes in cache if ( # Full block download (new_block_start < self._start and new_block_end > self._end) or new_block_start > self._end or new_block_end < self._start ): self._cache = self._fetch_range_proxy(new_block_start, new_block_end) elif new_block_end > self._end: prune_diff: int = new_block_start - self._start self._cache = self._cache[prune_diff:] + self._fetch_range_proxy(self._end, new_block_end) elif new_block_start < self._start: prune_diff = new_block_end - self._end self._cache = self._fetch_range_proxy(new_block_start, self._start) + self._cache[:prune_diff] else: raise RuntimeError("AWSWrangler's cache calculation error.") self._start = new_block_start self._end = new_block_end return None def tell(self) -> int: """Return the current file location.""" return self._loc def seek(self, loc: int, whence: int = 0) -> int: """Set current file location.""" if self.readable() is False: raise OSError(ESPIPE, "Seek only available in read mode") if whence == 0: loc_tmp: int = loc elif whence == 1: loc_tmp = self._loc + loc elif whence == 2: loc_tmp = self._size + loc else: raise ValueError(f"invalid whence ({whence}, should be 0, 1 or 2).") if loc_tmp < 0: raise ValueError("Seek before start of file") self._loc = loc_tmp return self._loc def flush(self, force: bool = False) -> None: """Write buffered data to S3.""" if self.closed: raise RuntimeError("I/O operation on closed file.") if self.writable() and self._buffer.closed is False: total_size: int = self._buffer.tell() if total_size < _MIN_WRITE_BLOCK and force is False: return None if total_size == 0: return None try: _logger.debug("Flushing: %s bytes", total_size) self._mpu = self._mpu or _utils.try_it( f=self._client.create_multipart_upload, # type: ignore[arg-type] ex=_S3_RETRYABLE_ERRORS, base=0.5, max_num_tries=6, Bucket=self._bucket, Key=self._key, **get_botocore_valid_kwargs( function_name="create_multipart_upload", s3_additional_kwargs=self._s3_additional_kwargs ), ) self._buffer.seek(0) for chunk_size in _utils.get_even_chunks_sizes( total_size=total_size, chunk_size=_MIN_WRITE_BLOCK, upper_bound=False ): _logger.debug("chunk_size: %s bytes", chunk_size) self._parts_count += 1 self._upload_proxy.upload( bucket=self._bucket, key=self._key, part=self._parts_count, upload_id=self._mpu["UploadId"], data=self._buffer.read(chunk_size), s3_client=self._client, boto3_kwargs=get_botocore_valid_kwargs( function_name="upload_part", s3_additional_kwargs=self._s3_additional_kwargs ), ) finally: # Ensure that the buffer is cleared (even in the event of an exception) so that # any partial data doesn't get written when close() is called. self._buffer.seek(0) self._buffer.truncate(0) self._buffer.close() self._buffer = io.BytesIO() return None def readable(self) -> bool: """Return whether this object is opened for reading.""" return "r" in self._mode def seekable(self) -> bool: """Return whether this object is opened for seeking.""" return self.readable() def writable(self) -> bool: """Return whether this object is opened for writing.""" return "w" in self._mode def close(self) -> None: """Clean up the cache.""" if self.closed: return None if self.writable(): _logger.debug("Closing: %s parts", self._parts_count) if self._parts_count > 0: self.flush(force=True) parts: list[dict[str, str | int]] = self._upload_proxy.close() part_info: dict[str, list[dict[str, Any]]] = {"Parts": parts} _logger.debug("Running complete_multipart_upload...") _utils.try_it( f=self._client.complete_multipart_upload, ex=_S3_RETRYABLE_ERRORS, base=0.5, max_num_tries=6, Bucket=self._bucket, Key=self._key, UploadId=self._mpu["UploadId"], MultipartUpload=part_info, **get_botocore_valid_kwargs( function_name="complete_multipart_upload", s3_additional_kwargs=self._s3_additional_kwargs ), ) _logger.debug("complete_multipart_upload done!") elif self._buffer.tell() > 0: _logger.debug("put_object") _utils.try_it( f=self._client.put_object, ex=_S3_RETRYABLE_ERRORS, base=0.5, max_num_tries=6, Bucket=self._bucket, Key=self._key, Body=self._buffer.getvalue(), **get_botocore_valid_kwargs( function_name="put_object", s3_additional_kwargs=self._s3_additional_kwargs ), ) self._parts_count = 0 self._upload_proxy.close() self._buffer.seek(0) self._buffer.truncate(0) self._buffer.close() elif self.readable(): self._cache = b"" else: raise RuntimeError(f"Invalid mode: {self._mode}") super().close() return None def read(self, length: int = -1) -> bytes: """Return cached data and fetch on demand chunks.""" if self.readable() is False: raise ValueError("File not in read mode.") if self.closed is True: raise ValueError("I/O operation on closed file.") if length < 0 or self._loc + length > self._size: length = self._size - self._loc self._fetch(self._loc, self._loc + length) out: bytes = self._cache[self._loc - self._start : self._loc - self._start + length] self._loc += len(out) return out def readline(self, length: int | None = -1) -> bytes: """Read until the next line terminator.""" length = -1 if length is None else length end: int = self._loc + self._s3_block_size end = self._size if end > self._size else end self._fetch(self._loc, end) while True: found: int = self._cache[self._loc - self._start :].find(self._newline.encode(encoding=self._encoding)) if 0 < length < found: return self.read(length + 1) if found >= 0: return self.read(found + 1) if self._end >= self._size: return self.read(-1) end = self._end + self._s3_half_block_size end = self._size if end > self._size else end self._fetch(self._loc, end) def write(self, data: bytes | bytearray | memoryview) -> int: # type: ignore[override] """Write data to buffer and only upload on close() or if buffer is greater than or equal to _MIN_WRITE_BLOCK.""" if self.writable() is False: raise RuntimeError("File not in write mode.") if self.closed: raise RuntimeError("I/O operation on closed file.") n: int = self._buffer.write(data) self._loc += n if self._buffer.tell() >= _MIN_WRITE_BLOCK: self.flush() return n @contextmanager @apply_configs def open_s3_object( path: str, mode: str = "wb", version_id: str | None = None, use_threads: bool | int = False, s3_additional_kwargs: dict[str, str] | None = None, s3_block_size: int = -1, # One shot download boto3_session: boto3.Session | None = None, s3_client: "S3Client" | None = None, newline: str | None = "\n", encoding: str | None = "utf-8", ) -> Iterator[_S3ObjectBase | io.TextIOWrapper]: """Return a _S3Object or TextIOWrapper based on the received mode.""" s3obj: _S3ObjectBase | None = None text_s3obj: io.TextIOWrapper | None = None s3_client = _utils.client(service_name="s3", session=boto3_session) if not s3_client else s3_client try: s3obj = _S3ObjectBase( path=path, s3_block_size=s3_block_size, mode=mode, version_id=version_id, use_threads=use_threads, s3_client=s3_client, s3_additional_kwargs=s3_additional_kwargs, encoding=encoding, newline=newline, ) if "b" in mode: # binary yield s3obj else: # text text_s3obj = io.TextIOWrapper( buffer=cast(BinaryIO, s3obj), encoding=encoding, newline=newline, line_buffering=False, write_through=False, ) yield text_s3obj finally: if text_s3obj is not None and text_s3obj.closed is False: text_s3obj.close() if s3obj is not None and s3obj.closed is False: s3obj.close()