From 15498035aeca4e7e7c44b0ab55f5c6c365e0a207 Mon Sep 17 00:00:00 2001
From: Ben Kehoe <ben@kehoe.io>
Date: Tue, 3 May 2022 14:13:48 -0400
Subject: [PATCH] Add Session.assume_role method

---
 boto3/__init__.py |  9 +++++++++
 boto3/session.py  | 50 ++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/boto3/__init__.py b/boto3/__init__.py
index 0cf09b7149..a2a95df1de 100644
--- a/boto3/__init__.py
+++ b/boto3/__init__.py
@@ -101,6 +101,15 @@ def resource(*args, **kwargs):
     return _get_default_session().resource(*args, **kwargs)
 
 
+def assume_role(*args, **kwargs):
+    """
+    Create a new session that assumes the given role.
+
+    See :py:meth:`boto3.session.Session.assume_role`.
+    """
+    return _get_default_session().assume_role(*args, **kwargs)
+
+
 # Set up logging to ``/dev/null`` like a library is supposed to.
 # https://docs.python.org/3.3/howto/logging.html#configuring-logging-for-a-library
 class NullHandler(logging.Handler):
diff --git a/boto3/session.py b/boto3/session.py
index bdda65ad41..2c8bb1b411 100644
--- a/boto3/session.py
+++ b/boto3/session.py
@@ -16,7 +16,8 @@
 
 import botocore.session
 from botocore.client import Config
-from botocore.exceptions import DataNotFoundError, UnknownServiceError
+from botocore.exceptions import DataNotFoundError, UnknownServiceError, NoCredentialsError
+from botocore.credentials import AssumeRoleCredentialFetcher, DeferredRefreshableCredentials
 
 import boto3
 import boto3.utils
@@ -477,6 +478,53 @@ def resource(
 
         return cls(client=client)
 
+    def assume_role(self, role_arn, extra_args=None):
+        """
+        Create a new session that assumes the given role.
+
+        :type role_arn: str
+        :param role_arn: The ARN of the role to be assumed.
+        :type extra_args: dict
+        :param extra_args: Any additional arguments to add to the assume
+            role request using the format of the botocore operation.
+            Possible keys include, but may not be limited to,
+            DurationSeconds, Policy, and RoleSessionName.
+        """
+        botocore_session = self._session
+
+        credentials = botocore_session.get_credentials()
+        if not credentials:
+            # Error out now rather than wait for the new session to get used
+            raise NoCredentialsError
+
+        credential_fetcher = AssumeRoleCredentialFetcher(
+            botocore_session.create_client,
+            credentials,
+            role_arn,
+            extra_args=extra_args
+        )
+
+        assumed_role_credentials = DeferredRefreshableCredentials(
+            credential_fetcher.fetch_credentials,
+            "assume-role"
+        )
+
+        assumed_role_botocore_session = botocore.session.get_session()
+        assumed_role_botocore_session._credentials = assumed_role_credentials
+
+        # note that if this session's region changes, it will not cascade
+        region_name = self.region_name
+
+        assumed_role_boto3_session = Session(
+            botocore_session=assumed_role_botocore_session,
+            region_name=region_name,
+        )
+
+        # provice traceability
+        assumed_role_boto3_session.assume_role_parent_session = self
+
+        return assumed_role_boto3_session
+
     def _register_default_handlers(self):
 
         # S3 customizations