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

Author: aquamatthias

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,69 +514,63 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> 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 = []
+        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}",
                 )
-                api_instances.append(api_instance)
+                instances.append(api_instance)
                 builder.add_node(api_instance, js)
-                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
+                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
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if self.api_endpoint_configuration:

plugins/aws/fix_plugin_aws/resource/athena.py

@@ -1,6 +1,5 @@
 from typing import ClassVar, Dict, Optional, List, Any
 from typing import Type
-from concurrent.futures import wait as futures_wait
 
 from attrs import define
 
@@ -135,6 +134,8 @@ 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
@@ -147,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
@@ -163,13 +165,9 @@ 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):
-                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())]
+                fetch_workgroup(name)
         return workgroups
 
     # noinspection PyUnboundLocalVariable
@@ -255,6 +253,8 @@ 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,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
@@ -281,14 +282,10 @@ 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":
-                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())]
+                fetch_data_catalog(name)
         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)
+        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
 

plugins/aws/fix_plugin_aws/resource/cognito.py

@@ -1,7 +1,5 @@
 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
@@ -245,41 +243,32 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        pool_instances: 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:
                 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}")
-                pool_instances.append(pool_instance)
+                instances.append(pool_instance)
                 builder.add_node(pool_instance, pool)
                 builder.submit_work(service_name, add_tags, pool_instance)
-                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
+                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/dynamodb.py

@@ -1,7 +1,6 @@
 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
@@ -402,27 +401,24 @@ def called_collect_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        def add_instance(table: str) -> Optional[AwsResource]:
+        instances = []
+
+        def add_instance(table: str) -> None:
             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):
-                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())]
+                add_instance(js)
         return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
@@ -493,7 +489,9 @@ def called_collect_apis(cls) -> List[AwsApiSpec]:
 
     @classmethod
     def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        def add_instance(table: Dict[str, str]) -> Optional[AwsResource]:
+        instances = []
+
+        def add_instance(table: Dict[str, str]) -> None:
             table_description = builder.client.get(
                 service_name,
                 "describe-global-table",
@@ -502,22 +500,17 @@ def add_instance(table: Dict[str, str]) -> Optional[AwsResource]:
             )
             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:
-            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())]
+            add_instance(js)
         return instances
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None: