fix(aws): paginate resource inventory using boto3 paginators

core/utils_aws.py

@@ -18,36 +18,88 @@
 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):
-        for attempt in range(max_retries):
-            try:
-                function_to_call = getattr(client, function_name)
-                if parameters:
-                    return function_to_call(**parameters, **kwargs)
-                else:
-                    return function_to_call(**kwargs)
-            except botocore.exceptions.ClientError as error:
-                error_code = error.response["Error"]["Code"]
-                # logger.warning(f"ClientError: {error_code}. Attempt {attempt + 1} of {max_retries}. Retrying in {retry_delay} seconds.")
-                if error_code in ["Throttling", "RequestLimitExceeded"]:
-                    time.sleep(retry_delay * (2**attempt))
-                    continue
-                else:
-                    raise
-            except botocore.exceptions.BotoCoreError:
-                # logger.warning(f"BotoCoreError. Attempt {attempt + 1} of {max_retries}. Retrying in {retry_delay} seconds.")
-                time.sleep(retry_delay * (2**attempt))
-                continue
-        raise Exception(f"Failed to call {function_name} after {max_retries} attempts")
+        merged = {**parameters, **kwargs} if parameters else kwargs
+        return _retry_call(function_to_call, max_retries, retry_delay, **merged)
 
-    return api_call  # Return the callable function
+    return api_call
+
+
+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
+        items.extend(page.get(result_key, []))
+    return items
+
+
+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()
+    if not isinstance(response, dict):
+        return []
+    return response.get(result_key, [])
 
 
 def convert_datetime(obj: Any) -> Any:
@@ -115,24 +167,7 @@ def build_aws_resource_inventory(
                     # logger.error(f"Operation {operation_name} does not exist for service {service_name}")
                     continue
 
-                # Make the API call
-                api_call = aws_api_call_with_retry(
-                    client, operation_name, {}, max_retries=3, retry_delay=2
-                )
-                response = api_call()
-
-                if isinstance(response, dict):
-                    response.pop("ResponseMetadata", None)
-                    resources = response.get(result_key.strip(), [])
-                    # Handle paginated results
-                    while "NextToken" in response:
-                        next_token = response["NextToken"]
-                        response = api_call(NextToken=next_token)
-                        response.pop("ResponseMetadata", None)
-                        resources.extend(response.get(result_key.strip(), []))
-                else:
-                    # logger.warning(f"No valid response found for {service_name} operation {operation_name}. Skipping.")
-                    continue
+                resources = paginate_or_call(client, operation_name, result_key.strip())
 
                 # Aggregate the resources
                 for resource in resources:

core/utils_egress_aws.py

@@ -5,6 +5,7 @@
 from datetime import datetime, timedelta, timezone
 from botocore.exceptions import BotoCoreError, ClientError
 
+from .utils_aws import paginate
 from .utils_egress import GIB, format_bytes, new_row
 
 logger = logging.getLogger("core.engine.egress.aws")
@@ -128,14 +129,6 @@ def fetch_latest_metric_values(
         return None
 
 
-def _paginate(client: Any, operation_name: str, result_key: str, **kwargs) -> list:
-    items = []
-    paginator = client.get_paginator(operation_name)
-    for page in paginator.paginate(**kwargs):
-        items.extend(page.get(result_key, []))
-    return items
-
-
 def _bucket_region(s3_client: Any, bucket_name: str) -> str:
     location = s3_client.get_bucket_location(Bucket=bucket_name).get(
         "LocationConstraint"
@@ -254,7 +247,7 @@ def _collect_ebs_volumes(
 ) -> list[dict[str, Any]]:
     ec2_client = session.client("ec2", region_name=region)
     rows = []
-    for volume in _paginate(ec2_client, "describe_volumes", "Volumes"):
+    for volume in paginate(ec2_client, "describe_volumes", "Volumes"):
         name = next(
             (tag["Value"] for tag in volume.get("Tags", []) if tag["Key"] == "Name"),
             volume["VolumeId"],
@@ -275,7 +268,7 @@ def _collect_ebs_snapshots(
 ) -> list[dict[str, Any]]:
     ec2_client = session.client("ec2", region_name=region)
     rows = []
-    for snapshot in _paginate(
+    for snapshot in paginate(
         ec2_client, "describe_snapshots", "Snapshots", OwnerIds=["self"]
     ):
         row = new_row(
@@ -306,7 +299,7 @@ def _collect_rds_instances(
     specs = []
     spec_rows: dict[str, tuple[dict[str, Any], int]] = {}
     for index, instance in enumerate(
-        _paginate(rds_client, "describe_db_instances", "DBInstances")
+        paginate(rds_client, "describe_db_instances", "DBInstances")
     ):
         identifier = instance["DBInstanceIdentifier"]
         row = new_row(
@@ -353,7 +346,7 @@ def _collect_dynamodb_tables(
 ) -> list[dict[str, Any]]:
     dynamodb_client = session.client("dynamodb", region_name=region)
     rows = []
-    for table_name in _paginate(dynamodb_client, "list_tables", "TableNames"):
+    for table_name in paginate(dynamodb_client, "list_tables", "TableNames"):
         try:
             table = dynamodb_client.describe_table(TableName=table_name).get(
                 "Table", {}
@@ -389,7 +382,7 @@ def _collect_backup_vaults(
 ) -> list[dict[str, Any]]:
     backup_client = session.client("backup", region_name=region)
     rows = []
-    for vault in _paginate(backup_client, "list_backup_vaults", "BackupVaultList"):
+    for vault in paginate(backup_client, "list_backup_vaults", "BackupVaultList"):
         vault_name = vault["BackupVaultName"]
         row = new_row(
             vault.get("BackupVaultArn", vault_name),