From d41709ea442575c338761c5b2b3d5171ae64976d Mon Sep 17 00:00:00 2001
From: 1101-1 <70093559+1101-1@users.noreply.github.com>
Date: Tue, 18 Jun 2024 15:38:24 +0500
Subject: [PATCH] [aws][fix]: Fix non acyclic graph connection (#2118)

---
 plugins/aws/fix_plugin_aws/resource/backup.py | 25 +++++++++++++++++--
 plugins/aws/test/collector_test.py            |  2 +-
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/plugins/aws/fix_plugin_aws/resource/backup.py b/plugins/aws/fix_plugin_aws/resource/backup.py
index ebe6db5837..dd4d03103e 100644
--- a/plugins/aws/fix_plugin_aws/resource/backup.py
+++ b/plugins/aws/fix_plugin_aws/resource/backup.py
@@ -5,6 +5,13 @@
 from attrs import define, field
 
 from fix_plugin_aws.resource.base import AwsResource, AwsApiSpec, GraphBuilder
+from fix_plugin_aws.resource.cloudformation import AwsCloudFormationStack
+from fix_plugin_aws.resource.dynamodb import AwsDynamoDbTable, AwsDynamoDbGlobalTable
+from fix_plugin_aws.resource.ec2 import AwsEc2Instance, AwsEc2Volume
+from fix_plugin_aws.resource.efs import AwsEfsFileSystem
+from fix_plugin_aws.resource.rds import AwsRdsCluster, AwsRdsInstance
+from fix_plugin_aws.resource.redshift import AwsRedshiftCluster
+from fix_plugin_aws.resource.s3 import AwsS3Bucket
 from fix_plugin_aws.utils import TagsValue
 from fixlib.baseresources import ModelReference
 from fixlib.json_bender import Bender, S, ForallBend, Bend
@@ -172,7 +179,6 @@ class AwsBackupProtectedResource(AwsResource):
         "resource_name": S("ResourceName"),
         "last_backup_vault_arn": S("LastBackupVaultArn"),
         "last_recovery_point_arn": S("LastRecoveryPointArn"),
-        "arn": S("ResourceArn"),
     }
     resource_arn: Optional[str] = field(default=None, metadata={"description": "An Amazon Resource Name (ARN) that uniquely identifies a resource. The format of the ARN depends on the resource type."})  # fmt: skip
     resource_type: Optional[str] = field(default=None, metadata={"description": "The type of Amazon Web Services resource; for example, an Amazon Elastic Block Store (Amazon EBS) volume or an Amazon Relational Database Service (Amazon RDS) database. For Windows Volume Shadow Copy Service (VSS) backups, the only supported resource type is Amazon EC2."})  # fmt: skip
@@ -183,7 +189,22 @@ class AwsBackupProtectedResource(AwsResource):
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if resource_arn := self.resource_arn:
-            builder.add_edge(self, clazz=AwsResource, arn=resource_arn)
+            builder.add_edge(
+                self,
+                clazz=(
+                    AwsS3Bucket,
+                    AwsEc2Instance,
+                    AwsEc2Volume,
+                    AwsRdsCluster,
+                    AwsRdsInstance,
+                    AwsDynamoDbTable,
+                    AwsDynamoDbGlobalTable,
+                    AwsEfsFileSystem,
+                    AwsRedshiftCluster,
+                    AwsCloudFormationStack,
+                ),
+                arn=resource_arn,
+            )
         if vault_arn := self.last_backup_vault_arn:
             builder.add_edge(self, reverse=True, clazz=AwsBackupVault, id=vault_arn)
         if recovery_point_arn := self.last_recovery_point_arn:
diff --git a/plugins/aws/test/collector_test.py b/plugins/aws/test/collector_test.py
index 58f0f7f318..4026b033f1 100644
--- a/plugins/aws/test/collector_test.py
+++ b/plugins/aws/test/collector_test.py
@@ -34,7 +34,7 @@ def count_kind(clazz: Type[AwsResource]) -> int:
     assert len(threading.enumerate()) == 1
     # ensure the correct number of nodes and edges
     assert count_kind(AwsResource) == 236
-    assert len(account_collector.graph.edges) == 540
+    assert len(account_collector.graph.edges) == 539
     assert len(account_collector.graph.deferred_edges) == 2
     for node in account_collector.graph.nodes:
         if isinstance(node, AwsRegion):