From 2ecac380c67202c16279137c14539ab2acd9bbe9 Mon Sep 17 00:00:00 2001
From: Lauren Yu <6631887+laurenyu@users.noreply.github.com>
Date: Wed, 15 Apr 2020 10:18:22 -0700
Subject: [PATCH] fix: use valid encryption key arg for S3 downloads

---
 src/sagemaker/s3.py   | 2 +-
 tests/unit/test_s3.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/sagemaker/s3.py b/src/sagemaker/s3.py
index 57fdef6c06..d81710c412 100644
--- 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
index ba995ef86f..c073417116 100644
--- 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},
     )