[feat][aws] Add more resources to done in parallel (#2066)

Co-authored-by: Matthias Veit <matthias_veit@yahoo.de>

Author: 1101-1

plugins/aws/fix_plugin_aws/resource/apigateway.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from typing import ClassVar, Dict, Optional, List, Type, Union, Any
-
+from concurrent.futures import wait as futures_wait
 from attrs import define, field
 from fix_plugin_aws.aws_client import AwsClient
 
@@ -514,63 +514,69 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        instances: List[AwsResource] = []
+        api_instances: List[AwsResource] = []
+
+        def add_instance(api_instance: AwsResource) -> List[AwsResource]:
+            instances: List[AwsResource] = []
+            for deployment in builder.client.list(service_name, "get-deployments", "items", restApiId=api_instance.id):
+                if deploy_instance := AwsApiGatewayDeployment.from_api(deployment, builder):
+                    deploy_instance.set_arn(
+                        builder=builder,
+                        account="",
+                        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(
+                        service_name,
+                        "get-stages",
+                        "item",
+                        restApiId=api_instance.id,
+                        deploymentId=deploy_instance.id,
+                    ):
+                        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)
+            for authorizer in builder.client.list(service_name, "get-authorizers", "items", restApiId=api_instance.id):
+                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
+
+        futures = []
         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)
+                api_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
-                ):
-                    if deploy_instance := AwsApiGatewayDeployment.from_api(deployment, builder):
-                        deploy_instance.set_arn(
-                            builder=builder,
-                            account="",
-                            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(
-                            service_name,
-                            "get-stages",
-                            "item",
-                            restApiId=api_instance.id,
-                            deploymentId=deploy_instance.id,
-                        ):
-                            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)
-                for authorizer in builder.client.list(
-                    service_name, "get-authorizers", "items", restApiId=api_instance.id
-                ):
-                    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
+                future = builder.submit_work(service_name, add_instance, api_instance)
+                futures.append(future)
+        futures_wait(futures)
+        instances: List[AwsResource] = [result for future in futures for result in future.result()]
+        return api_instances + instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if self.api_endpoint_configuration:

plugins/aws/fix_plugin_aws/resource/athena.py

@@ -1,5 +1,6 @@
 from typing import ClassVar, Dict, Optional, List, Any
 from typing import Type
+from concurrent.futures import wait as futures_wait
 
 from attrs import define
 
@@ -134,8 +135,6 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     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
@@ -148,7 +147,6 @@ 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,9 +163,13 @@ def add_tags(data_catalog: AwsAthenaWorkGroup) -> None:
             if tags:
                 data_catalog.tags = bend(ToDict(), tags)
 
+        futures = []
         for js in json:
             if (name := js.get("Name")) is not None and isinstance(name, str):
-                fetch_workgroup(name)
+                future = builder.submit_work(service_name, fetch_workgroup, name)
+                futures.append(future)
+        futures_wait(futures)
+        workgroups: List[AwsResource] = [result for future in futures if (result := future.result())]
         return workgroups
 
     # noinspection PyUnboundLocalVariable
@@ -253,8 +255,6 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     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,
@@ -266,7 +266,6 @@ 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
@@ -282,10 +281,14 @@ def add_tags(data_catalog: AwsAthenaDataCatalog) -> None:
             if tags:
                 data_catalog.tags = bend(S("Tags", default=[]) >> ToDict(), tags[0])
 
+        futures = []
         for js in json:
             # 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)
+                future = builder.submit_work(service_name, fetch_data_catalog, name)
+                futures.append(future)
+        futures_wait(futures)
+        catalogs: List[AwsResource] = [result for future in futures if (result := future.result())]
         return catalogs
 
     def update_resource_tag(self, client: AwsClient, key: str, value: str) -> bool:

plugins/aws/fix_plugin_aws/resource/cloudtrail.py

@@ -257,7 +257,7 @@ def collect_tags(trail: AwsCloudTrail) -> None:
                 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
