add support for assuming a role arn with sts

Author: icholy

dvc_s3/__init__.py

@@ -219,6 +219,10 @@ def _prepare_credentials(self, **config):
             # Enable bucket region caching
             login_info["cache_regions"] = config.get("cache_regions", True)
 
+        role_arn = config.get("role_arn")
+        if role_arn:
+            self._assume_role(login_info, role_arn)
+
         config_path = config.get("configpath")
         if config_path:
             os.environ.setdefault("AWS_CONFIG_FILE", config_path)
@@ -229,6 +233,29 @@ def _prepare_credentials(self, **config):
             splitter="dot",
         )
 
+    @staticmethod
+    def _assume_role(login_info, role_arn):
+        import botocore.session
+
+        session = botocore.session.Session(
+            profile=login_info.get("profile"),
+        )
+        sts = session.create_client(
+            "sts",
+            aws_access_key_id=login_info.get("key"),
+            aws_secret_access_key=login_info.get("secret"),
+            aws_session_token=login_info.get("token"),
+        )
+        resp = sts.assume_role(
+            RoleArn=role_arn,
+            RoleSessionName="dvc",
+        )
+        creds = resp["Credentials"]
+        login_info["key"] = creds["AccessKeyId"]
+        login_info["secret"] = creds["SecretAccessKey"]
+        login_info["token"] = creds["SessionToken"]
+        login_info["profile"] = None
+
     @wrap_prop(threading.Lock())
     @cached_property
     def fs(self):

dvc_s3/tests/test_s3.py

@@ -136,3 +136,31 @@ def test_key_id_and_secret():
     assert fs.fs_args["key"] == key_id
     assert fs.fs_args["secret"] == key_secret
     assert fs.fs_args["token"] == session_token
+
+
+def test_assume_role(monkeypatch):
+    from unittest.mock import MagicMock
+
+    role_arn = "arn:aws:iam::123456789012:role/test-role"
+
+    mock_sts = MagicMock()
+    mock_sts.assume_role.return_value = {
+        "Credentials": {
+            "AccessKeyId": "assumed-key-id",
+            "SecretAccessKey": "assumed-secret",
+            "SessionToken": "assumed-session-token",
+        }
+    }
+    mock_session = MagicMock()
+    mock_session.create_client.return_value = mock_sts
+    monkeypatch.setattr("botocore.session.Session", lambda **kwargs: mock_session)
+
+    fs = S3FileSystem(url=url, role_arn=role_arn)
+    assert fs.fs_args["key"] == "assumed-key-id"
+    assert fs.fs_args["secret"] == "assumed-secret"
+    assert fs.fs_args["token"] == "assumed-session-token"
+    assert "profile" not in fs.fs_args
+
+    mock_sts.assume_role.assert_called_once_with(
+        RoleArn=role_arn, RoleSessionName="dvc"
+    )