Adds support for loading AWS bearer tokens via session constructor params

Author: jacoblee93

.changes/next-version.json

@@ -0,0 +1,7 @@
+[
+  {
+    "category": "Session",
+    "description": "Add support for ``aws_bearer_token`` parameter in Session, Client, and Resource to enable per-session/client bearer token authentication for services that support it (e.g. Bedrock). This allows safe multi-tenant usage without relying on global environment variables like ``AWS_BEARER_TOKEN_BEDROCK``.",
+    "type": "feature"
+  }
+]

boto3/session.py

@@ -50,6 +50,9 @@ class Session:
                          the default profile is used.
     :type aws_account_id: string
     :param aws_account_id: AWS account ID
+    :type aws_bearer_token: string
+    :param aws_bearer_token: Bearer token for authentication with services
+                             that support bearer token auth (e.g. Bedrock).
     """
 
     def __init__(
@@ -61,6 +64,7 @@ def __init__(
         botocore_session=None,
         profile_name=None,
         aws_account_id=None,
+        aws_bearer_token=None,
     ):
         if botocore_session is not None:
             self._session = botocore_session
@@ -100,6 +104,9 @@ def __init__(
         if region_name is not None:
             self._session.set_config_variable('region', region_name)
 
+        # Store the bearer token for per-session bearer token authentication
+        self._aws_bearer_token = aws_bearer_token
+
         self.resource_factory = ResourceFactory(
             self._session.get_component('event_emitter')
         )
@@ -243,6 +250,7 @@ def client(
         aws_session_token=None,
         config=None,
         aws_account_id=None,
+        aws_bearer_token=None,
     ):
         """
         Create a low-level service client by name.
