s3: Allow authentication using session token (#5240)

* S3: Allow authentication using session token

* Add a little test

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,

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)
 

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):