Merge pull request #167 from Skyscanner/fix-throttling-on-aws-functions

Add handling for AWS throttling when listing exports

CHANGELOG.md

@@ -1,6 +1,11 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+## [1.0.2] - 2021-03-25
+### Improvements
+- Handle AWS throttling errors when listing exports for a given account and region
+- If we get a throttling error, we actually sleep for some time before retrying (before we were sleeping for 0 seconds)
+
 ## [1.0.1] - 2021-03-25
 ### Improvements
 - Decrease logging level when loading external filters

cfripper/__version__.py

@@ -1,3 +1,3 @@
-VERSION = (1, 0, 1)
+VERSION = (1, 0, 2)
 
 __version__ = ".".join(map(str, VERSION))

cfripper/boto3_client.py

@@ -47,13 +47,14 @@ def get_template(self) -> Optional[Dict]:
                     logger.warning(
                         f"No template body found for stack: {self.stack_id} on {self.account_id} - {self.region}"
                     )
-                    sleep(i)
+                    sleep((i + 1) * 2)
             except ClientError as e:
                 if e.response["Error"]["Code"] == "ValidationError":
                     logger.exception(f"There is no stack: {self.stack_id} on {self.account_id} - {self.region}")
+                    return stack_content
                 elif e.response["Error"]["Code"] == "Throttling":
                     logger.warning(f"AWS Throttling: {self.stack_id} on {self.account_id} - {self.region}")
-                    sleep(i)
+                    sleep((i + 1) * 2)
                 else:
                     logger.exception(
                         "Unexpected error occurred when getting stack template for:"
@@ -82,15 +83,25 @@ def download_template_to_dictionary(self, s3_url):
 
     def get_exports(self) -> Dict[str, str]:
         client = self.session.client("cloudformation", region_name=self.region)
-        try:
-            return {export["Name"]: export["Value"] for export in client.list_exports()["Exports"]}
-        except ClientError as e:
-            if e.response["Error"]["Code"] == "AccessDenied":
-                logger.warning(f"Access Denied for obtaining AWS Export values! ({self.account_id} - {self.region})")
-            else:
-                logger.exception(
-                    f"Unhandled ClientError getting AWS Export values! ({self.account_id} - {self.region})"
-                )
-        except Exception:
-            logger.exception(f"Unknown exception getting AWS Export values! ({self.account_id} - {self.region})")
-        return {}
+        export_values = {}
+        i = 0
+        while not export_values and i < self.N_RETRIES:
+            try:
+                export_values = {export["Name"]: export["Value"] for export in client.list_exports().get("Exports", [])}
+            except ClientError as e:
+                if e.response["Error"]["Code"] == "AccessDenied":
+                    logger.warning(
+                        f"Access Denied for obtaining AWS Export values! ({self.account_id} - {self.region})"
+                    )
+                    return export_values
+                elif e.response["Error"]["Code"] == "Throttling":
+                    logger.warning(f"AWS Throttling: {self.stack_id} on {self.account_id} - {self.region}")
+                    sleep((i + 1) * 2)
+                else:
+                    logger.exception(
+                        f"Unhandled ClientError getting AWS Export values! ({self.account_id} - {self.region})"
+                    )
+            except Exception:
+                logger.exception(f"Unknown exception getting AWS Export values! ({self.account_id} - {self.region})")
+            i += 1
+        return export_values

tests/test_boto3_client.py

@@ -9,6 +9,7 @@
 from cfripper.boto3_client import Boto3Client
 from cfripper.model.utils import InvalidURLException, convert_json_or_yaml_to_dict
 
+CLIENT_ERROR_ACCESS_DENIED = ClientError({"Error": {"Code": "AccessDenied"}}, "list_exports")
 CLIENT_ERROR_THROTTLING = ClientError({"Error": {"Code": "Throttling"}}, "get_template")
 CLIENT_ERROR_VALIDATION = ClientError({"Error": {"Code": "ValidationError"}}, "get_template")
 DUMMY_CLIENT_ERROR = ClientError({"Error": {"Code": "Exception"}}, "get_template")
@@ -41,11 +42,11 @@ def boto3_client():
             [],
         ),
         (
-            [CLIENT_ERROR_VALIDATION] * 10,
+            [CLIENT_ERROR_VALIDATION],
             None,
-            [call(f"Stack: stack-id on 123456789 - eu-west-1 get_template Attempt #{i}") for i in range(5)],
+            [call("Stack: stack-id on 123456789 - eu-west-1 get_template Attempt #0")],
             [],
-            [call("There is no stack: stack-id on 123456789 - eu-west-1") for _ in range(5)],
+            [call("There is no stack: stack-id on 123456789 - eu-west-1")],
         ),
         (
             [CLIENT_ERROR_THROTTLING, {"A": "a"}],
@@ -237,6 +238,55 @@ def test_urlencoded_url(s3_bucket, boto3_client):
     assert result["hello"] == "this is valid json"
 
 
+@pytest.mark.parametrize(
+    "aws_responses, expected_exports, mocked_warning_logs, mocked_exceptions",
+    [
+        ([[{"Name": "A", "Value": "a"}]], {"A": "a"}, [], [],),
+        (
+            [[{"Foo": "Bar"}], [{"Name": "A", "Value": "a"}]],
+            {"A": "a"},
+            [],
+            [call("Unknown exception getting AWS Export values! (123456789 - eu-west-1)")],
+        ),
+        (
+            [CLIENT_ERROR_ACCESS_DENIED],
+            {},
+            [call("Access Denied for obtaining AWS Export values! (123456789 - eu-west-1)")],
+            [],
+        ),
+        (
+            [CLIENT_ERROR_THROTTLING, [{"Name": "A", "Value": "a"}]],
+            {"A": "a"},
+            [call("AWS Throttling: stack-id on 123456789 - eu-west-1")],
+            [],
+        ),
+        (
+            [DUMMY_CLIENT_ERROR, [{"Name": "A", "Value": "a"}]],
+            {"A": "a"},
+            [],
+            [call("Unhandled ClientError getting AWS Export values! (123456789 - eu-west-1)")],
+        ),
+    ],
+)
+@patch("logging.Logger.warning")
+@patch("logging.Logger.exception")
+def test_get_exports(
+    patched_exceptions,
+    patched_logger_warning,
+    aws_responses,
+    expected_exports,
+    mocked_warning_logs,
+    mocked_exceptions,
+    boto3_client,
+):
+    with patch.object(boto3_client, "session") as session_mock:
+        session_mock.client().list_exports().get.side_effect = aws_responses
+        exports = boto3_client.get_exports()
+        assert exports == expected_exports
+    assert patched_logger_warning.mock_calls == mocked_warning_logs
+    assert patched_exceptions.mock_calls == mocked_exceptions
+
+
 @mock_cloudformation
 def test_export_values(boto3_client: Boto3Client):
     cf_client = boto3_client.session.client("cloudformation", "eu-west-1")