@@ -314,9 +322,25 @@ def client(
         :param aws_account_id: The account id to use when creating
             the client.  Same semantics as aws_access_key_id above.
 
+        :type aws_bearer_token: string
+        :param aws_bearer_token: The bearer token to use for
+            authentication with services that support bearer token auth
+            (e.g. Bedrock). If provided, this takes precedence over any
+            bearer token resolved from environment variables and any
+            session-level token. Same semantics as aws_access_key_id
+            above.
+
         :return: Service client instance
 
         """
+        # Determine the effective bearer token to use
+        # Precedence: client parameter > session parameter > environment variable
+        effective_aws_bearer_token = (
+            aws_bearer_token
+            if aws_bearer_token is not None
+            else self._aws_bearer_token
+        )
+
         create_client_kwargs = {
             'region_name': region_name,
             'api_version': api_version,
@@ -328,11 +352,16 @@ def client(
             'aws_session_token': aws_session_token,
             'config': config,
             'aws_account_id': aws_account_id,
+            'aws_bearer_token': effective_aws_bearer_token,
         }
         if aws_account_id is None:
             # Remove aws_account_id for arbitrary
             # botocore version mismatches in AWS Lambda.
             del create_client_kwargs['aws_account_id']
+        if effective_aws_bearer_token is None:
+            # Remove aws_bearer_token for arbitrary
+            # botocore version mismatches in AWS Lambda.
+            del create_client_kwargs['aws_bearer_token']
 
         return self._session.create_client(
             service_name, **create_client_kwargs
@@ -350,6 +379,7 @@ def resource(
         aws_secret_access_key=None,
         aws_session_token=None,
         config=None,
+        aws_bearer_token=None,
     ):
         """
         Create a resource service client by name.
@@ -419,6 +449,14 @@ def resource(
             <https://docs.aws.amazon.com/botocore/latest/reference/config.html>`_
             for more details.
 
+        :type aws_bearer_token: string
+        :param aws_bearer_token: The bearer token to use for
+            authentication with services that support bearer token auth
+            (e.g. Bedrock). If provided, this takes precedence over any
+            bearer token resolved from environment variables and any
+            session-level token. Same semantics as aws_access_key_id
+            above.
+
         :return: Subclass of :py:class:`~boto3.resources.base.ServiceResource`
         """
         try:
@@ -483,6 +521,7 @@ def resource(
             aws_secret_access_key=aws_secret_access_key,
             aws_session_token=aws_session_token,
             config=config,
+            aws_bearer_token=aws_bearer_token,
         )
         service_model = client.meta.service_model
 

tests/unit/test_session.py

@@ -272,6 +272,116 @@ def test_create_client_with_args(self):
             config=None,
         )
 
+    def test_session_with_aws_bearer_token(self):
+        session = Session(aws_bearer_token='test-token')
+        assert session._aws_bearer_token == 'test-token'
+
+    def test_session_with_aws_bearer_token_none(self):
+        session = Session()
+        assert session._aws_bearer_token is None
+
+    def test_create_client_with_aws_bearer_token(self):
+        bc_session = self.bc_session_cls.return_value
+
+        session = Session(region_name='us-east-1')
+        session.client(
+            'bedrock-runtime',
+            region_name='us-west-2',
+            aws_bearer_token='test-bearer-token',
+        )
+
+        bc_session.create_client.assert_called_with(
+            'bedrock-runtime',
+            aws_secret_access_key=None,
+            aws_access_key_id=None,
+            endpoint_url=None,
+            use_ssl=True,
+            aws_session_token=None,
+            verify=None,
+            region_name='us-west-2',
+            api_version=None,
+            config=None,
+            aws_bearer_token='test-bearer-token',
+        )
+
+    def test_create_client_with_session_level_aws_bearer_token(self):
+        bc_session = self.bc_session_cls.return_value
+
+        session = Session(
+            region_name='us-east-1',
+            aws_bearer_token='session-level-token',
+        )
+        session.client('bedrock-runtime', region_name='us-west-2')
+
+        bc_session.create_client.assert_called_with(
+            'bedrock-runtime',
+            aws_secret_access_key=None,
+            aws_access_key_id=None,
+            endpoint_url=None,
+            use_ssl=True,
+            aws_session_token=None,
+            verify=None,
+            region_name='us-west-2',
+            api_version=None,
+            config=None,
+            aws_bearer_token='session-level-token',
+        )
+
+    def test_create_client_aws_bearer_token_client_overrides_session(self):
+        bc_session = self.bc_session_cls.return_value
+
+        session = Session(
+            region_name='us-east-1',
+            aws_bearer_token='session-level-token',
+        )
+        session.client(
+            'bedrock-runtime',
+            region_name='us-west-2',
+            aws_bearer_token='client-level-token',
+        )
+
+        bc_session.create_client.assert_called_with(
+            'bedrock-runtime',
+            aws_secret_access_key=None,
+            aws_access_key_id=None,
+            endpoint_url=None,
+            use_ssl=True,
+            aws_session_token=None,
+            verify=None,
+            region_name='us-west-2',
+            api_version=None,
+            config=None,
+            aws_bearer_token='client-level-token',
+        )
+
+    def test_create_client_aws_bearer_token_empty_string_overrides_session(self):
+        bc_session = self.bc_session_cls.return_value
+
+        session = Session(
+            region_name='us-east-1',
+            aws_bearer_token='session-level-token',
+        )
+        # Empty string is an explicit value and should override session token
+        session.client(
+            'bedrock-runtime',
+            region_name='us-west-2',
+            aws_bearer_token='',
+        )
+
+        bc_session.create_client.assert_called_with(
+            'bedrock-runtime',
+            aws_secret_access_key=None,
+            aws_access_key_id=None,
+            endpoint_url=None,
+            use_ssl=True,
+            aws_session_token=None,
+            verify=None,
+            region_name='us-west-2',
+            api_version=None,
+            config=None,
+            aws_bearer_token='',
+        )
+
     def test_create_client_with_aws_account_id(self):
         bc_session = self.bc_session_cls.return_value
 
@@ -324,6 +434,7 @@ def test_create_resource_with_args(self):
             region_name=None,
             api_version='2014-11-02',
             config=mock.ANY,
+            aws_bearer_token=None,
         )
         client_config = session.client.call_args[1]['config']
         assert client_config.user_agent_extra == 'Resource'
@@ -356,6 +467,7 @@ def test_create_resource_with_config(self):
             region_name=None,
             api_version='2014-11-02',
             config=mock.ANY,
+            aws_bearer_token=None,
         )
         client_config = session.client.call_args[1]['config']
         assert client_config.user_agent_extra == 'Resource'
@@ -388,6 +500,7 @@ def test_create_resource_with_config_override_user_agent_extra(self):
             region_name=None,
             api_version='2014-11-02',
             config=mock.ANY,
+            aws_bearer_token=None,
         )
         client_config = session.client.call_args[1]['config']
         assert client_config.user_agent_extra == 'foo'