Merge customizations for S3

Author: aws-sdk-python-automation

.changes/next-release/enhancement-s3copies-74028.json

@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "``s3`` copies",
+  "description": "Adds ``all`` option to ``--copy-props`` for ``cp``, ``mv``, and ``sync`` commands. When set, S3 to S3 copy operations will copy object annotations, metadata, and tags."
+}

awscli/botocore/data/s3/2006-03-01/paginators-1.sdk-extras.json

@@ -32,6 +32,16 @@
           "Prefix"
         ]
       },
+      "ListObjectAnnotations": {
+        "non_aggregate_keys": [
+          "AnnotationCount",
+          "AnnotationPrefix",
+          "Bucket",
+          "Key",
+          "ObjectVersionId",
+          "RequestCharged"
+        ]
+      },
       "ListParts": {
         "non_aggregate_keys": [
           "ChecksumAlgorithm",

awscli/customizations/s3/subcommands.py

@@ -452,7 +452,7 @@
 
 COPY_PROPS = {
     'name': 'copy-props',
-    'choices': ['none', 'metadata-directive', 'default'],
+    'choices': ['none', 'metadata-directive', 'default', 'all'],
     'default': 'default',
     'help_text': (
         'Determines which properties are copied from the source S3 object. '
@@ -468,6 +468,9 @@
         '<li>``default`` - The default value. Copies tags and properties '
         'covered under the ``metadata-directive`` value from the '
         'source S3 object.</li>'
+        '<li>``all`` - Copies annotations in addition to the tags and '
+        'properties covered under the ``default`` value from the source '
+        'S3 object.</li>'
         '</ul>'
         'In order to copy the appropriate properties for multipart copies, '
         'some of the options may require additional API calls if a multipart '
@@ -476,8 +479,11 @@
         '<li>``metadata-directive`` may require additional ``HeadObject`` '
         'API calls.</li>'
         '<li>``default`` may require additional ``HeadObject``, '
-        '``GetObjectTagging``, and ``PutObjectTagging`` API calls. Note this'
-        ' list of API calls may grow in the future in order to ensure '
+        '``GetObjectTagging``, and ``PutObjectTagging`` API calls.</li>'
+        '<li>``all`` may require the additional API calls covered under the '
+        '``default`` value, as well as ``ListObjectAnnotations``, '
+        '``GetObjectAnnotation``, and ``PutObjectAnnotation`` API calls. Note '
+        'this list of API calls may grow in the future in order to ensure '
         'multipart copies preserve the exact properties a ``CopyObject`` '
         'API call would preserve.</li>'
         '</ul>'

awscli/customizations/s3/subscribers.py

@@ -258,12 +258,14 @@ def _get_none_subscribers(self, fileinfo):
         return [
             ReplaceMetadataDirectiveSubscriber(),
             ReplaceTaggingDirectiveSubscriber(),
+            ExcludeAnnotationDirectiveSubscriber(),
         ]
 
     def _get_metadata_directive_subscribers(self, fileinfo):
         return [
             self._create_metadata_directive_props_subscriber(fileinfo),
             ReplaceTaggingDirectiveSubscriber(),
+            ExcludeAnnotationDirectiveSubscriber(),
         ]
 
     def _get_default_subscribers(self, fileinfo):
@@ -275,6 +277,19 @@ def _get_default_subscribers(self, fileinfo):
                 self._cli_params,
                 source_client=fileinfo.source_client,
             ),
+            ExcludeAnnotationDirectiveSubscriber(),
+        ]
+
+    def _get_all_subscribers(self, fileinfo):
+        return [
+            self._create_metadata_directive_props_subscriber(fileinfo),
+            SetTagsSubscriber(
+                self._client,
+                self._transfer_config,
+                self._cli_params,
+                source_client=fileinfo.source_client,
+            ),
+            self._create_annotations_subscriber(fileinfo),
         ]
 
     def _create_metadata_directive_props_subscriber(self, fileinfo):
@@ -289,20 +304,43 @@ def _create_metadata_directive_props_subscriber(self, fileinfo):
             )
         return SetMetadataDirectivePropsSubscriber(**subscriber_kwargs)
 
+    def _create_annotations_subscriber(self, fileinfo):
+        kwargs = {
+            'client': self._client,
+            'transfer_config': self._transfer_config,
+            'cli_params': self._cli_params,
+            'source_client': fileinfo.source_client,
+        }
+        if not self._cli_params.get('dir_op'):
+            kwargs['head_object_response'] = (
+                fileinfo.associated_response_data
+            )
+        return SetAnnotationsSubscriber(**kwargs)
+
 
-class ReplaceDirectiveSubscriber(BaseSubscriber):
+class SetDirectiveSubscriber(BaseSubscriber):
     _DIRECTIVE_PARAM = ''
+    _DIRECTIVE_VALUE = ''
 
     def on_queued(self, future, **kwargs):
-        future.meta.call_args.extra_args[self._DIRECTIVE_PARAM] = 'REPLACE'
+        future.meta.call_args.extra_args[self._DIRECTIVE_PARAM] = (
+            self._DIRECTIVE_VALUE
+        )
 
 
-class ReplaceMetadataDirectiveSubscriber(ReplaceDirectiveSubscriber):
+class ReplaceMetadataDirectiveSubscriber(SetDirectiveSubscriber):
     _DIRECTIVE_PARAM = 'MetadataDirective'
+    _DIRECTIVE_VALUE = 'REPLACE'
 
 
-class ReplaceTaggingDirectiveSubscriber(ReplaceDirectiveSubscriber):
+class ReplaceTaggingDirectiveSubscriber(SetDirectiveSubscriber):
     _DIRECTIVE_PARAM = 'TaggingDirective'
+    _DIRECTIVE_VALUE = 'REPLACE'
+
+
+class ExcludeAnnotationDirectiveSubscriber(SetDirectiveSubscriber):
+    _DIRECTIVE_PARAM = 'AnnotationDirective'
+    _DIRECTIVE_VALUE = 'EXCLUDE'
 
 
 class SetMetadataDirectivePropsSubscriber(BaseSubscriber):
@@ -442,3 +480,137 @@ def _serialize_to_header_value(self, tags):
 
     def _is_multipart_copy(self, future):
         return future.meta.size >= self._transfer_config.multipart_threshold
+
+
+class AnnotationCopyError(Exception):
+    def __init__(self, bucket, key, succeeded, failed):
+        succeeded_names = ', '.join(succeeded) or '(none)'
+        failed_descriptions = '; '.join(
+            f'{name}: {error}' for name, error in failed
+        )
+        super().__init__(
+            f'Failed to copy all annotations to s3://{bucket}/{key}. '
+            f'The object was copied successfully and was not deleted. '
+            f'Annotations written: {succeeded_names}. '
+            f'Annotations that failed: {failed_descriptions}.'
+        )
+
+
+class SetAnnotationsSubscriber(OnDoneFilteredSubscriber):
+    _ANNOTATIONS_CONTEXT_KEY = 'CopySourceAnnotations'
+
+    def __init__(
+        self, client, transfer_config, cli_params, source_client,
+        head_object_response=None,
+    ):
+        self._client = client
+        self._transfer_config = transfer_config
+        self._cli_params = cli_params
+        self._source_client = source_client
+        self._head_object_response = head_object_response
+
+    def on_queued(self, future, **kwargs):
+        # Annotations only need to be copied for multipart copies. Single-part
+        # copies carry them over server-side.
+        if not self._is_multipart_copy(future):
+            return
+        bucket, key = self._get_bucket_key_from_copy_source(future)
+        source_version_id = self._get_source_version_id()
+        annotation_names = self._list_annotation_names(
+            bucket, key, source_version_id
+        )
+        if not annotation_names:
+            return
+        annotations = {}
+        for name in annotation_names:
+            annotations[name] = self._get_annotation(
+                bucket, key, name, source_version_id
+            )
+        future.meta.user_context[self._ANNOTATIONS_CONTEXT_KEY] = annotations
+
+    def _on_success(self, future):
+        annotations = future.meta.user_context.get(
+            self._ANNOTATIONS_CONTEXT_KEY
+        )
+        if not annotations:
+            return
+        bucket = future.meta.call_args.bucket
+        key = future.meta.call_args.key
+        dest_etag, dest_version_id = self._get_dest_object_identity(future)
+        succeeded = []
+        failed = []
+        for name, payload in annotations.items():
+            try:
+                self._put_annotation(
+                    bucket, key, name, payload, dest_etag, dest_version_id
+                )
+                succeeded.append(name)
+            except Exception as e:
+                failed.append((name, str(e)))
+        if failed:
+            future.set_exception(
+                AnnotationCopyError(bucket, key, succeeded, failed)
+            )
+
+    def _get_dest_object_identity(self, future):
+        response = future.result()
+        return response.get('ETag'), response.get('VersionId')
+
+    def _get_source_version_id(self):
+        if self._head_object_response is not None:
+            return self._head_object_response.get('VersionId')
+        return None
+
+    def _list_annotation_names(self, bucket, key, version_id=None):
+        extra_args = {}
+        utils.RequestParamsMapper.map_list_object_annotations_params(
+            extra_args, self._cli_params
+        )
+        if version_id is not None:
+            extra_args['VersionId'] = version_id
+        paginator = self._source_client.get_paginator(
+            'list_object_annotations'
+        )
+        names = []
+        for page in paginator.paginate(Bucket=bucket, Key=key, **extra_args):
+            for annotation in page.get('Annotations', []):
+                names.append(annotation['AnnotationName'])
+        return names
+
+    def _get_annotation(self, bucket, key, name, version_id=None):
+        extra_args = {}
+        utils.RequestParamsMapper.map_get_object_annotation_params(
+            extra_args, self._cli_params
+        )
+        if version_id is not None:
+            extra_args['VersionId'] = version_id
+        response = self._source_client.get_object_annotation(
+            Bucket=bucket, Key=key, AnnotationName=name, **extra_args
+        )
+        return response['AnnotationPayload'].read()
+
+    def _put_annotation(
+        self, bucket, key, name, payload, dest_etag, dest_version_id
+    ):
+        extra_args = {}
+        utils.RequestParamsMapper.map_put_object_annotation_params(
+            extra_args, self._cli_params
+        )
+        if dest_etag is not None:
+            extra_args['ObjectIfMatch'] = dest_etag
+        if dest_version_id is not None:
+            extra_args['VersionId'] = dest_version_id
+        self._client.put_object_annotation(
+            Bucket=bucket,
+            Key=key,
+            AnnotationName=name,
+            AnnotationPayload=payload,
+            **extra_args,
+        )
+
+    def _get_bucket_key_from_copy_source(self, future):
+        copy_source = future.meta.call_args.copy_source
+        return copy_source['Bucket'], copy_source['Key']
+
+    def _is_multipart_copy(self, future):
+        return future.meta.size >= self._transfer_config.multipart_threshold

awscli/customizations/s3/utils.py

@@ -512,6 +512,21 @@ def map_put_object_tagging_params(cls, request_params, cli_params):
         """Map CLI params to PutObjectTagging request params"""
         cls._set_request_payer_param(request_params, cli_params)
 
+    @classmethod
+    def map_list_object_annotations_params(cls, request_params, cli_params):
+        """Map CLI params to ListObjectAnnotations request params"""
+        cls._set_request_payer_param(request_params, cli_params)
+
+    @classmethod
+    def map_get_object_annotation_params(cls, request_params, cli_params):
+        """Map CLI params to GetObjectAnnotation request params"""
+        cls._set_request_payer_param(request_params, cli_params)
+
+    @classmethod
+    def map_put_object_annotation_params(cls, request_params, cli_params):
+        """Map CLI params to PutObjectAnnotation request params"""
+        cls._set_request_payer_param(request_params, cli_params)
+
     @classmethod
     def map_copy_object_params(cls, request_params, cli_params):
         """Map CLI params to CopyObject request params"""

awscli/s3transfer/copies.py

@@ -69,6 +69,7 @@ class CopySubmissionTask(SubmissionTask):
         'CopySourceSSECustomerKeyMD5',
         'MetadataDirective',
         'TaggingDirective',
+        'AnnotationDirective',
         'IfNoneMatch',
     ]
 
