[aws][feat]: Update collect method implementation (#2051)

Author: 1101-1

plugins/aws/fix_plugin_aws/resource/apigateway.py

@@ -513,14 +513,16 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
+        instances: List[AwsResource] = []
         for js in json:
             if api_instance := cls.from_api(js, builder):
                 api_instance.set_arn(
                     builder=builder,
                     account="",
                     resource=f"/restapis/{api_instance.id}",
                 )
+                instances.append(api_instance)
                 builder.add_node(api_instance, js)
                 for deployment in builder.client.list(
                     service_name, "get-deployments", "items", restApiId=api_instance.id
@@ -532,6 +534,7 @@ def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) ->
                             resource=f"/restapis/{api_instance.id}/deployments/{deploy_instance.id}",
                         )
                         deploy_instance.api_link = api_instance.id
+                        instances.append(deploy_instance)
                         builder.add_node(deploy_instance, deployment)
                         builder.add_edge(api_instance, EdgeType.default, node=deploy_instance)
                         for stage in builder.client.list(
@@ -544,6 +547,7 @@ def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) ->
                             stage["syntheticId"] = f'{api_instance.id}_{stage["stageName"]}'  # create unique id
                             if stage_instance := AwsApiGatewayStage.from_api(stage, builder):
                                 stage_instance.api_link = api_instance.id
+                                instances.append(stage_instance)
                                 builder.add_node(stage_instance, stage)
                                 # reference kinds for this edge are maintained in AwsApiGatewayDeployment.reference_kinds # noqa: E501
                                 builder.add_edge(deploy_instance, EdgeType.default, node=stage_instance)
@@ -552,18 +556,21 @@ def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) ->
                 ):
                     if auth_instance := AwsApiGatewayAuthorizer.from_api(authorizer, builder):
                         auth_instance.api_link = api_instance.id
+                        instances.append(auth_instance)
                         builder.add_node(auth_instance, authorizer)
                         builder.add_edge(api_instance, EdgeType.default, node=auth_instance)
                 for resource in builder.client.list(service_name, "get-resources", "items", restApiId=api_instance.id):
                     if resource_instance := AwsApiGatewayResource.from_api(resource, builder):
                         resource_instance.api_link = api_instance.id
+                        instances.append(resource_instance)
                         if resource_instance.resource_methods:
                             for method in resource_instance.resource_methods:
                                 mapped = bend(AwsApiGatewayMethod.mapping, resource["resourceMethods"][method])
                                 if gm := parse_json(mapped, AwsApiGatewayMethod, builder):
                                     resource_instance.resource_methods[method] = gm
                         builder.add_node(resource_instance, resource)
                         builder.add_edge(api_instance, EdgeType.default, node=resource_instance)
+        return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if self.api_endpoint_configuration:

plugins/aws/fix_plugin_aws/resource/athena.py

@@ -133,7 +133,9 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
+        workgroups: List[AwsResource] = []
+
         def fetch_workgroup(name: str) -> Optional[AwsAthenaWorkGroup]:
             result = builder.client.get(
                 aws_service=service_name, action="get-work-group", result_name="WorkGroup", WorkGroup=name
@@ -146,6 +148,7 @@ def fetch_workgroup(name: str) -> Optional[AwsAthenaWorkGroup]:
                     builder=builder,
                     resource=f"workgroup/{workgroup.name}",
                 )
+                workgroups.append(workgroup)
                 builder.add_node(workgroup, result)
                 builder.submit_work(service_name, add_tags, workgroup)
                 return workgroup
@@ -165,6 +168,7 @@ def add_tags(data_catalog: AwsAthenaWorkGroup) -> None:
         for js in json:
             if (name := js.get("Name")) is not None and isinstance(name, str):
                 fetch_workgroup(name)
