diff --git a/dvc/config.py b/dvc/config.py index c6ebc6399..0854ae930 100644 --- a/dvc/config.py +++ b/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, diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 0b8237549..6366fc426 100644 --- a/dvc/remote/s3.py +++ b/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) diff --git a/tests/unit/remote/test_s3.py b/tests/unit/remote/test_s3.py new file mode 100644 index 000000000..57fcfcec7 --- /dev/null +++ b/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)