From a8900c40a9533d70d3e8e7a33853b6dd9d62ba33 Mon Sep 17 00:00:00 2001
From: Ilya Konstantinov <ilya.konstantinov@gmail.com>
Date: Mon, 20 Feb 2023 19:55:10 -0800
Subject: [PATCH] Use model's AWS credentials in threads

When a model specifies custom AWS credentials (instead of global ones), they should be used when creating new sessions in threads. Previously, threads would always use the global credentials.
---
 docs/release_notes.rst         | 13 +++++++++++++
 pynamodb/__init__.py           |  2 +-
 pynamodb/connection/base.py    | 13 ++++++++++++-
 pynamodb/connection/table.py   | 11 +++++------
 tests/test_table_connection.py | 12 +++++++++++-
 5 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/docs/release_notes.rst b/docs/release_notes.rst
index cdc58a81f..d90ca2d58 100644
--- a/docs/release_notes.rst
+++ b/docs/release_notes.rst
@@ -3,6 +3,19 @@
 Release Notes
 =============
 
+v5.4.1
+----------
+* Use model's AWS credentials in threads (#1164)
+
+  A model can specify custom AWS credentials in the ``Meta`` class (in lieu of "global"
+  AWS credentials from the environment). Previously those model-specific credentials
+  were not used from within new threads.
+
+Contributors to this release:
+
+* @atsuoishimoto
+
+
 v5.4.0
 ----------
 * Expose transaction cancellation reasons in
diff --git a/pynamodb/__init__.py b/pynamodb/__init__.py
index 195411627..41fa4926c 100644
--- a/pynamodb/__init__.py
+++ b/pynamodb/__init__.py
@@ -7,4 +7,4 @@
 """
 __author__ = 'Jharrod LaFon'
 __license__ = 'MIT'
-__version__ = '5.4.0'
+__version__ = '5.4.1'
diff --git a/pynamodb/connection/base.py b/pynamodb/connection/base.py
index b06aec88c..62b0d9488 100644
--- a/pynamodb/connection/base.py
+++ b/pynamodb/connection/base.py
@@ -257,7 +257,10 @@ def __init__(self,
                  max_retry_attempts: Optional[int] = None,
                  base_backoff_ms: Optional[int] = None,
                  max_pool_connections: Optional[int] = None,
-                 extra_headers: Optional[Mapping[str, str]] = None):
+                 extra_headers: Optional[Mapping[str, str]] = None,
+                 aws_access_key_id: Optional[str] = None,
+                 aws_secret_access_key: Optional[str] = None,
+                 aws_session_token: Optional[str] = None):
         self._tables: Dict[str, MetaTable] = {}
         self.host = host
         self._local = local()
@@ -298,6 +301,10 @@ def __init__(self,
         else:
             self._extra_headers = get_settings_value('extra_headers')
 
+        self._aws_access_key_id = aws_access_key_id
+        self._aws_secret_access_key = aws_secret_access_key
+        self._aws_session_token = aws_session_token
+
     def __repr__(self) -> str:
         return "Connection<{}>".format(self.client.meta.endpoint_url)
 
@@ -558,6 +565,10 @@ def session(self) -> botocore.session.Session:
         # botocore client creation is not thread safe as of v1.2.5+ (see issue #153)
         if getattr(self._local, 'session', None) is None:
             self._local.session = get_session()
+            if self._aws_access_key_id and self._aws_secret_access_key:
+                self._local.session.set_credentials(self._aws_access_key_id,
+                                                        self._aws_secret_access_key,
+                                                        self._aws_session_token)
         return self._local.session
 
     @property
diff --git a/pynamodb/connection/table.py b/pynamodb/connection/table.py
index fb7720e00..22d16cad2 100644
--- a/pynamodb/connection/table.py
+++ b/pynamodb/connection/table.py
@@ -41,15 +41,14 @@ def __init__(
                                      max_retry_attempts=max_retry_attempts,
                                      base_backoff_ms=base_backoff_ms,
                                      max_pool_connections=max_pool_connections,
-                                     extra_headers=extra_headers)
+                                     extra_headers=extra_headers,
+                                     aws_access_key_id=aws_access_key_id,
+                                     aws_secret_access_key=aws_secret_access_key,
+                                     aws_session_token=aws_session_token)
+
         if meta_table is not None:
             self.connection.add_meta_table(meta_table)
 
-        if aws_access_key_id and aws_secret_access_key:
-            self.connection.session.set_credentials(aws_access_key_id,
-                                                    aws_secret_access_key,
-                                                    aws_session_token)
-
     def get_meta_table(self) -> MetaTable:
         """
         Returns a MetaTable
diff --git a/tests/test_table_connection.py b/tests/test_table_connection.py
index 87e955fd6..c26bef24b 100644
--- a/tests/test_table_connection.py
+++ b/tests/test_table_connection.py
@@ -2,6 +2,7 @@
 Test suite for the table class
 """
 from unittest import TestCase
+from concurrent.futures import ThreadPoolExecutor
 
 from pynamodb.connection import TableConnection
 from pynamodb.constants import DEFAULT_REGION
@@ -39,7 +40,16 @@ def test_connection_session_set_credentials(self):
             aws_access_key_id='access_key_id',
             aws_secret_access_key='secret_access_key')
 
-        credentials = conn.connection.session.get_credentials()
+        def get_credentials():
+            return conn.connection.session.get_credentials()
+
+        credentials = get_credentials()
+        self.assertEqual(credentials.access_key, 'access_key_id')
+        self.assertEqual(credentials.secret_key, 'secret_access_key')
+
+        with ThreadPoolExecutor() as executor:
+            fut = executor.submit(get_credentials)
+            credentials = fut.result()
 
         self.assertEqual(credentials.access_key, 'access_key_id')
         self.assertEqual(credentials.secret_key, 'secret_access_key')