diff --git a/dvc/config_schema.py b/dvc/config_schema.py index 095ebc733..a94b50f92 100644 --- a/dvc/config_schema.py +++ b/dvc/config_schema.py @@ -141,6 +141,7 @@ SCHEMA = { "endpointurl": str, "access_key_id": str, "secret_access_key": str, + "session_token": str, Optional("listobjects", default=False): Bool, # obsoleted Optional("use_ssl", default=True): Bool, "sse": str, diff --git a/dvc/tree/s3.py b/dvc/tree/s3.py index 1064348e5..4d1acb541 100644 --- a/dvc/tree/s3.py +++ b/dvc/tree/s3.py @@ -54,6 +54,7 @@ class S3Tree(BaseTree): self.access_key_id = config.get("access_key_id") self.secret_access_key = config.get("secret_access_key") + self.session_token = config.get("session_token") shared_creds = config.get("credentialpath") if shared_creds: @@ -73,6 +74,8 @@ class S3Tree(BaseTree): session_opts["aws_access_key_id"] = self.access_key_id if self.secret_access_key: session_opts["aws_secret_access_key"] = self.secret_access_key + if self.session_token: + session_opts["aws_session_token"] = self.session_token session = boto3.session.Session(**session_opts) diff --git a/tests/unit/tree/test_s3.py b/tests/unit/tree/test_s3.py index 44962117f..caf776f14 100644 --- a/tests/unit/tree/test_s3.py +++ b/tests/unit/tree/test_s3.py @@ -9,6 +9,7 @@ prefix = "some/prefix" url = f"s3://{bucket_name}/{prefix}" key_id = "key-id" key_secret = "key-secret" +session_token = "session-token" @pytest.fixture(autouse=True) @@ -65,10 +66,16 @@ def test_sse_kms_key_id(dvc): def test_key_id_and_secret(dvc): tree = S3Tree( dvc, - {"url": url, "access_key_id": key_id, "secret_access_key": key_secret}, + { + "url": url, + "access_key_id": key_id, + "secret_access_key": key_secret, + "session_token": session_token, + }, ) assert tree.access_key_id == key_id assert tree.secret_access_key == key_secret + assert tree.session_token == session_token def test_get_s3_no_credentials(mocker):