Merge pull request #1091 from pyathena-dev/refactor/1086-get-put-object

Move object GET and PUT requests into S3Core

Author: laughingman7743

docs/filesystem.md

@@ -300,7 +300,7 @@ directories below the bucket level) and is always a no-op.
 ## Typed S3 operations
 
 `S3FileSystem.core` is an `S3Core`, the typed operations that the filesystem sends
-its listing, lookup, delete, multipart upload and copy requests with. It can also be
+its listing, lookup, read, write, delete, multipart upload and copy requests with. It can also be
 built on a boto3 S3 client. Each operation sends one request (one per page for the
 iterators and `list_object_annotations()`); `plan_multipart_copy()` and
 `copy_object_annotation()`, described below, send several. The requests are sent with
@@ -326,6 +326,18 @@ for page in core.list_objects("YOUR_S3_BUCKET", prefix="path/to/", delimiter="/"
 
 `list_object_versions()` and `list_buckets()` return pages in the same way.
 
+`get_object()` reads an object, a version of it, or a byte range and returns the bytes.
+A range has an exclusive end; `None` as the end reads to the end of the object, and a
+negative start without an end reads the last bytes. The response body is closed after
+the read; a failure to read it raises the botocore exception and is not retried.
+`put_object()` writes bytes and returns an `S3PutObject`.
+
+```python
+data = core.get_object(S3Path.parse("s3://YOUR_S3_BUCKET/path/to/object"), (0, 1024))
+result = core.put_object(S3Path("YOUR_S3_BUCKET", "path/to/copy"), data)
+print(result.etag, result.version_id)
+```
+
 `delete_objects()` deletes an `S3DeleteBatch`, the objects of one bucket that one
 DeleteObjects request accepts. `S3DeleteBatch.from_paths()` groups paths into batches
 of up to 1,000 objects per bucket. The objects that S3 could not delete are in the

pyathena/filesystem/s3.py

@@ -44,7 +44,6 @@
     S3Object,
     S3ObjectType,
     S3ObjectVersion,
-    S3PutObject,
     S3StorageClass,
 )
 from pyathena.filesystem.s3_path import S3Path
@@ -1403,7 +1402,9 @@ def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]:
         if not s3_path.key:
             raise ValueError("Cannot touch the bucket.")
 
-        object_ = self._put_object(bucket=s3_path.bucket, key=s3_path.key, body=None, **kwargs)
+        # The empty body is passed, so that a body keyword argument raises
+        # TypeError.
+        object_ = self.core.put_object(s3_path, None, **kwargs)
         self.invalidate_cache(path)
         return object_.to_dict()
 