+        futures_wait(futures)
         instances: List[AwsResource] = [result for future in futures if (result := future.result())]
         return instances
 

plugins/aws/fix_plugin_aws/resource/cognito.py

@@ -1,5 +1,7 @@
 from attrs import define, field
 from typing import ClassVar, Dict, List, Optional, Type, Tuple, Any
+from concurrent.futures import wait as futures_wait
+
 from fix_plugin_aws.aws_client import AwsClient
 from fix_plugin_aws.resource.base import AwsApiSpec, AwsResource, GraphBuilder
 from fix_plugin_aws.resource.iam import AwsIamRole
@@ -243,32 +245,41 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        instances: List[AwsResource] = []
+        pool_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:
                 pool.tags = tags
 
+        def fetch_additional_instances(pool_instance: AwsCognitoUser) -> List[AwsResource]:
+            instances: List[AwsResource] = []
+            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
+
+        futures = []
         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)
+                pool_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
+                future = builder.submit_work(service_name, fetch_additional_instances, pool_instance)
+                futures.append(future)
+        futures_wait(futures)
+        user_and_groups: List[AwsResource] = [result for future in futures for result in future.result()]
+        return pool_instances + user_and_groups
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if self.lambda_config:

plugins/aws/fix_plugin_aws/resource/dynamodb.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import ClassVar, Dict, List, Optional, Type, Any
 from attrs import define, field
+from concurrent.futures import Future, wait as futures_wait
 from fix_plugin_aws.aws_client import AwsClient
 from fix_plugin_aws.resource.base import AwsApiSpec, AwsResource, GraphBuilder
 from fix_plugin_aws.resource.kinesis import AwsKinesisStream
@@ -401,24 +402,27 @@ def called_collect_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        instances = []
-
-        def add_instance(table: str) -> None:
+        def add_instance(table: str) -> Optional[AwsResource]:
             table_description = builder.client.get(service_name, "describe-table", "Table", TableName=table)
             if table_description is not None:
                 if instance := cls.from_api(table_description, builder):
-                    instances.append(instance)
                     builder.add_node(instance, table_description)
                     builder.submit_work(service_name, add_tags, instance)
+                    return instance
+            return None
 
         def add_tags(table: AwsDynamoDbTable) -> None:
             tags = builder.client.list(service_name, "list-tags-of-resource", "Tags", ResourceArn=table.arn)
             if tags:
                 table.tags = bend(ToDict(), tags)
 
+        futures: List[Future[Optional[AwsResource]]] = []
         for js in json:
             if isinstance(js, str):
-                add_instance(js)
+                future = builder.submit_work(service_name, add_instance, js)
+                futures.append(future)
+        futures_wait(futures)
+        instances: List[AwsResource] = [result for future in futures if (result := future.result())]
         return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
@@ -489,9 +493,7 @@ def called_collect_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        instances = []
-
-        def add_instance(table: Dict[str, str]) -> None:
+        def add_instance(table: Dict[str, str]) -> Optional[AwsResource]:
             table_description = builder.client.get(
                 service_name,
                 "describe-global-table",
@@ -500,17 +502,22 @@ def add_instance(table: Dict[str, str]) -> None:
             )
             if table_description:
                 if instance := cls.from_api(table_description, builder):
-                    instances.append(instance)
                     builder.add_node(instance, table_description)
                     builder.submit_work(service_name, add_tags, instance)
+                    return instance
+            return None
 
         def add_tags(table: AwsDynamoDbGlobalTable) -> None:
             tags = builder.client.list(service_name, "list-tags-of-resource", "Tags", ResourceArn=table.arn)
             if tags:
                 table.tags = bend(ToDict(), tags)
 
+        futures = []
         for js in json:
-            add_instance(js)
+            future = builder.submit_work(service_name, add_instance, js)
+            futures.append(future)
+        futures_wait(futures)
+        instances: List[AwsResource] = [result for future in futures if (result := future.result())]
         return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None: