From ab638f050f9bad90ef50a0ed628365459acf33af Mon Sep 17 00:00:00 2001
From: davidlm <davidlm@amazon.com>
Date: Wed, 4 Oct 2023 17:10:53 -0400
Subject: [PATCH] add account_id to boto3 session

---
 boto3/session.py           | 20 +++++++++++++++++++-
 tests/unit/test_session.py | 13 ++++++++++++-
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/boto3/session.py b/boto3/session.py
index bdda65ad41..fbcd49cb7b 100644
--- a/boto3/session.py
+++ b/boto3/session.py
@@ -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__(
@@ -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
@@ -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:
@@ -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.
@@ -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
 
         """
@@ -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(
@@ -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.
@@ -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:
@@ -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
 
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index 728feb33c7..8592f5b4b5 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -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
@@ -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
@@ -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):
@@ -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'
@@ -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'
@@ -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'