@@ -1873,7 +1874,7 @@ def pipe_file(
             # buffered path does.
             value = bytes(value)
 
-        self._put_object(bucket=s3_path.bucket, key=s3_path.key, body=value, **request_kwargs)
+        self.core.put_object(s3_path, value, **request_kwargs)
         self.invalidate_cache(path)
 
     def _finish_multipart_upload(
@@ -2012,13 +2013,7 @@ def cat_file(
                     return b""
                 ranges = (start, end)
         try:
-            return self._get_object(
-                bucket=s3_path.bucket,
-                key=s3_path.key,
-                ranges=ranges,
-                version_id=version_id,
-                **kwargs,
-            )[1]
+            return self.core.get_object(s3_path.with_version_id(version_id), ranges, **kwargs)
         except OSError as e:
             if (
                 ranges
@@ -2869,69 +2864,6 @@ def _open(
             **kwargs,
         )
 
-    def _get_object(
-        self,
-        bucket: str,
-        key: str,
-        ranges: tuple[int, int | None] | None = None,
-        version_id: str | None = None,
-        **kwargs,
-    ) -> tuple[int, bytes]:
-        """Read an object or a byte range of it with GetObject.
-
-        Args:
-            bucket: The bucket name.
-            key: The object key.
-            ranges: The ``(start, end)`` byte range to read, with an exclusive
-                end or ``None`` to read to the end of the object (the last
-                ``-start`` bytes for a negative start), or ``None`` to read
-                the whole object.
-            version_id: The version ID to read, or ``None`` for the latest.
-            **kwargs: Additional parameters passed to the GetObject API.
-
-        Returns:
-            Tuple of the start of the range as given (0 for the whole
-            object) and the bytes read.
-
-        Raises:
-            ValueError: If the range is empty. S3 ignores a range whose last
-                byte precedes its first byte and returns the whole object.
-        """
-        request = {"Bucket": bucket, "Key": key}
-        if ranges:
-            if ranges[1] is not None and ranges[0] >= ranges[1]:
-                raise ValueError(f"Invalid empty range: {ranges}.")
-            range_ = S3File._format_ranges(ranges)
-            request.update({"Range": range_})
-        else:
-            ranges = (0, 0)
-            range_ = "bytes=0-"
-        if version_id:
-            request.update({"VersionId": version_id})
-
-        _logger.debug(f"Get object: {S3Path(bucket, key, version_id).uri} range={range_}")
-        response = self._call(
-            self._client.get_object,
-            # The fields of the request take precedence over inherited
-            # parameters of the same name.
-            **{**kwargs, **request},
-        )
-        return ranges[0], cast(bytes, response["Body"].read())
-
-    def _put_object(self, bucket: str, key: str, body: bytes | None, **kwargs) -> S3PutObject:
-        request: dict[str, Any] = {"Bucket": bucket, "Key": key}
-        if body:
-            request.update({"Body": body})
-
-        _logger.debug(f"Put object: s3://{bucket}/{key}")
-        response = self._call(
-            self._client.put_object,
-            # The fields of the request take precedence over inherited
-            # parameters of the same name.
-            **{**kwargs, **request},
-        )
-        return S3PutObject(response)
-
     def _call(self, method: str | Callable[..., Any], **kwargs) -> dict[str, Any]:
         """Send a request with the core (see :meth:`S3Core.call`).
 
@@ -3364,11 +3296,8 @@ def commit(self) -> None:
             # Upload files smaller than block size.
             self.buffer.seek(0)
             data = self.buffer.read()
-            self.fs._put_object(
-                bucket=self.bucket,
-                key=self.key,
-                body=data,
-                **self._get_request_kwargs("put_object"),
+            self.fs.core.put_object(
+                S3Path(self.bucket, self.key), data, **self._get_request_kwargs("put_object")
             )
         else:
             if not self.multipart_upload:
@@ -3488,45 +3417,18 @@ def _fetch_range(self, start: int, end: int) -> bytes:
         ranges = self._get_ranges(
             start, end, max_workers=self.max_workers, worker_block_size=self.blocksize
         )
+        path = S3Path(self.bucket, self.key, self.version_id)
+        request_kwargs = self._get_request_kwargs("get_object")
         if len(ranges) > 1:
             futures = [
-                self._executor.submit(
-                    self.fs._get_object,
-                    bucket=self.bucket,
-                    key=self.key,
-                    ranges=r,
-                    version_id=self.version_id,
-                    **self._get_request_kwargs("get_object"),
-                )
+                self._executor.submit(self.fs.core.get_object, path, r, **request_kwargs)
                 for r in ranges
             ]
-            object_ = self._merge_objects([f.result() for f in as_completed(futures)])
-        else:
-            object_ = self.fs._get_object(
-                self.bucket,
-                self.key,
-                ranges[0],
-                self.version_id,
-                **self._get_request_kwargs("get_object"),
-            )[1]
-        return object_
-
-    @staticmethod
-    def _format_ranges(ranges: tuple[int, int | None]) -> str:
-        """Format a byte range as the value of an HTTP ``Range`` header.
-
-        Args:
-            ranges: The ``(start, end)`` byte range, with an exclusive end or
-                ``None`` for the end of the object. A negative start with no
-                end selects the last ``-start`` bytes.
-
-        Returns:
-            The range, such as ``bytes=0-99``, ``bytes=100-`` or ``bytes=-8``.
-        """
-        start, end = ranges
-        if end is None:
-            return f"bytes={start}" if start < 0 else f"bytes={start}-"
-        return f"bytes={start}-{end - 1}"
+            for future in as_completed(futures):
+                # The failure of the first read to finish is raised.
+                future.result()
+            return b"".join(f.result() for f in futures)
+        return self.fs.core.get_object(path, ranges[0], **request_kwargs)
 
     @staticmethod
     def _get_ranges(
@@ -3548,8 +3450,3 @@ def _get_ranges(
         else:
             ranges.append((start, end))
         return ranges
-
-    @staticmethod
-    def _merge_objects(objects: list[tuple[int, bytes]]) -> bytes:
-        objects.sort(key=lambda x: x[0])
-        return b"".join([obj for start, obj in objects])

pyathena/filesystem/s3_core.py

@@ -9,6 +9,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import math
 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
@@ -27,6 +28,7 @@
     S3MultipartUpload,
     S3MultipartUploadPart,
     S3ObjectVersion,
+    S3PutObject,
 )
 from pyathena.filesystem.s3_path import S3Path
 from pyathena.util import RetryConfig, override, retry_api_call
@@ -606,6 +608,92 @@ def head_bucket(self, bucket: str, **params) -> S3Bucket:
         response = self.call(self._client.head_bucket, Bucket=bucket, **params)
         return S3Bucket(name=bucket, bucket_region=response.get("BucketRegion"))
 
+    def get_object(
+        self, path: S3Path, range_: tuple[int, int | None] | None = None, **params
+    ) -> bytes:
+        """Read an object, a version of it, or a byte range of it with GetObject.
+
+        The body of the response is read whole and then closed, also when
+        the read fails. A failure to read the body is neither translated nor
+        retried: the botocore exception, such as ``ReadTimeoutError`` or
+        ``IncompleteReadError``, propagates.
+
+        Args:
+            path: The path of the object, with the version ID to read, if
+                any.
+            range_: The ``(start, end)`` byte range to read, with an exclusive
+                end, or None as the end to read to the end of the object. A
+                negative start without an end reads the last ``-start``
+                bytes, or the whole object if it is shorter. None sends no
+                range of its own, so the whole object is read unless
+                ``params`` or ``request_kwargs`` have ``Range``.
+            **params: Additional request parameters. The fields that the
+                other arguments set take precedence over parameters of the
+                same name.
+
+        Returns:
+            The bytes read.
+
+        Raises:
+            ValueError: If the path has no key, or the range is empty or has
+                both a negative start and an end, which S3 would ignore and
+                return the whole object for.
+            FileNotFoundError: If the object or version does not exist.
+            OSError: If the range starts at or past the end of the object.
+                Its ``__cause__`` is the ``ClientError`` with the
+                ``InvalidRange`` code.
+        """
+        if not path.key:
+            raise ValueError(f"The path has no key: {path.uri}.")
+        request: dict[str, Any] = {"Bucket": path.bucket, "Key": path.key}
+        if path.version_id:
+            request.update({"VersionId": path.version_id})
+        if range_ is not None:
+            start, end = range_
+            if end is None:
+                request.update({"Range": f"bytes={start}" if start < 0 else f"bytes={start}-"})
+            elif start < 0 or start >= end:
+                raise ValueError(f"Invalid range: {range_}.")
+            else:
+                request.update({"Range": f"bytes={start}-{end - 1}"})
+        request = {**params, **request}
+        _logger.debug(f"Get object: {path.uri} range={request.get('Range')}")
+        response = self.call(self._client.get_object, **request)
+        # Read through the StreamingBody, which verifies the length and the
+        # checksum of the data; entering it would return the raw stream.
+        with contextlib.closing(response["Body"]) as body:
+            return cast(bytes, body.read())
+
+    def put_object(self, path: S3Path, body: bytes | None = None, **params) -> S3PutObject:
+        """Write an object with PutObject.
+
+        Args:
+            path: The path of the object to write, without a version ID.
+            body: The data to write. None or empty bytes send no body of
+                their own, so an empty object is written unless ``params`` or
+                ``request_kwargs`` have ``Body``.
+            **params: Additional request parameters. The fields that the
+                other arguments set take precedence over parameters of the
+                same name.
+
+        Returns:
+            The result of the write.
+
+        Raises:
+            ValueError: If the path has no key, or has a version ID, which a
+                write cannot replace.
+        """
+        if not path.key:
+            raise ValueError(f"The path has no key: {path.uri}.")
+        if path.version_id:
+            raise ValueError(f"Cannot write to a version: {path.uri}.")
+        request: dict[str, Any] = {"Bucket": path.bucket, "Key": path.key}
+        if body:
+            request.update({"Body": body})
+        _logger.debug(f"Put object: {path.uri}")
+        response = self.call(self._client.put_object, **{**params, **request})
+        return S3PutObject(response)
+
     def delete_object(self, path: S3Path, **params) -> None:
         """Delete an object, or a version of it, with DeleteObject.