Add account ID to boto3 session and client

boto3/session.py

@@ -16,7 +16,11 @@
 
 import botocore.session
 from botocore.client import Config
-from botocore.exceptions import DataNotFoundError, UnknownServiceError
+from botocore.exceptions import (
+    DataNotFoundError,
+    NoCredentialsError,
+    UnknownServiceError,
+)
 
 import boto3
 import boto3.utils
@@ -44,6 +48,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__(
@@ -54,6 +60,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
@@ -74,9 +81,22 @@ def __init__(
         if profile_name is not None:
             self._session.set_config_variable('profile', profile_name)
 
-        if aws_access_key_id or aws_secret_access_key or aws_session_token:
+        creds = (
+            aws_access_key_id,
+            aws_secret_access_key,
+            aws_session_token,
+            aws_account_id,
+        )
+        if any(creds):
+            if self._account_id_set_without_credentials(
+                aws_account_id, aws_access_key_id, aws_secret_access_key
+            ):
+                raise NoCredentialsError()
             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:
@@ -224,6 +244,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.
@@ -291,6 +312,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 account id to use when creating
+            the client.  Same semantics as aws_access_key_id above.
+
         :return: Service client instance
 
         """
@@ -305,6 +330,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(
@@ -527,3 +553,12 @@ def _register_default_handlers(self):
                 event_emitter=self.events,
             ),
         )
+
+    def _account_id_set_without_credentials(
+        self, account_id, access_key, secret_key
+    ):
+        if account_id is None:
+            return False
+        elif access_key is None or secret_key is None:
+            return True
+        return False

tests/unit/test_session.py

@@ -13,7 +13,7 @@
 import pytest
 from botocore import loaders
 from botocore.client import Config
-from botocore.exceptions import UnknownServiceError
+from botocore.exceptions import NoCredentialsError, UnknownServiceError
 
 from boto3 import __version__
 from boto3.exceptions import ResourceNotExistsError
@@ -70,17 +70,47 @@ def test_credentials_can_be_set(self):
 
         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', None
+        )
+
+    def test_credentials_can_be_set_with_account_id(self):
+        bc_session = self.bc_session_cls.return_value
+
+        # Set values in constructor
+        Session(
+            aws_access_key_id='key',
+            aws_secret_access_key='secret',
+            aws_session_token='token',
+            aws_account_id='account',
+        )
+
+        assert self.bc_session_cls.called
+        assert bc_session.set_credentials.called
+        bc_session.set_credentials.assert_called_with(
+            'key', 'secret', 'token', 'account'
+        )
+
+    def test_account_id_set_without_credentials(self):
+        bc_session = self.bc_session_cls.return_value
+
+        with pytest.raises(NoCredentialsError) as e:
+            Session(aws_account_id='account_id')
+
+        assert not bc_session.set_credentials.called
+        assert 'Unable to locate credentials' in str(e.value)
 
     def test_can_get_credentials(self):
         access_key = 'foo'
         secret_key = 'bar'
         token = 'baz'
+        account_id = 'bin'
 
         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
@@ -89,12 +119,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
@@ -240,6 +272,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):