Use model's AWS credentials in threads (pynamodb#1160)

atsuoishimoto · web-flow · commit 78a0faf0c4a2 · 2023-02-20T19:15:18.000-08:00
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.
diff --git 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
@@ -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
@@ -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.connection.base import MetaTable
@@ -38,7 +39,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')
