Use multipart upload objects and fix checksum completion and cleanup races

docs/filesystem.md

@@ -348,6 +348,29 @@ instead. `copy_object_annotation()` copies one annotation onto the destination a
 the upload completes, with GetObjectAnnotation and PutObjectAnnotation. The
 filesystems' `cp_file()`, `copy()` and `mv()` run these plans.
 
+The multipart primitives take the `S3MultipartUpload` returned by creation,
+keeping its bucket, key, upload ID and checksum configuration together.
+The core retains no upload state.
+`upload_part()` passes the upload's checksum algorithm to the SDK;
+completion selects the matching part checksum and sends the upload's
+`ChecksumType` when present, including `FULL_OBJECT`.
+Without a creation algorithm, completion sends the ETag and part number without
+checksums that the SDK may add to part uploads.
+`abort_multipart_upload()` also accepts uploads returned by
+`list_multipart_uploads()`.
+
+```python
+upload = core.create_multipart_upload(
+    S3Path("YOUR_S3_BUCKET", "path/to/object"), ChecksumAlgorithm="SHA256"
+)
+try:
+    part = core.upload_part(upload, 1, b"data")
+    result = core.complete_multipart_upload(upload, [part])
+except Exception:
+    core.abort_multipart_upload(upload)
+    raise
+```
+
 ## Async filesystem
 
 `AioS3FileSystem` provides the same functionality on top of fsspec's

pyathena/filesystem/s3.py