+        return workgroups
 
     # noinspection PyUnboundLocalVariable
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
@@ -248,7 +252,9 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
+        catalogs: List[AwsResource] = []
+
         def fetch_data_catalog(data_catalog_name: str) -> Optional[AwsAthenaDataCatalog]:
             result = builder.client.get(
                 aws_service=service_name,
@@ -260,6 +266,7 @@ def fetch_data_catalog(data_catalog_name: str) -> Optional[AwsAthenaDataCatalog]
                 return None
             if catalog := AwsAthenaDataCatalog.from_api(result, builder):
                 catalog.set_arn(builder=builder, resource=f"datacatalog/{catalog.name}")
+                catalogs.append(catalog)
                 builder.add_node(catalog, result)
                 builder.submit_work(service_name, add_tags, catalog)
                 return catalog
@@ -279,6 +286,7 @@ def add_tags(data_catalog: AwsAthenaDataCatalog) -> None:
             # we filter out the default data catalog as it is not possible to do much with it
             if (name := js.get("CatalogName")) is not None and isinstance(name, str) and name != "AwsDataCatalog":
                 fetch_data_catalog(name)
+        return catalogs
 
     def update_resource_tag(self, client: AwsClient, key: str, value: str) -> bool:
         client.call(

plugins/aws/fix_plugin_aws/resource/base.py

@@ -202,10 +202,10 @@ def collect_resources(cls: Type[AwsResource], builder: GraphBuilder) -> None:
                     expected_errors=spec.expected_errors,
                     **kwargs,
                 )
-                cls.collect(items, builder)
+                collected = cls.collect(items, builder)
                 if builder.config.collect_usage_metrics:
                     try:
-                        cls.collect_usage_metrics(builder)
+                        cls.collect_usage_metrics(builder, collected)
                     except Exception as e:
                         log.warning(
                             f"Failed to collect usage metrics for {cls.__name__} in region {builder.region.id}: {e}"
@@ -220,20 +220,26 @@ def collect_resources(cls: Type[AwsResource], builder: GraphBuilder) -> None:
                 raise
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
         # Default behavior: iterate over json snippets and for each:
         # - bend the json
         # - transform the result into a resource
         # - add the resource to the graph
+        # - return a list of resources
         # In case additional work needs to be done, override this method.
+        instances = []
         for js in json:
             if instance := cls.from_api(js, builder):
                 # post process
                 instance.post_process(builder, js)
                 builder.add_node(instance, js)
+                instances.append(instance)
+        return instances
 
     @classmethod
-    def collect_usage_metrics(cls: Type[AwsResource], builder: GraphBuilder) -> None:
+    def collect_usage_metrics(
+        cls: Type[AwsResource], builder: GraphBuilder, collected_resources: List[AwsResource]
+    ) -> None:
         # Default behavior: do nothing
         pass
 

plugins/aws/fix_plugin_aws/resource/cloudformation.py

@@ -1,6 +1,7 @@
 import time
 from datetime import datetime
 from typing import Any, ClassVar, Dict, Literal, Optional, List, Type, cast
+from concurrent.futures import Future, wait as futures_wait
 
 from attrs import define, field
 
@@ -282,10 +283,14 @@ class AwsCloudFormationStackSet(AwsResource):
     stack_set_parameters: Optional[Dict[str, Any]] = None
 
     @classmethod
-    def collect(cls, json: List[Json], builder: GraphBuilder) -> None:
-        def stack_set_instances(ss: AwsCloudFormationStackSet) -> None:
+    def collect(cls, json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
+        stack_items: List[AwsResource] = []
+
+        def stack_set_instances(ss: AwsCloudFormationStackSet) -> List[AwsCloudFormationStackInstanceSummary]:
+            instances = []
             for sij in builder.client.list(service_name, "list-stack-instances", "Summaries", StackSetName=ss.name):
                 if sii := AwsCloudFormationStackInstanceSummary.from_api(sij, builder):
+                    instances.append(sii)
                     builder.add_node(sii, sij)
                     builder.add_edge(ss, node=sii)
                     builder.graph.add_deferred_edge(
@@ -294,11 +299,18 @@ def stack_set_instances(ss: AwsCloudFormationStackSet) -> None:
                             f'is(aws_cloudformation_stack) and reported.id="{sii.stack_instance_stack_id}"'
                         ),
                     )
+            return instances
 
+        futures: List[Future[List[AwsCloudFormationStackInstanceSummary]]] = []
         for js in json:
             if stack_set := cls.from_api(js, builder):
+                stack_items.append(stack_set)
                 builder.add_node(stack_set, js)
-                builder.submit_work(service_name, stack_set_instances, stack_set)
+                future = builder.submit_work(service_name, stack_set_instances, stack_set)
+                futures.append(future)
+        futures_wait(futures)
+        stack_instances: List[AwsResource] = [result for future in futures for result in future.result()]
+        return stack_items + stack_instances
 
     def _modify_tag(self, client: AwsClient, key: str, value: Optional[str], mode: Literal["update", "delete"]) -> bool:
         tags = dict(self.tags)

plugins/aws/fix_plugin_aws/resource/cloudfront.py

@@ -23,7 +23,9 @@
 
 class CloudFrontResource:
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:  # type: ignore
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:  # type: ignore
+        instances = []
+
         def add_tags(res: AwsResource) -> None:
             tags = builder.client.get(
                 service_name, "list-tags-for-resource", "Tags", Resource=res.arn, expected_errors=["InvalidArgument"]
@@ -33,9 +35,11 @@ def add_tags(res: AwsResource) -> None:
 
         for js in json:
             if instance := cls.from_api(js, builder):
+                instances.append(instance)
                 if instance.arn:
                     builder.submit_work(service_name, add_tags, instance)
                 builder.add_node(instance, js)
+        return instances
 
     @staticmethod
     def delete_cloudfront_resource(client: AwsClient, resource: str, rid: str) -> bool:
@@ -642,11 +646,6 @@ def fetch_distribution(did: str) -> None:
                 aws_service=service_name, action="list-distributions", result_name="DistributionList.Items"
             ):
                 builder.submit_work(service_name, fetch_distribution, item["Id"])
-            if builder.config.collect_usage_metrics:
-                try:
-                    cls.collect_usage_metrics(builder)
-                except Exception as e:
-                    log.warning(f"Failed to collect usage metrics for {cls.__name__}: {e}")
         except Boto3Error as e:
             msg = f"Error while collecting {cls.__name__} in region {builder.region.name}: {e}"
             builder.core_feedback.error(msg, log)

plugins/aws/fix_plugin_aws/resource/cloudtrail.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 from typing import ClassVar, Dict, Optional, Type, List, Any
+from concurrent.futures import wait as futures_wait
 
 from attr import define, field as attrs_field, field
 
@@ -195,8 +196,8 @@ def called_collect_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
-        def collect_trail(trail_arn: str) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
+        def collect_trail(trail_arn: str) -> Optional[AwsCloudTrail]:
             if trail_raw := builder.client.get(service_name, "get-trail", "Trail", Name=trail_arn):
                 if instance := AwsCloudTrail.from_api(trail_raw, builder):
                     builder.add_node(instance, js)
@@ -206,6 +207,8 @@ def collect_trail(trail_arn: str) -> None:
                         collect_event_selectors(instance)
                     if instance.trail_has_insight_selectors:
                         collect_insight_selectors(instance)
+                    return instance
+            return None
 
         def collect_event_selectors(trail: AwsCloudTrail) -> None:
             if esj := builder.client.get(service_name, "get-event-selectors", TrailName=trail.arn):
@@ -241,17 +244,22 @@ def collect_tags(trail: AwsCloudTrail) -> None:
             ):
                 trail.tags = bend(S("TagsList", default=[]) >> ToDict(), tr)
 
+        futures = []
         for js in json:
             arn = js["TrailARN"]
             # list trails will return multi account trails in all regions
             if js["HomeRegion"] == builder.region.name and builder.account.id in arn:
                 # only collect trails in the current account and current region
-                builder.submit_work(service_name, collect_trail, arn)
+                future = builder.submit_work(service_name, collect_trail, arn)
+                futures.append(future)
             else:
                 # add a deferred edge to the trails in another account or region
                 builder.add_deferred_edge(
                     builder.region, EdgeType.default, f'is(aws_cloud_trail) and reported.arn=="{arn}"'
                 )
+        futures_wait(futures)  # only continue, when all task definitions are collected
+        instances: List[AwsResource] = [result for future in futures if (result := future.result())]
+        return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if s3 := self.trail_s3_bucket_name:

plugins/aws/fix_plugin_aws/resource/cloudwatch.py

@@ -255,16 +255,19 @@ class AwsCloudwatchAlarm(CloudwatchTaggable, AwsResource):
     cloudwatch_threshold_metric_id: Optional[str] = field(default=None)
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
         def add_tags(alarm: AwsCloudwatchAlarm) -> None:
             tags = builder.client.list(service_name, "list-tags-for-resource", "Tags", ResourceARN=alarm.arn)
             if tags:
                 alarm.tags = bend(ToDict(), tags)
 
+        instances = []
         for js in json:
             if instance := cls.from_api(js, builder):
+                instances.append(instance)
                 builder.add_node(instance, js)
                 builder.submit_work(service_name, add_tags, instance)
+        return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         super().connect_in_graph(builder, source)

plugins/aws/fix_plugin_aws/resource/cognito.py

@@ -242,7 +242,9 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
+        instances: List[AwsResource] = []
+
         def add_tags(pool: AwsCognitoUserPool) -> None:
             tags = builder.client.get(service_name, "list-tags-for-resource", "Tags", ResourceArn=pool.arn)
             if tags:
@@ -251,18 +253,22 @@ def add_tags(pool: AwsCognitoUserPool) -> None:
         for pool in json:
             if pool_instance := cls.from_api(pool, builder):
                 pool_instance.set_arn(builder=builder, resource=f"userpool/{pool_instance.id}")
+                instances.append(pool_instance)
                 builder.add_node(pool_instance, pool)
                 builder.submit_work(service_name, add_tags, pool_instance)
                 for user in builder.client.list(service_name, "list-users", "Users", UserPoolId=pool_instance.id):
                     if user_instance := AwsCognitoUser.from_api(user, builder):
                         user_instance.pool_name = pool_instance.name
                         user_instance._pool_id = pool_instance.id
+                        instances.append(user_instance)
                         builder.add_node(user_instance, user)
                         builder.add_edge(from_node=pool_instance, edge_type=EdgeType.default, node=user_instance)
                 for group in builder.client.list(service_name, "list-groups", "Groups", UserPoolId=pool_instance.id):
                     if group_instance := AwsCognitoGroup.from_api(group, builder):
+                        instances.append(group_instance)
                         builder.add_node(group_instance, group)
                         builder.add_edge(from_node=pool_instance, edge_type=EdgeType.default, node=group_instance)
+        return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if self.lambda_config:

plugins/aws/fix_plugin_aws/resource/config.py

@@ -84,9 +84,11 @@ class AwsConfigRecorder(AwsResource):
     recorder_status: Optional[AwsConfigRecorderStatus] = field(default=None, metadata=dict(ignore_history=True))
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
         # get all statuses
         statuses: Dict[str, AwsConfigRecorderStatus] = {}
+
+        instances: List[AwsResource] = []
         for r in builder.client.list(
             service_name, "describe-configuration-recorder-status", "ConfigurationRecordersStatus"
         ):
@@ -95,10 +97,12 @@ def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) ->
 
         for js in json:
             if instance := AwsConfigRecorder.from_api(js, builder):
+                instances.append(instance)
                 if status := statuses.get(instance.id):
                     instance.recorder_status = status
                     instance.mtime = status.last_status_change_time
                 builder.add_node(instance, js)
+        return instances
 
     def delete_resource(self, client: AwsClient, graph: Graph) -> bool:
         client.call(service_name, "delete-configuration-recorder", self.name)