Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .changes/next-version.json
Original file line number Diff line number Diff line change
@@ -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"
}
]
39 changes: 39 additions & 0 deletions boto3/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand All @@ -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
Expand Down Expand Up @@ -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')
)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down
113 changes: 113 additions & 0 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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'
Expand Down