[aws][feat] Bedrock resources collection (#2190)

Co-authored-by: Matthias Veit <aquamatthias@users.noreply.github.com>
Co-authored-by: Matthias Veit <matthias_veit@yahoo.de>

Author: 1101-1

fixlib/fixlib/baseresources.py

@@ -1513,6 +1513,31 @@ class BaseManagedKubernetesClusterProvider(BaseResource):
     endpoint: Optional[str] = field(default=None, metadata={"description": "The kubernetes API endpoint"})
 
 
+@define(eq=False, slots=False)
+class BaseAIResource(BaseResource):
+    kind: ClassVar[str] = "ai_resource"
+    kind_display: ClassVar[str] = "AI resource"
+    kind_description: ClassVar[str] = "An AI Resource."
+    metadata: ClassVar[Dict[str, Any]] = {"icon": "resource", "group": "ai"}
+    _categories: ClassVar[List[Category]] = [Category.ai, Category.compute]
+
+
+@define(eq=False, slots=False)
+class BaseAIJob(BaseAIResource):
+    kind: ClassVar[str] = "ai_job"
+    kind_display: ClassVar[str] = "AI Job"
+    kind_description: ClassVar[str] = "An AI Job resource."
+    metadata: ClassVar[Dict[str, Any]] = {"icon": "job", "group": "ai"}
+
+
+@define(eq=False, slots=False)
+class BaseAIModel(BaseAIResource):
+    kind: ClassVar[str] = "ai_model"
+    kind_display: ClassVar[str] = "AI Model"
+    kind_description: ClassVar[str] = "An AI Model resource."
+    metadata: ClassVar[Dict[str, Any]] = {"icon": "resource", "group": "ai"}
+
+
 @define(eq=False, slots=False)
 class UnknownCloud(BaseCloud):
     kind: ClassVar[str] = "unknown_cloud"

plugins/aws/fix_plugin_aws/collector.py

@@ -48,6 +48,7 @@
     acm,
     waf,
     backup,
+    bedrock,
 )
 from fix_plugin_aws.resource.base import AwsAccount, AwsApiSpec, AwsRegion, AwsResource, GraphBuilder
 from fixlib.baseresources import Cloud, EdgeType, BaseOrganizationalRoot, BaseOrganizationalUnit
@@ -107,6 +108,7 @@
     + redshift.resources
     + backup.resources
     + amazonq.resources
+    + bedrock.resources
 )
 all_resources: List[Type[AwsResource]] = global_resources + regional_resources
 
@@ -259,6 +261,8 @@ def get_last_run() -> Optional[datetime]:
 
             # wait for all futures to finish
             shared_queue.wait_for_submitted_work()
+            # remove unused nodes
+            self.remove_unused()
             self.core_feedback.progress_done(self.account.dname, 1, 1, context=[self.cloud.id])
             self.error_accumulator.report_all(global_builder.core_feedback)
 
@@ -330,6 +334,32 @@ def collect_and_set_metrics(
 
             builder.submit_work("cloudwatch", collect_and_set_metrics, start, region, queries)
 
+    def remove_unused(self) -> None:
+        remove_nodes = []
+
+        def rm_nodes(cls, ignore_kinds: Optional[Type[Any]] = None, check_pred: bool = True) -> None:  # type: ignore
+            for node in self.graph.nodes:
+                if not isinstance(node, cls):
+                    continue
+                if check_pred:
+                    nodes = list(self.graph.predecessors(node))
+                else:
+                    nodes = list(self.graph.successors(node))
+                if ignore_kinds is not None:
+                    nodes = [n for n in nodes if not isinstance(n, ignore_kinds)]
+                if not nodes:
+                    remove_nodes.append(node)
+            log.debug(f"Removing {len(remove_nodes)} unreferenced nodes of type {cls}")
+            removed = set()
+            for node in remove_nodes:
+                if node in removed:
+                    continue
+                removed.add(node)
+                self.graph.remove_node(node)
+            remove_nodes.clear()
+
+        rm_nodes(bedrock.AwsBedrockFoundationModel, check_pred=False)
+
     # TODO: move into separate AwsAccountSettings
     def update_account(self) -> None:
         log.info(f"Collecting AWS IAM Account Summary in account {self.account.dname}")

plugins/aws/fix_plugin_aws/resource/amazonq.py

@@ -143,7 +143,7 @@ def add_tags(tag_resource: AwsResource) -> None:
                     resourceARN=tag_resource.arn,
                 )
                 if tags:
