From 029313f3f96535dad846fa25415df22290cf4dc5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bence=20H=C3=A9zs=C5=91?= Date: Wed, 22 Jul 2026 15:21:39 +0200 Subject: [PATCH 1/2] fix aws resource inventory pagination --- core/utils_aws.py | 111 +++++++++++++++++++++++++-------------- core/utils_egress_aws.py | 19 +++---- 2 files changed, 79 insertions(+), 51 deletions(-) diff --git a/core/utils_aws.py b/core/utils_aws.py index 759a9a9..8044bae 100644 --- a/core/utils_aws.py +++ b/core/utils_aws.py @@ -18,6 +18,29 @@ 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, @@ -25,29 +48,56 @@ def aws_api_call_with_retry( 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 api_call # Return the callable function +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 +165,9 @@ 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 + resources = paginate_or_call( + client, operation_name, result_key.strip() ) - 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 # Aggregate the resources for resource in resources: diff --git a/core/utils_egress_aws.py b/core/utils_egress_aws.py index b7d16fb..dc9010e 100644 --- a/core/utils_egress_aws.py +++ b/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), From 21b76631c6a4fdac27e8bcd8a082c540d596a19b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bence=20H=C3=A9zs=C5=91?= Date: Wed, 22 Jul 2026 15:37:49 +0200 Subject: [PATCH 2/2] fix aws resource inventory pagination - fix formatting --- core/utils_aws.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/core/utils_aws.py b/core/utils_aws.py index 8044bae..9681837 100644 --- a/core/utils_aws.py +++ b/core/utils_aws.py @@ -38,7 +38,9 @@ def _retry_call( 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") + raise Exception( + f"Failed to call {getattr(func, '__name__', repr(func))} after {max_retries} attempts" + ) def aws_api_call_with_retry( @@ -165,9 +167,7 @@ def build_aws_resource_inventory( # logger.error(f"Operation {operation_name} does not exist for service {service_name}") continue - resources = paginate_or_call( - client, operation_name, result_key.strip() - ) + resources = paginate_or_call(client, operation_name, result_key.strip()) # Aggregate the resources for resource in resources: