fix(aws): retry throttled paginated calls at the client layer

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

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"]