fix: Cache AWS Creds instead of getting a Session on every invoke (aws#1494)

jfuss · web-flow · commit 0ce744f00f05 · 2019-12-11T11:00:43.000-06:00
Addresses Issue aws#1375. Due to Boto3.session.Session objects not being thread safe, we create a Session object on every invoke to get credentials to mount into the container. This works well except for customers using MFA. When using MFA, this dehavior requires customers to provide MFA code/pin on every invoke. This creates much longer debug sessions for customers and pauses invokes until the MFA code/pin is provided. This PR changes that behavior to one we had between version v0.8.0 and v0.16.0, which is caching the creds/session to be reused on other invokes until the credentials expire. Once the credentials expire, customers will need to restart the command (which was the same behavior on version 0.8.0 to 0.16.0.
diff --git a/samcli/commands/local/lib/local_lambda.py b/samcli/commands/local/lib/local_lambda.py
@@ -53,6 +53,8 @@ def __init__(
         self.aws_region = aws_region
         self.env_vars_values = env_vars_values or {}
         self.debug_context = debug_context
+        self._boto3_session_creds = None
+        self._boto3_region = None
 
     def invoke(self, function_name, event, stdout=None, stderr=None):
         """
@@ -205,6 +207,22 @@ def _make_env_vars(self, function):
             aws_creds=aws_creds,
         )
 
+    def _get_session_creds(self):
+        if self._boto3_session_creds is None:
+            # to pass command line arguments for region & profile to setup boto3 default session
+            LOG.debug("Loading AWS credentials from session with profile '%s'", self.aws_profile)
+            session = boto3.session.Session(profile_name=self.aws_profile, region_name=self.aws_region)
+
+            # check for region_name in session and cache
+            if hasattr(session, "region_name") and session.region_name:
+                self._boto3_region = session.region_name
+
+            # don't set cached session creds if there is not a session
+            if session:
+                self._boto3_session_creds = session.get_credentials()
+
+        return self._boto3_session_creds
+
     def get_aws_creds(self):
         """
         Returns AWS credentials obtained from the shell environment or given profile
@@ -215,26 +233,17 @@ def get_aws_creds(self):
         """
         result = {}
 
-        # to pass command line arguments for region & profile to setup boto3 default session
-        LOG.debug("Loading AWS credentials from session with profile '%s'", self.aws_profile)
-        # boto3.session.Session is not thread safe. To ensure we do not run into a race condition with start-lambda
-        # or start-api, we create the session object here on every invoke.
-        session = boto3.session.Session(profile_name=self.aws_profile, region_name=self.aws_region)
-
-        if not session:
-            return result
-
         # Load the credentials from profile/environment
-        creds = session.get_credentials()
+        creds = self._get_session_creds()
+
+        # After loading credentials, region name might be available here.
+        if self._boto3_region:
+            result["region"] = self._boto3_region
 
         if not creds:
-            # If we were unable to load credentials, then just return empty. We will use the default
+            # If we were unable to load credentials, then just return result. We will use the default
             return result
 
-        # After loading credentials, region name might be available here.
-        if hasattr(session, "region_name") and session.region_name:
-            result["region"] = session.region_name
-
         # Only add the key, if its value is present
         if hasattr(creds, "access_key") and creds.access_key:
             result["key"] = creds.access_key
diff --git a/tests/unit/commands/local/lib/test_local_lambda.py b/tests/unit/commands/local/lib/test_local_lambda.py
@@ -144,7 +144,7 @@ def test_must_work_with_no_credentials(self, boto3_mock):
         boto3_mock.session.Session.return_value = mock_session
         mock_session.get_credentials.return_value = None
 
-        expected = {}
+        expected = {"region": mock_session.region_name}
         actual = self.local_lambda.get_aws_creds()
         self.assertEqual(expected, actual)
 