@@ -1784,18 +1784,14 @@ def _copy_object_with_multipart_upload(
                     wait([creation])
                     if creation.exception() is None:
                         self._abort_multipart_upload(
-                            plan.destination.bucket,
-                            cast(str, plan.destination.key),
-                            cast(str, creation.result().upload_id),
+                            creation.result(),
                             plan.abort_params,
                         )
                 raise
-            upload_id = cast(str, multipart_upload.upload_id)
             futures = [
                 executor.submit(
                     self.core.upload_part_copy,
-                    path=plan.destination,
-                    upload_id=upload_id,
+                    upload=multipart_upload,
                     part_number=i + 1,
                     source=plan.source,
                     range_=range_,
@@ -1804,9 +1800,7 @@ def _copy_object_with_multipart_upload(
                 for i, range_ in enumerate(plan.ranges)
             ]
             completed = self._finish_multipart_upload(
-                bucket=plan.destination.bucket,
-                key=cast(str, plan.destination.key),
-                upload_id=upload_id,
+                upload=multipart_upload,
                 futures=futures,
                 # Filtered again for the completion and the abort, which
                 # leaves the plan's parameters of each unchanged.
@@ -1930,9 +1924,7 @@ def pipe_file(
 
     def _finish_multipart_upload(
         self,
-        bucket: str,
-        key: str,
-        upload_id: str,
+        upload: S3MultipartUpload,
         futures: list[Future[S3MultipartUploadPart]],
         request_kwargs: Mapping[str, Any] | None = None,
         abort: bool = True,
@@ -1946,9 +1938,7 @@ def _finish_multipart_upload(
         false. The original error is then re-raised.
 
         Args:
-            bucket: S3 bucket name.
-            key: Object key being uploaded.
-            upload_id: Unique identifier for the multipart upload.
+            upload: The multipart upload returned by creation.
             futures: Futures of the part uploads, in part-number order.
             request_kwargs: Parameters of the upload, such as
                 ``RequestPayer`` or the SSE-C parameters; the completion and
@@ -1964,8 +1954,7 @@ def _finish_multipart_upload(
             # The futures are in part-number order.
             parts = [future.result() for future in futures]
             return self.core.complete_multipart_upload(
-                S3Path(bucket, key),
-                upload_id,
+                upload,
                 parts,
                 **self.core.operation_params("complete_multipart_upload", request_kwargs),
             )
@@ -1976,33 +1965,31 @@ def _finish_multipart_upload(
             # be stored after the abort, so wait for the parts that could not
             # be cancelled first.
             wait([future for future in futures if not future.cancel()])
-            self._abort_multipart_upload(bucket, key, upload_id, request_kwargs)
+            self._abort_multipart_upload(upload, request_kwargs)
             raise
 
     def _abort_multipart_upload(
-        self, bucket: str, key: str, upload_id: str, request_kwargs: Mapping[str, Any]
+        self, upload: S3MultipartUpload, request_kwargs: Mapping[str, Any]
     ) -> None:
         """Abort a failed multipart upload, logging an error of the abort.
 
         The error of the abort is logged instead of raised, so that the
         caller can re-raise the error that made the upload fail.
 
         Args:
-            bucket: S3 bucket name.
-            key: Object key being uploaded.
-            upload_id: Unique identifier for the multipart upload.
+            upload: The multipart upload returned by creation.
             request_kwargs: Parameters of the upload; the abort receives
                 those that it accepts.
         """
         try:
             self.core.abort_multipart_upload(
-                S3Path(bucket, key),
-                upload_id,
+                upload,
                 **self.core.operation_params("abort_multipart_upload", request_kwargs),
             )
         except Exception:
             _logger.exception(
-                f"Failed to abort multipart upload {upload_id} to s3://{bucket}/{key}."
+                f"Failed to abort multipart upload {upload.upload_id} "
+                f"to s3://{upload.bucket}/{upload.key}."
             )
 
     def cat_file(
@@ -2623,6 +2610,9 @@ def object_version_info(
     def clear_multipart_uploads(self, path: str) -> None:
         """Abort any incomplete multipart uploads in the bucket.
 
+        Uploads completed or aborted after listing are already cleared.
+        All abort results are checked before another abort error is raised.
+
         Args:
             path: S3 bucket or key path (e.g., "bucket", "s3://bucket" or
                 "s3://bucket/prefix"). If the path contains a key, only the
@@ -2632,17 +2622,30 @@ def clear_multipart_uploads(self, path: str) -> None:
         uploads = self.list_multipart_uploads(path)
         if not uploads:
             return
+        error: Exception | None = None
         with self._create_executor(max_workers=self.max_workers) as executor:
             futures = [
                 executor.submit(
                     self.core.abort_multipart_upload,
-                    S3Path(cast(str, upload.bucket), cast(str, upload.key)),
-                    cast(str, upload.upload_id),
+                    upload,
                 )
                 for upload in uploads
             ]
             for future in as_completed(futures):
-                future.result()
+                try:
+                    future.result()
+                except Exception as e:
+                    cause = e.__cause__
+                    if (
+                        isinstance(e, FileNotFoundError)
+                        and isinstance(cause, botocore.exceptions.ClientError)
+                        and S3ClientError(cause).code == "NoSuchUpload"
+                    ):
+                        continue
+                    if error is None:
+                        error = e
+        if error is not None:
+            raise error
 
     def created(self, path: str) -> datetime:
         """Return the creation time of the path.
@@ -3287,8 +3290,7 @@ def _initiate_upload(self) -> None:
                     self.multipart_upload_parts.append(
                         self._executor.submit(
                             self.fs.core.upload_part_copy,
-                            path=path,
-                            upload_id=cast(str, self.multipart_upload.upload_id),
+                            upload=self.multipart_upload,
                             part_number=i + 1,
                             # The existing object is copied into the upload.
                             source=path,
@@ -3300,8 +3302,7 @@ def _initiate_upload(self) -> None:
                 self.multipart_upload_parts.append(
                     self._executor.submit(
                         self.fs.core.upload_part_copy,
-                        path=path,
-                        upload_id=cast(str, self.multipart_upload.upload_id),
+                        upload=self.multipart_upload,
                         part_number=1,
                         source=path,
                         **self._get_request_kwargs("upload_part_copy"),
@@ -3365,8 +3366,7 @@ def _upload_chunk(self, final: bool = False) -> bool:
                 self.multipart_upload_parts.append(
                     self._executor.submit(
                         self.fs.core.upload_part,
-                        path=S3Path(self.bucket, self.key),
-                        upload_id=cast(str, self.multipart_upload.upload_id),
+                        upload=self.multipart_upload,
                         part_number=part_number,
                         body=upload,
                         **self._get_request_kwargs("upload_part"),
@@ -3421,9 +3421,7 @@ def commit(self) -> None:
             upload_id = cast(str, self.multipart_upload.upload_id)
             try:
                 self.fs._finish_multipart_upload(
-                    bucket=self.bucket,
-                    key=self.key,
-                    upload_id=upload_id,
+                    upload=self.multipart_upload,
                     futures=self.multipart_upload_parts,
                     request_kwargs=self.s3_additional_kwargs,
                     abort=False,
@@ -3456,8 +3454,7 @@ def discard(self) -> None:
             # be cancelled first.
             wait([f for f in self.multipart_upload_parts if not f.cancel()])
             self.fs.core.abort_multipart_upload(
-                S3Path(self.bucket, self.key),
-                cast(str, self.multipart_upload.upload_id),
+                self.multipart_upload,
                 **self._get_request_kwargs("abort_multipart_upload"),
             )
 

pyathena/filesystem/s3_async.py

@@ -585,7 +585,7 @@ async def _copy_object_with_multipart_upload(
             # See S3FileSystem._copy_object_with_multipart_upload.
             await asyncio.to_thread(self.core.copy_object, plan.source, plan.destination, **kwargs)
             return
-        upload_id: str
+        multipart_upload: S3MultipartUpload
 
         semaphore = asyncio.Semaphore(max_workers)
         failed = False
@@ -599,8 +599,7 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart
                 try:
                     return await asyncio.to_thread(
                         self.core.upload_part_copy,
-                        path=plan.destination,
-                        upload_id=upload_id,
+                        upload=multipart_upload,
                         part_number=i + 1,
                         source=plan.source,
                         range_=range_,
@@ -630,9 +629,7 @@ async def _abort() -> None:
                     return
             await asyncio.to_thread(
                 self._sync_fs._abort_multipart_upload,
-                plan.destination.bucket,
-                cast(str, plan.destination.key),
-                cast(str, creation.result().upload_id),
+                creation.result(),
                 plan.abort_params,
             )
 
@@ -648,7 +645,6 @@ async def _abort() -> None:
             # shield keeps a cancellation from cancelling the creation, whose
             # thread would keep running, so that _abort() can wait for it.
             multipart_upload = await asyncio.shield(creation)
-            upload_id = cast(str, multipart_upload.upload_id)
             tasks = [asyncio.ensure_future(_upload_part(i, r)) for i, r in enumerate(plan.ranges)]
             # Unlike gather, wait does not cancel the parts when this task is
             # cancelled; their threads would keep copying, so they are waited
@@ -662,8 +658,7 @@ async def _abort() -> None:
             completion = asyncio.ensure_future(
                 asyncio.to_thread(
                     self.core.complete_multipart_upload,
-                    plan.destination,
-                    upload_id,
+                    multipart_upload,
                     cast(list[S3MultipartUploadPart], parts),
                     **plan.complete_params,
                 )