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.

Author: ikonst

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

pynamodb/__init__.py

@@ -7,4 +7,4 @@
 """
 __author__ = 'Jharrod LaFon'
 __license__ = 'MIT'
-__version__ = '5.4.0'
+__version__ = '5.4.1'

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

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

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')