From fd3c84d682aba203f068e5a91062090a84923a6b Mon Sep 17 00:00:00 2001 From: dlg227 Date: Fri, 17 May 2019 16:40:53 -0400 Subject: [PATCH 1/2] remote: add basic support for s3 server side encryption. Fixes #1967 --- dvc/config.py | 3 +++ dvc/remote/s3.py | 28 +++++++++++++++++++++------- 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/dvc/config.py b/dvc/config.py index 8222bf093..1fe903fd0 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -208,6 +208,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes SECTION_AWS_REGION = "region" SECTION_AWS_PROFILE = "profile" SECTION_AWS_USE_SSL = "use_ssl" + SECTION_AWS_SSE = "sse" SECTION_AWS_SCHEMA = { SECTION_AWS_STORAGEPATH: str, Optional(SECTION_AWS_REGION): str, @@ -216,6 +217,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes Optional(SECTION_AWS_ENDPOINT_URL): str, Optional(SECTION_AWS_LIST_OBJECTS, default=False): BOOL_SCHEMA, Optional(SECTION_AWS_USE_SSL, default=True): BOOL_SCHEMA, + Optional(SECTION_AWS_SSE): str, } # backward compatibility @@ -256,6 +258,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes Optional(SECTION_AWS_ENDPOINT_URL): str, Optional(SECTION_AWS_LIST_OBJECTS, default=False): BOOL_SCHEMA, Optional(SECTION_AWS_USE_SSL, default=True): BOOL_SCHEMA, + Optional(SECTION_AWS_SSE): str, Optional(SECTION_GCP_PROJECTNAME): str, Optional(SECTION_CACHE_TYPE): SECTION_CACHE_TYPE_SCHEMA, Optional(SECTION_CACHE_PROTECTED, default=False): BOOL_SCHEMA, diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 99ff523d2..d1fcd1b4a 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -64,6 +64,12 @@ class RemoteS3(RemoteBASE): self.use_ssl = config.get(Config.SECTION_AWS_USE_SSL, True) + self.extra_args = {} + + self.sse = config.get(Config.SECTION_AWS_SSE, "") + if self.sse: + self.extra_args["ServerSideEncryption"] = self.sse + shared_creds = config.get(Config.SECTION_AWS_CREDENTIALPATH) if shared_creds: os.environ.setdefault("AWS_SHARED_CREDENTIALS_FILE", shared_creds) @@ -112,9 +118,11 @@ class RemoteS3(RemoteBASE): return obj @classmethod - def _copy_multipart(cls, s3, from_info, to_info, size, n_parts): + def _copy_multipart( + cls, s3, from_info, to_info, size, n_parts, extra_args={} + ): mpu = s3.create_multipart_upload( - Bucket=to_info.bucket, Key=to_info.path + Bucket=to_info.bucket, Key=to_info.path, **extra_args ) mpu_id = mpu["UploadId"] @@ -154,7 +162,7 @@ class RemoteS3(RemoteBASE): ) @classmethod - def _copy(cls, s3, from_info, to_info): + def _copy(cls, s3, from_info, to_info, extra_args={}): # NOTE: object's etag depends on the way it was uploaded to s3 or the # way it was copied within the s3. More specifically, it depends on # the chunk size that was used to transfer it, which would affect @@ -182,10 +190,12 @@ class RemoteS3(RemoteBASE): _, _, parts_suffix = etag.partition("-") if parts_suffix: n_parts = int(parts_suffix) - cls._copy_multipart(s3, from_info, to_info, size, n_parts) + cls._copy_multipart( + s3, from_info, to_info, size, n_parts, extra_args + ) else: source = {"Bucket": from_info.bucket, "Key": from_info.path} - s3.copy(source, to_info.bucket, to_info.path) + s3.copy(source, to_info.bucket, to_info.path, ExtraArgs=extra_args) cached_etag = cls.get_etag(s3, to_info.bucket, to_info.path) if etag != cached_etag: @@ -193,7 +203,7 @@ class RemoteS3(RemoteBASE): def copy(self, from_info, to_info, s3=None): s3 = s3 if s3 else self.s3 - self._copy(s3, from_info, to_info) + self._copy(s3, from_info, to_info, extra_args=self.extra_args) def remove(self, path_info): if path_info.scheme != "s3": @@ -257,7 +267,11 @@ class RemoteS3(RemoteBASE): try: s3.upload_file( - from_info.path, to_info.bucket, to_info.path, Callback=cb + from_info.path, + to_info.bucket, + to_info.path, + Callback=cb, + ExtraArgs=self.extra_args, ) except Exception: msg = "failed to upload '{}'".format(from_info.path) From fa6b764f97bd7f14fd10eefa8105128bc371b0ed Mon Sep 17 00:00:00 2001 From: dlg227 Date: Sat, 8 Jun 2019 17:10:58 -0400 Subject: [PATCH 2/2] refactor of remote s3 extra_args --- dvc/remote/s3.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index d1fcd1b4a..900f8b7bb 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -117,19 +117,16 @@ class RemoteS3(RemoteBASE): ) return obj - @classmethod - def _copy_multipart( - cls, s3, from_info, to_info, size, n_parts, extra_args={} - ): + def _copy_multipart(self, s3, from_info, to_info, size, n_parts): mpu = s3.create_multipart_upload( - Bucket=to_info.bucket, Key=to_info.path, **extra_args + Bucket=to_info.bucket, Key=to_info.path, **self.extra_args ) mpu_id = mpu["UploadId"] parts = [] byte_position = 0 for i in range(1, n_parts + 1): - obj = cls.get_head_object( + obj = self.get_head_object( s3, from_info.bucket, from_info.path, PartNumber=i ) part_size = obj["ContentLength"] @@ -161,8 +158,7 @@ class RemoteS3(RemoteBASE): MultipartUpload={"Parts": parts}, ) - @classmethod - def _copy(cls, s3, from_info, to_info, extra_args={}): + def _copy(self, s3, from_info, to_info): # NOTE: object's etag depends on the way it was uploaded to s3 or the # way it was copied within the s3. More specifically, it depends on # the chunk size that was used to transfer it, which would affect @@ -183,27 +179,27 @@ class RemoteS3(RemoteBASE): # preserve etag, we need to transfer each part separately, so the # object is transfered in the same chunks as it was originally. - obj = cls.get_head_object(s3, from_info.bucket, from_info.path) + obj = self.get_head_object(s3, from_info.bucket, from_info.path) etag = obj["ETag"].strip('"') size = obj["ContentLength"] _, _, parts_suffix = etag.partition("-") if parts_suffix: n_parts = int(parts_suffix) - cls._copy_multipart( - s3, from_info, to_info, size, n_parts, extra_args - ) + self._copy_multipart(s3, from_info, to_info, size, n_parts) else: source = {"Bucket": from_info.bucket, "Key": from_info.path} - s3.copy(source, to_info.bucket, to_info.path, ExtraArgs=extra_args) + s3.copy( + source, to_info.bucket, to_info.path, ExtraArgs=self.extra_args + ) - cached_etag = cls.get_etag(s3, to_info.bucket, to_info.path) + cached_etag = self.get_etag(s3, to_info.bucket, to_info.path) if etag != cached_etag: raise ETagMismatchError(etag, cached_etag) def copy(self, from_info, to_info, s3=None): s3 = s3 if s3 else self.s3 - self._copy(s3, from_info, to_info, extra_args=self.extra_args) + self._copy(s3, from_info, to_info) def remove(self, path_info): if path_info.scheme != "s3":