From 7ca3a1119cd9dd804be3c9163863385b731dac5a Mon Sep 17 00:00:00 2001
From: Jack Wells <jd.wells@gmail.com>
Date: Mon, 11 May 2020 12:52:55 -0700
Subject: [PATCH] Added S3 SSEKMSKeyId config option (#3762)

* Added S3 SSEKMSKeyId config option

* test: unit: s3: add SSEKMSKeyId test

Co-authored-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/config.py                | 1 +
 dvc/remote/s3.py             | 4 ++++
 tests/unit/remote/test_s3.py | 5 +++++
 3 files changed, 10 insertions(+)

diff --git a/dvc/config.py b/dvc/config.py
index dde2cb081..1d96d7555 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -151,6 +151,7 @@ SCHEMA = {
                     Optional("listobjects", default=False): Bool,
                     Optional("use_ssl", default=True): Bool,
                     "sse": str,
+                    "sse_kms_key_id": str,
                     "acl": str,
                     "grant_read": str,
                     "grant_read_acp": str,
diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py
index 0ebb4099f..8d92642e9 100644
--- a/dvc/remote/s3.py
+++ b/dvc/remote/s3.py
@@ -46,6 +46,10 @@ class S3Remote(BaseRemote):
         if self.sse:
             self.extra_args["ServerSideEncryption"] = self.sse
 
+        self.sse_kms_key_id = config.get("sse_kms_key_id")
+        if self.sse_kms_key_id:
+            self.extra_args["SSEKMSKeyId"] = self.sse_kms_key_id
+
         self.acl = config.get("acl")
         if self.acl:
             self.extra_args["ACL"] = self.acl
diff --git a/tests/unit/remote/test_s3.py b/tests/unit/remote/test_s3.py
index c32091e0a..b9ee1ebd3 100644
--- a/tests/unit/remote/test_s3.py
+++ b/tests/unit/remote/test_s3.py
@@ -53,3 +53,8 @@ def test_grants_mutually_exclusive_acl_error(dvc, grants):
 
         with pytest.raises(ConfigError):
             S3Remote(dvc, config)
+
+
+def test_sse_kms_key_id(dvc):
+    remote = S3Remote(dvc, {"url": url, "sse_kms_key_id": "key"})
+    assert remote.extra_args["SSEKMSKeyId"] == "key"