Merge pull request #2931 from rbvesb/dvc-s3-acl

remote: Extend S3 remote support for custom ACL

dvc/config.py

@@ -173,6 +173,10 @@ class Config(object):  # pylint: disable=too-many-instance-attributes
     SECTION_AWS_USE_SSL = "use_ssl"
     SECTION_AWS_SSE = "sse"
     SECTION_AWS_ACL = "acl"
+    SECTION_AWS_GRANT_READ = "grant_read"
+    SECTION_AWS_GRANT_READ_ACP = "grant_read_acp"
+    SECTION_AWS_GRANT_WRITE_ACP = "grant_write_acp"
+    SECTION_AWS_GRANT_FULL_CONTROL = "grant_full_control"
 
     # gcp specific options
     SECTION_GCP_CREDENTIALPATH = CREDENTIALPATH
@@ -212,6 +216,10 @@ class Config(object):  # pylint: disable=too-many-instance-attributes
         Optional(SECTION_AWS_USE_SSL, default=True): Bool,
         SECTION_AWS_SSE: str,
         SECTION_AWS_ACL: str,
+        SECTION_AWS_GRANT_READ: str,
+        SECTION_AWS_GRANT_READ_ACP: str,
+        SECTION_AWS_GRANT_WRITE_ACP: str,
+        SECTION_AWS_GRANT_FULL_CONTROL: str,
         SECTION_GCP_PROJECTNAME: str,
         SECTION_CACHE_TYPE: supported_cache_type,
         Optional(SECTION_CACHE_PROTECTED, default=False): Bool,

dvc/remote/s3.py

@@ -7,7 +7,7 @@ import threading
 
 from funcy import cached_property, wrap_prop
 
-from dvc.config import Config
+from dvc.config import Config, ConfigError
 from dvc.exceptions import DvcException
 from dvc.exceptions import ETagMismatchError
 from dvc.path_info import CloudURLInfo
@@ -53,6 +53,8 @@ class RemoteS3(RemoteBASE):
         if self.acl:
             self.extra_args["ACL"] = self.acl
 
+        self._append_aws_grants_to_extra_args(config)
+
         shared_creds = config.get(Config.SECTION_AWS_CREDENTIALPATH)
         if shared_creds:
             os.environ.setdefault("AWS_SHARED_CREDENTIALS_FILE", shared_creds)
@@ -302,3 +304,34 @@ class RemoteS3(RemoteBASE):
                 continue
 
             yield path_info.replace(path=fname)
+
+    def _append_aws_grants_to_extra_args(self, config):
+        # Keys for extra_args can be one of the following list:
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/s3.html#boto3.s3.transfer.S3Transfer.ALLOWED_UPLOAD_ARGS
+        """
+          ALLOWED_UPLOAD_ARGS = [
+            'ACL', 'CacheControl', 'ContentDisposition', 'ContentEncoding',
+            'ContentLanguage', 'ContentType', 'Expires', 'GrantFullControl',
+            'GrantRead', 'GrantReadACP', 'GrantWriteACP', 'Metadata',
+            'RequestPayer', 'ServerSideEncryption', 'StorageClass',
+            'SSECustomerAlgorithm', 'SSECustomerKey', 'SSECustomerKeyMD5',
+            'SSEKMSKeyId', 'WebsiteRedirectLocation'
+          ]
+        """
+
+        grants = {
+            Config.SECTION_AWS_GRANT_FULL_CONTROL: "GrantFullControl",
+            Config.SECTION_AWS_GRANT_READ: "GrantRead",
+            Config.SECTION_AWS_GRANT_READ_ACP: "GrantReadACP",
+            Config.SECTION_AWS_GRANT_WRITE_ACP: "GrantWriteACP",
+        }
+
+        for grant_option, extra_args_key in grants.items():
+            if config.get(grant_option, ""):
+                if self.acl:
+                    raise ConfigError(
+                        "`acl` and `grant_*` AWS S3 config options "
+                        "are mutually exclusive"
+                    )
+
+                self.extra_args[extra_args_key] = config.get(grant_option)

tests/unit/remote/test_s3.py

@@ -0,0 +1,56 @@
+import pytest
+
+from dvc.config import ConfigError
+from dvc.remote.s3 import RemoteS3
+
+
+bucket_name = "bucket-name"
+prefix = "some/prefix"
+url = "s3://{}/{}".format(bucket_name, prefix)
+
+
+@pytest.fixture(autouse=True)
+def grants():
+    return {
+        "grant_read": "id=read-permission-id,id=other-read-permission-id",
+        "grant_read_acp": "id=read-acp-permission-id",
+        "grant_write_acp": "id=write-acp-permission-id",
+        "grant_full_control": "id=full-control-permission-id",
+    }
+
+
+def test_init():
+    config = {"url": url}
+    remote = RemoteS3(None, config)
+
+    assert remote.path_info == url
+
+
+def test_grants():
+    config = {
+        "url": url,
+        "grant_read": "id=read-permission-id,id=other-read-permission-id",
+        "grant_read_acp": "id=read-acp-permission-id",
+        "grant_write_acp": "id=write-acp-permission-id",
+        "grant_full_control": "id=full-control-permission-id",
+    }
+    remote = RemoteS3(None, config)
+
+    assert (
+        remote.extra_args["GrantRead"]
+        == "id=read-permission-id,id=other-read-permission-id"
+    )
+    assert remote.extra_args["GrantReadACP"] == "id=read-acp-permission-id"
+    assert remote.extra_args["GrantWriteACP"] == "id=write-acp-permission-id"
+    assert (
+        remote.extra_args["GrantFullControl"]
+        == "id=full-control-permission-id"
+    )
+
+
+def test_grants_mutually_exclusive_acl_error(grants):
+    for grant_option, grant_value in grants.items():
+        config = {"url": url, "acl": "public-read", grant_option: grant_value}
+
+        with pytest.raises(ConfigError):
+            RemoteS3(None, config)