From 4c7bbca95f9cf8baa1d48b42df6b3ffd2ad101ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bence=20H=C3=A9zs=C5=91?= Date: Thu, 23 Jul 2026 15:59:16 +0200 Subject: [PATCH] fix aws pagination throttle bug --- core/utils_aws.py | 78 +++------------- core/utils_egress_aws.py | 26 ++++-- tests/test_utils_aws.py | 189 ++++++++++++++++----------------------- 3 files changed, 107 insertions(+), 186 deletions(-) diff --git a/core/utils_aws.py b/core/utils_aws.py index 9681837..e37bfd4 100644 --- a/core/utils_aws.py +++ b/core/utils_aws.py @@ -1,81 +1,30 @@ # core/utils_aws.py import boto3 -import botocore import json import os -import time import logging import sqlite3 -from collections.abc import Callable from typing import Any from datetime import date, datetime, timezone from collections import defaultdict from dateutil.relativedelta import relativedelta -from botocore.exceptions import NoCredentialsError, ClientError +from botocore.config import Config from .utils_db import connect, load_data logger = logging.getLogger("core.engine.aws") - -def _retry_call( - func: Callable[..., Any], - max_retries: int, - retry_delay: int, - *args: Any, - **kwargs: Any, -) -> Any: - for attempt in range(max_retries): - try: - return func(*args, **kwargs) - except botocore.exceptions.ClientError as error: - error_code = error.response["Error"]["Code"] - if error_code in ["Throttling", "RequestLimitExceeded"]: - time.sleep(retry_delay * (2**attempt)) - continue - else: - raise - except botocore.exceptions.BotoCoreError: - time.sleep(retry_delay * (2**attempt)) - continue - raise Exception( - f"Failed to call {getattr(func, '__name__', repr(func))} after {max_retries} attempts" - ) - - -def aws_api_call_with_retry( - client: Any, - function_name: str, - parameters: dict[str, Any], - max_retries: int, - retry_delay: int, -) -> Callable[..., Any]: - function_to_call = getattr(client, function_name) - - def api_call(*args, **kwargs): - merged = {**parameters, **kwargs} if parameters else kwargs - return _retry_call(function_to_call, max_retries, retry_delay, **merged) - - return api_call +AWS_RETRY_CONFIG = Config(retries={"mode": "adaptive", "max_attempts": 8}) def paginate( client: Any, operation_name: str, result_key: str, - max_retries: int = 3, - retry_delay: int = 2, **kwargs: Any, ) -> list: - """Iterate all pages of a boto3 operation, retrying throttled page fetches.""" - paginator = client.get_paginator(operation_name) - page_iter = iter(paginator.paginate(**kwargs)) items: list = [] - while True: - try: - page = _retry_call(next, max_retries, retry_delay, page_iter) - except StopIteration: - break + for page in client.get_paginator(operation_name).paginate(**kwargs): items.extend(page.get(result_key, [])) return items @@ -84,19 +33,12 @@ def paginate_or_call( client: Any, operation_name: str, result_key: str, - max_retries: int = 3, - retry_delay: int = 2, **kwargs: Any, ) -> list: """paginate() when boto3 supports it for this operation, else a single call.""" if client.can_paginate(operation_name): - return paginate( - client, operation_name, result_key, max_retries, retry_delay, **kwargs - ) - api_call = aws_api_call_with_retry( - client, operation_name, kwargs, max_retries, retry_delay - ) - response = api_call() + return paginate(client, operation_name, result_key, **kwargs) + response = getattr(client, operation_name)(**kwargs) if not isinstance(response, dict): return [] return response.get(result_key, []) @@ -162,7 +104,9 @@ def build_aws_resource_inventory( # logger.info(f"Processing service {service_name} with operation {operation_name}") try: - client = session.client(service_name, region_name=region) + client = session.client( + service_name, region_name=region, config=AWS_RETRY_CONFIG + ) if not hasattr(client, operation_name): # logger.error(f"Operation {operation_name} does not exist for service {service_name}") continue @@ -182,7 +126,7 @@ def build_aws_resource_inventory( } ) - except (NoCredentialsError, ClientError, Exception): + except Exception: # logger.error(f"Error while processing {service_name}", exc_info=True) continue @@ -264,7 +208,9 @@ def build_aws_cost_inventory( aws_session_token=provider_details.get("sessionToken"), region_name=provider_details["region"], ) - cost_explorer = session.client("ce", region_name="us-east-1") + cost_explorer = session.client( + "ce", region_name="us-east-1", config=AWS_RETRY_CONFIG + ) db_path = os.path.join(report_path, "data", "assessment.db") diff --git a/core/utils_egress_aws.py b/core/utils_egress_aws.py index dc9010e..ef25ef9 100644 --- a/core/utils_egress_aws.py +++ b/core/utils_egress_aws.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta, timezone from botocore.exceptions import BotoCoreError, ClientError -from .utils_aws import paginate +from .utils_aws import AWS_RETRY_CONFIG, paginate from .utils_egress import GIB, format_bytes, new_row logger = logging.getLogger("core.engine.egress.aws") @@ -165,8 +165,10 @@ def _list_buckets_in_region(s3_client: Any, region: str) -> list[str]: def _collect_s3_buckets( session: Any, region: str, code: str, entry: dict[str, Any] ) -> list[dict[str, Any]]: - s3_client = session.client("s3", region_name=region) - cloudwatch = session.client("cloudwatch", region_name=region) + s3_client = session.client("s3", region_name=region, config=AWS_RETRY_CONFIG) + cloudwatch = session.client( + "cloudwatch", region_name=region, config=AWS_RETRY_CONFIG + ) bucket_names = _list_buckets_in_region(s3_client, region) @@ -245,7 +247,7 @@ def _collect_s3_buckets( def _collect_ebs_volumes( session: Any, region: str, code: str, entry: dict[str, Any] ) -> list[dict[str, Any]]: - ec2_client = session.client("ec2", region_name=region) + ec2_client = session.client("ec2", region_name=region, config=AWS_RETRY_CONFIG) rows = [] for volume in paginate(ec2_client, "describe_volumes", "Volumes"): name = next( @@ -266,7 +268,7 @@ def _collect_ebs_volumes( def _collect_ebs_snapshots( session: Any, region: str, code: str, entry: dict[str, Any] ) -> list[dict[str, Any]]: - ec2_client = session.client("ec2", region_name=region) + ec2_client = session.client("ec2", region_name=region, config=AWS_RETRY_CONFIG) rows = [] for snapshot in paginate( ec2_client, "describe_snapshots", "Snapshots", OwnerIds=["self"] @@ -292,8 +294,10 @@ def _collect_ebs_snapshots( def _collect_rds_instances( session: Any, region: str, code: str, entry: dict[str, Any] ) -> list[dict[str, Any]]: - rds_client = session.client("rds", region_name=region) - cloudwatch = session.client("cloudwatch", region_name=region) + rds_client = session.client("rds", region_name=region, config=AWS_RETRY_CONFIG) + cloudwatch = session.client( + "cloudwatch", region_name=region, config=AWS_RETRY_CONFIG + ) rows = [] specs = [] @@ -344,7 +348,9 @@ def _collect_rds_instances( def _collect_dynamodb_tables( session: Any, region: str, code: str, entry: dict[str, Any] ) -> list[dict[str, Any]]: - dynamodb_client = session.client("dynamodb", region_name=region) + dynamodb_client = session.client( + "dynamodb", region_name=region, config=AWS_RETRY_CONFIG + ) rows = [] for table_name in paginate(dynamodb_client, "list_tables", "TableNames"): try: @@ -380,7 +386,9 @@ def _collect_dynamodb_tables( def _collect_backup_vaults( session: Any, region: str, code: str, entry: dict[str, Any] ) -> list[dict[str, Any]]: - backup_client = session.client("backup", region_name=region) + backup_client = session.client( + "backup", region_name=region, config=AWS_RETRY_CONFIG + ) rows = [] for vault in paginate(backup_client, "list_backup_vaults", "BackupVaultList"): vault_name = vault["BackupVaultName"] diff --git a/tests/test_utils_aws.py b/tests/test_utils_aws.py index c7b5614..3c3eb50 100644 --- a/tests/test_utils_aws.py +++ b/tests/test_utils_aws.py @@ -8,9 +8,10 @@ import botocore.exceptions from core.utils_aws import ( - aws_api_call_with_retry, convert_datetime, get_missing_months_aws, + paginate, + paginate_or_call, ) @@ -83,116 +84,6 @@ def test_returns_all_when_none_processed(self, mock_dt): self.assertEqual(len(missing), 6) -class AwsApiCallWithRetryTests(unittest.TestCase): - def test_successful_call_returns_result(self): - mock_client = MagicMock() - mock_client.describe_instances.return_value = {"Reservations": []} - - api_call = aws_api_call_with_retry( - mock_client, "describe_instances", {}, max_retries=3, retry_delay=0 - ) - result = api_call() - - self.assertEqual(result, {"Reservations": []}) - mock_client.describe_instances.assert_called_once() - - def test_passes_parameters_to_function(self): - mock_client = MagicMock() - mock_client.list_buckets.return_value = {"Buckets": []} - - params = {"MaxItems": 10} - api_call = aws_api_call_with_retry( - mock_client, "list_buckets", params, max_retries=1, retry_delay=0 - ) - api_call() - - mock_client.list_buckets.assert_called_once_with(MaxItems=10) - - def test_passes_kwargs_from_caller(self): - mock_client = MagicMock() - mock_client.describe_instances.return_value = {"Reservations": []} - - api_call = aws_api_call_with_retry( - mock_client, "describe_instances", {}, max_retries=1, retry_delay=0 - ) - api_call(NextToken="abc123") - - mock_client.describe_instances.assert_called_once_with(NextToken="abc123") - - @patch("core.utils_aws.time.sleep") - def test_retries_on_throttling(self, mock_sleep): - mock_client = MagicMock() - throttle_error = botocore.exceptions.ClientError( - {"Error": {"Code": "Throttling", "Message": "Rate exceeded"}}, - "DescribeInstances", - ) - mock_client.describe_instances.side_effect = [ - throttle_error, - {"Reservations": ["instance-1"]}, - ] - - api_call = aws_api_call_with_retry( - mock_client, "describe_instances", {}, max_retries=3, retry_delay=1 - ) - result = api_call() - - self.assertEqual(result, {"Reservations": ["instance-1"]}) - self.assertEqual(mock_client.describe_instances.call_count, 2) - mock_sleep.assert_called_once() - - def test_raises_non_throttling_client_error_immediately(self): - mock_client = MagicMock() - access_denied = botocore.exceptions.ClientError( - {"Error": {"Code": "AccessDenied", "Message": "Not authorized"}}, - "DescribeInstances", - ) - mock_client.describe_instances.side_effect = access_denied - - api_call = aws_api_call_with_retry( - mock_client, "describe_instances", {}, max_retries=3, retry_delay=0 - ) - - with self.assertRaises(botocore.exceptions.ClientError): - api_call() - - mock_client.describe_instances.assert_called_once() - - @patch("core.utils_aws.time.sleep") - def test_raises_after_max_retries_exhausted(self, mock_sleep): - mock_client = MagicMock() - throttle_error = botocore.exceptions.ClientError( - {"Error": {"Code": "Throttling", "Message": "Rate exceeded"}}, - "DescribeInstances", - ) - mock_client.describe_instances.side_effect = throttle_error - - api_call = aws_api_call_with_retry( - mock_client, "describe_instances", {}, max_retries=2, retry_delay=0 - ) - - with self.assertRaises(Exception) as ctx: - api_call() - - self.assertIn("Failed to call", str(ctx.exception)) - self.assertEqual(mock_client.describe_instances.call_count, 2) - - @patch("core.utils_aws.time.sleep") - def test_retries_on_botocore_error(self, mock_sleep): - mock_client = MagicMock() - mock_client.describe_instances.side_effect = [ - botocore.exceptions.BotoCoreError(), - {"Reservations": []}, - ] - - api_call = aws_api_call_with_retry( - mock_client, "describe_instances", {}, max_retries=3, retry_delay=1 - ) - result = api_call() - - self.assertEqual(result, {"Reservations": []}) - self.assertEqual(mock_client.describe_instances.call_count, 2) - - class BuildAwsCostInventoryErrorTests(unittest.TestCase): @patch("core.utils_aws.connect") @patch("core.utils_aws.boto3.Session") @@ -315,5 +206,81 @@ def test_outer_exception_is_logged_silently(self, mock_session_cls, mock_load_da ) +class PaginateTests(unittest.TestCase): + def _fake_client(self, pages): + """Build a stub client whose paginator yields the given pages.""" + paginator = MagicMock() + paginator.paginate.return_value = iter(pages) + client = MagicMock() + client.get_paginator.return_value = paginator + return client + + def test_collects_items_across_pages(self): + client = self._fake_client( + [{"Items": [1, 2, 3]}, {"Items": [4, 5]}, {"Items": [6]}] + ) + self.assertEqual(paginate(client, "any_op", "Items"), [1, 2, 3, 4, 5, 6]) + + def test_missing_result_key_treated_as_empty(self): + client = self._fake_client([{"Items": [1]}, {}]) + self.assertEqual(paginate(client, "any_op", "Items"), [1]) + + def test_forwards_kwargs_to_paginator(self): + client = self._fake_client([{"Items": []}]) + paginate(client, "any_op", "Items", MaxResults=50) + client.get_paginator.return_value.paginate.assert_called_once_with( + MaxResults=50 + ) + + def test_throttling_mid_pagination_does_not_silently_truncate(self): + """Regression: retrying next() on a dead PageIterator generator + used to return a truncated prefix as if it were the full result. + With retries pushed down to the client (AWS_RETRY_CONFIG), the + paginator generator never sees the throttle and delivers every page. + Here we assert paginate() does not swallow a mid-iteration error.""" + + def flaky_pages(): + yield {"Items": ["r1", "r2"]} + raise botocore.exceptions.ClientError( + {"Error": {"Code": "Throttling", "Message": "slow down"}}, + "ListSomething", + ) + + paginator = MagicMock() + paginator.paginate.return_value = flaky_pages() + client = MagicMock() + client.get_paginator.return_value = paginator + + with self.assertRaises(botocore.exceptions.ClientError): + paginate(client, "any_op", "Items") + + +class PaginateOrCallTests(unittest.TestCase): + def test_uses_paginator_when_available(self): + paginator = MagicMock() + paginator.paginate.return_value = iter([{"Items": [1, 2]}, {"Items": [3]}]) + client = MagicMock() + client.can_paginate.return_value = True + client.get_paginator.return_value = paginator + + self.assertEqual(paginate_or_call(client, "list_things", "Items"), [1, 2, 3]) + client.can_paginate.assert_called_once_with("list_things") + + def test_falls_back_to_single_call_when_not_paginable(self): + client = MagicMock() + client.can_paginate.return_value = False + client.list_things.return_value = {"Items": ["a", "b"]} + + self.assertEqual(paginate_or_call(client, "list_things", "Items"), ["a", "b"]) + client.list_things.assert_called_once_with() + + def test_non_dict_response_returns_empty_list(self): + client = MagicMock() + client.can_paginate.return_value = False + client.list_things.return_value = None + + self.assertEqual(paginate_or_call(client, "list_things", "Items"), []) + + if __name__ == "__main__": unittest.main()