fix: use valid encryption key arg for S3 downloads (#1414)

laurenyu · web-flow · commit 9f81006a1f79 · 2020-04-15T11:16:47.000-07:00
diff --git a/src/sagemaker/s3.py b/src/sagemaker/s3.py
@@ -110,7 +110,7 @@ def download(s3_uri, local_path, kms_key=None, session=None):
         sagemaker_session = session or Session()
         bucket, key_prefix = parse_s3_url(url=s3_uri)
         if kms_key is not None:
-            extra_args = {"SSEKMSKeyId": kms_key}
+            extra_args = {"SSECustomerKey": kms_key}
         else:
             extra_args = None
 
diff --git a/tests/unit/test_s3.py b/tests/unit/test_s3.py
@@ -91,5 +91,5 @@ def test_download_with_kms_key(sagemaker_session):
         path="/path/for/download/",
         bucket=BUCKET_NAME,
         key_prefix=os.path.join(CURRENT_JOB_NAME, SOURCE_NAME),
-        extra_args={"SSEKMSKeyId": KMS_KEY},
+        extra_args={"SSECustomerKey": KMS_KEY},
     )

Original file line number	Diff line number	Diff line change
`@@ -91,5 +91,5 @@ def test_download_with_kms_key(sagemaker_session):`
`91`	`91`	`path="/path/for/download/",`
`92`	`92`	`bucket=BUCKET_NAME,`
`93`	`93`	`key_prefix=os.path.join(CURRENT_JOB_NAME, SOURCE_NAME),`
`94`		`- extra_args={"SSEKMSKeyId": KMS_KEY},`
	`94`	`+ extra_args={"SSECustomerKey": KMS_KEY},`
`95`	`95`	`)`