-                    tag_resource.tags.update(tags)
+                    tag_resource.tags.update(tags[0])
             else:
                 tags = builder.client.list(
                     service_name,
@@ -356,9 +356,9 @@ class AwsQBusinessDataSource(AmazonQTaggable, AwsResource):
     metadata: ClassVar[Dict[str, Any]] = {"icon": "bucket", "group": "ai"}
     # Collected via AwsQBusinessApplication()
     aws_metadata: ClassVar[Dict[str, Any]] = {
-        "provider_link_tpl": "https://{region_id}.console.aws.amazon.com/amazonq/business/applications/{app_id}/indices/{indice_id}/datasources/{id}/details?region={region}",  # fmt: skip
+        "provider_link_tpl": "https://{region_id}.console.aws.amazon.com/amazonq/business/applications/{application_id}/indices/{indice_id}/datasources/{id}/details?region={region}",  # fmt: skip
         "arn_tpl": "arn:{partition}:qbusiness:{region}:{account}:application/{application_id}/index/{indice_id}/data-source/{id}",
-        "extra_args": ["application_id", "indice_id"],
+        "extra_args_for_arn": ["application_id", "indice_id"],
     }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("dataSourceId"),
@@ -528,7 +528,7 @@ class AwsQBusinessIndice(AmazonQTaggable, AwsResource):
     metadata: ClassVar[Dict[str, Any]] = {"icon": "config", "group": "ai"}
     aws_metadata: ClassVar[Dict[str, Any]] = {
         "arn_tpl": "arn:{partition}:qbusiness:{region}:{account}:application/{application_id}/index/{id}",
-        "extra_args": ["application_id"],
+        "extra_args_for_arn": ["application_id"],
     }
     # Collected via AwsQBusinessApplication()
     mapping: ClassVar[Dict[str, Bender]] = {
@@ -753,7 +753,7 @@ class AwsQBusinessPlugin(AmazonQTaggable, AwsResource):
     metadata: ClassVar[Dict[str, Any]] = {"icon": "resource", "group": "ai"}
     aws_metadata: ClassVar[Dict[str, Any]] = {
         "arn_tpl": "arn:{partition}:qbusiness:{region}:{account}:application/{application_id}/plugin/{id}",
-        "extra_args": ["application_id"],
+        "extra_args_for_arn": ["application_id"],
     }
     # Collected via AwsQBusinessApplication()
     mapping: ClassVar[Dict[str, Bender]] = {
@@ -819,7 +819,7 @@ class AwsQBusinessRetriever(AmazonQTaggable, AwsResource):
     metadata: ClassVar[Dict[str, Any]] = {"icon": "application", "group": "ai"}
     aws_metadata: ClassVar[Dict[str, Any]] = {
         "arn_tpl": "arn:{partition}:qbusiness:{region}:{account}:application/{application_id}/retriever/{id}",
-        "extra_args": ["application_id"],
+        "extra_args_for_arn": ["application_id"],
     }
     # Collected via AwsQBusinessApplication()
     mapping: ClassVar[Dict[str, Bender]] = {
@@ -876,7 +876,7 @@ class AwsQBusinessWebExperience(AmazonQTaggable, AwsResource):
     metadata: ClassVar[Dict[str, Any]] = {"icon": "application", "group": "ai"}
     aws_metadata: ClassVar[Dict[str, Any]] = {
         "arn_tpl": "arn:{partition}:qbusiness:{region}:{account}:application/{application_id}/web-experience/{id}",
-        "extra_args": ["application_id"],
+        "extra_args_for_arn": ["application_id"],
     }
     # Collected via AwsQBusinessApplication()
     mapping: ClassVar[Dict[str, Bender]] = {

plugins/aws/fix_plugin_aws/resource/base.py

@@ -520,7 +520,7 @@ def add_node(
                 }
 
                 # Add any additional dynamic arguments from the metadata (if they exist)
-                if extra_args := meta.get("extra_args"):
+                if extra_args := meta.get("extra_args_for_arn"):
                     for extra_arg in extra_args:
                         args[extra_arg] = getattr(node, extra_arg)