Skip to content

Commit

Permalink
add account_id to boto3 session
Browse files Browse the repository at this point in the history
  • Loading branch information
davidlm committed Oct 5, 2023
1 parent cbe9683 commit ab638f0
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 2 deletions.
20 changes: 19 additions & 1 deletion boto3/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@ class Session:
:type profile_name: string
:param profile_name: The name of a profile to use. If not given, then
the default profile is used.
:type aws_account_id: string
:param aws_account_id: AWS account ID
"""

def __init__(
Expand All @@ -54,6 +56,7 @@ def __init__(
region_name=None,
botocore_session=None,
profile_name=None,
aws_account_id=None,
):
if botocore_session is not None:
self._session = botocore_session
Expand All @@ -78,7 +81,10 @@ def __init__(

if aws_access_key_id or aws_secret_access_key or aws_session_token:
self._session.set_credentials(
aws_access_key_id, aws_secret_access_key, aws_session_token
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
aws_account_id,
)

if region_name is not None:
Expand Down Expand Up @@ -226,6 +232,7 @@ def client(
aws_secret_access_key=None,
aws_session_token=None,
config=None,
aws_account_id=None,
):
"""
Create a low-level service client by name.
Expand Down Expand Up @@ -293,6 +300,10 @@ def client(
<https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html>`_
for more details.
:type aws_account_id: string
:param aws_account_id: The AWS account ID to use when creating the
client. Same semantics as aws_access_key_id above.
:return: Service client instance
"""
Expand All @@ -307,6 +318,7 @@ def client(
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
config=config,
aws_account_id=aws_account_id,
)

def resource(
Expand All @@ -321,6 +333,7 @@ def resource(
aws_secret_access_key=None,
aws_session_token=None,
config=None,
aws_account_id=None,
):
"""
Create a resource service client by name.
Expand Down Expand Up @@ -390,6 +403,10 @@ def resource(
<https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html>`_
for more details.
:type aws_account_id: string
:param aws_account_id: The AWS account ID to use when creating the
client. Same semantics as aws_access_key_id above.
:return: Subclass of :py:class:`~boto3.resources.base.ServiceResource`
"""
try:
Expand Down Expand Up @@ -454,6 +471,7 @@ def resource(
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
config=config,
aws_account_id=aws_account_id,
)
service_model = client.meta.service_model

Expand Down
13 changes: 12 additions & 1 deletion tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,21 +66,26 @@ def test_credentials_can_be_set(self):
aws_access_key_id='key',
aws_secret_access_key='secret',
aws_session_token='token',
aws_account_id='account_id',
)

assert self.bc_session_cls.called
assert bc_session.set_credentials.called
bc_session.set_credentials.assert_called_with('key', 'secret', 'token')
bc_session.set_credentials.assert_called_with(
'key', 'secret', 'token', 'account_id'
)

def test_can_get_credentials(self):
access_key = 'foo'
secret_key = 'bar'
token = 'baz'
account_id = 'account_id'

creds = mock.Mock()
creds.access_key = access_key
creds.secret_key = secret_key
creds.token = token
creds.account_id = account_id

bc_session = self.bc_session_cls.return_value
bc_session.get_credentials.return_value = creds
Expand All @@ -89,12 +94,14 @@ def test_can_get_credentials(self):
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
aws_session_token=token,
aws_account_id=account_id,
)

credentials = session.get_credentials()
assert credentials.access_key == access_key
assert credentials.secret_key == secret_key
assert credentials.token == token
assert credentials.account_id == account_id

def test_profile_can_be_set(self):
bc_session = self.bc_session_cls.return_value
Expand Down Expand Up @@ -240,6 +247,7 @@ def test_create_client_with_args(self):
region_name='us-west-2',
api_version=None,
config=None,
aws_account_id=None,
)

def test_create_resource_with_args(self):
Expand Down Expand Up @@ -268,6 +276,7 @@ def test_create_resource_with_args(self):
region_name=None,
api_version='2014-11-02',
config=mock.ANY,
aws_account_id=None,
)
client_config = session.client.call_args[1]['config']
assert client_config.user_agent_extra == 'Resource'
Expand Down Expand Up @@ -300,6 +309,7 @@ def test_create_resource_with_config(self):
region_name=None,
api_version='2014-11-02',
config=mock.ANY,
aws_account_id=None,
)
client_config = session.client.call_args[1]['config']
assert client_config.user_agent_extra == 'Resource'
Expand Down Expand Up @@ -332,6 +342,7 @@ def test_create_resource_with_config_override_user_agent_extra(self):
region_name=None,
api_version='2014-11-02',
config=mock.ANY,
aws_account_id=None,
)
client_config = session.client.call_args[1]['config']
assert client_config.user_agent_extra == 'foo'
Expand Down

0 comments on commit ab638f0

Please sign in to comment.