@@ -81,7 +82,6 @@ class CopySubmissionTask(SubmissionTask):
         'IfNoneMatch',
     ]
 
-
     def _submit(
         self, client, config, osutil, request_executor, transfer_future
     ):

awscli/s3transfer/tasks.py

@@ -376,7 +376,7 @@ def _main(self, client, bucket, key, upload_id, parts, extra_args):
         :param extra_args:  A dictionary of any extra arguments that may be
             used in completing the multipart transfer.
         """
-        client.complete_multipart_upload(
+        return client.complete_multipart_upload(
             Bucket=bucket,
             Key=key,
             UploadId=upload_id,

tests/functional/botocore/test_paginator_config.py

@@ -111,6 +111,8 @@
     's3.ListObjectsV2.KeyCount',
     's3.ListObjectsV2.Name',
     's3.ListObjectsV2.EncodingType',
+    's3.ListObjectAnnotations.ContinuationToken',
+    's3.ListObjectAnnotations.MaxAnnotationResults',
     's3.ListParts.PartNumberMarker',
     's3.ListParts.AbortDate',
     's3.ListParts.MaxParts',

tests/functional/s3/__init__.py

@@ -116,6 +116,15 @@ def get_object_tagging_response(self, tags):
     def put_object_tagging_response(self):
         return 'PutObjectTagging', self.empty_response()
 
+    def list_object_annotations_response(self, names):
+        return {'Annotations': [{'AnnotationName': name} for name in names]}
+
+    def get_object_annotation_response(self, payload):
+        return {'AnnotationPayload': payload}
+
+    def put_object_annotation_response(self):
+        return self.empty_response()
+
     def empty_response(self):
         return {}
 
@@ -162,6 +171,7 @@ def copy_object_request(
             'Key': key,
             'CopySource': {'Bucket': source_bucket, 'Key': source_key},
         }
+        override_kwargs.setdefault('AnnotationDirective', 'EXCLUDE')
         params.update(override_kwargs)
         return 'CopyObject', params
 
@@ -229,6 +239,34 @@ def put_object_tagging_request(self, bucket, key, tags):
             },
         }
 
+    def list_object_annotations_request(self, bucket, key, **override_kwargs):
+        params = {'Bucket': bucket, 'Key': key}
+        params.update(override_kwargs)
+        return 'ListObjectAnnotations', params
+
+    def get_object_annotation_request(
+        self, bucket, key, annotation_name, **override_kwargs
+    ):
+        params = {
+            'Bucket': bucket,
+            'Key': key,
+            'AnnotationName': annotation_name,
+        }
+        params.update(override_kwargs)
+        return 'GetObjectAnnotation', params
+
+    def put_object_annotation_request(
+        self, bucket, key, annotation_name, payload, **override_kwargs
+    ):
+        params = {
+            'Bucket': bucket,
+            'Key': key,
+            'AnnotationName': annotation_name,
+            'AnnotationPayload': payload,
+        }
+        params.update(override_kwargs)
+        return 'PutObjectAnnotation', params
+
     def no_such_key_error_response(self):
         return {
             'Error': {