[azure] [feat] Add new edges between resources (#1954)

Author: 1101-1

plugins/azure/fix_plugin_azure/resource/base.py

@@ -48,16 +48,12 @@ class AzureResource(BaseResource):
     def resource_subscription_id(self) -> Optional[str]:
         return self.extract_part("subscriptionId")
 
-    def resource_group_name(self) -> Optional[str]:
-        return self.extract_part("resourceGroupName")
-
     def extract_part(self, part: str) -> Optional[str]:
         """
         Extracts a specific part from a resource ID.
 
-        The function takes a resource ID and a specified part to extract, such as 'subscriptionId'
-        or 'resourceGroupName'. The resource ID is expected to follow the Azure Resource Manager
-        path format.
+        The function takes a resource ID and a specified part to extract, such as 'subscriptionId'.
+        The resource ID is expected to follow the Azure Resource Manager path format.
 
         Example:
         For the resource ID "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/...",
@@ -78,14 +74,6 @@ def extract_part(self, part: str) -> Optional[str]:
             if index := id_parts.index("subscriptions"):
                 return id_parts[index + 1]
             return None
-
-        elif part == "resourceGroupName":
-            if "resourceGroups" not in id_parts:
-                return None
-            if index := id_parts.index("resourceGroups"):
-                return id_parts[index + 1]
-            return None
-
         else:
             return None
 

plugins/azure/fix_plugin_azure/resource/compute.py

@@ -68,7 +68,7 @@ class AzureAvailabilitySet(AzureResource):
         expect_array=True,
     )
     reference_kinds: ClassVar[ModelReference] = {
-        "successors": {"default": ["azure_proximity_placement_group", "azure_virtual_machine"]},
+        "successors": {"default": ["azure_proximity_placement_group", "azure_virtual_machine_base"]},
     }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
@@ -95,7 +95,7 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
             )
         if virtual_machines := self.virtual_machines_availability:
             for vm_id in virtual_machines:
-                builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachine, id=vm_id)
+                builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachineBase, id=vm_id)
 
 
 @define(eq=False, slots=False)
@@ -120,7 +120,7 @@ class AzureCapacityReservationGroup(AzureResource):
         expect_array=True,
     )
     reference_kinds: ClassVar[ModelReference] = {
-        "successors": {"default": ["azure_virtual_machine"]},
+        "successors": {"default": ["azure_virtual_machine_base"]},
     }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
@@ -140,7 +140,7 @@ class AzureCapacityReservationGroup(AzureResource):
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if virtual_machines := self.virtual_machines_associated:
             for vm_id in virtual_machines:
-                builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachine, id=vm_id)
+                builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachineBase, id=vm_id)
 
 
 @define(eq=False, slots=False)
@@ -1755,7 +1755,7 @@ class AzureRestorePointCollection(AzureResource):
         expect_array=True,
     )
     reference_kinds: ClassVar[ModelReference] = {
-        "successors": {"default": ["azure_virtual_machine"]},
+        "successors": {"default": ["azure_virtual_machine_base"]},
     }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
@@ -1773,7 +1773,7 @@ class AzureRestorePointCollection(AzureResource):
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if (source_id := self.source) and (vm_id := source_id.id):
-            builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachine, id=vm_id)
+            builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureVirtualMachineBase, id=vm_id)
 
 
 @define(eq=False, slots=False)
@@ -2515,17 +2515,8 @@ class AzureVirtualMachineIdentity:
 
 
 @define(eq=False, slots=False)
-class AzureVirtualMachine(AzureResource, BaseInstance):
-    kind: ClassVar[str] = "azure_virtual_machine"
-    api_spec: ClassVar[AzureApiSpec] = AzureApiSpec(
-        service="compute",
-        version="2023-03-01",
-        path="/subscriptions/{subscriptionId}/providers/Microsoft.Compute/virtualMachines",
-        path_parameters=["subscriptionId"],
-        query_parameters=["api-version"],
-        access_path="value",
-        expect_array=True,
-    )
+class AzureVirtualMachineBase(AzureResource, BaseInstance):
+    kind: ClassVar[str] = "azure_virtual_machine_base"
     reference_kinds: ClassVar[ModelReference] = {
         "predecessors": {
             "default": [
@@ -2780,6 +2771,20 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
             )
 
 
+@define(eq=False, slots=False)
+class AzureVirtualMachine(AzureVirtualMachineBase):
+    kind: ClassVar[str] = "azure_virtual_machine"
+    api_spec: ClassVar[AzureApiSpec] = AzureApiSpec(
+        service="compute",
+        version="2023-03-01",
+        path="/subscriptions/{subscriptionId}/providers/Microsoft.Compute/virtualMachines",
+        path_parameters=["subscriptionId"],
+        query_parameters=["api-version"],
+        access_path="value",
+        expect_array=True,
+    )
+
+
 @define(eq=False, slots=False)
 class AzureRollingUpgradePolicy:
     kind: ClassVar[str] = "azure_rolling_upgrade_policy"
@@ -3229,7 +3234,8 @@ class AzureVirtualMachineScaleSet(AzureResource, BaseAutoScalingGroup):
         expect_array=True,
     )
     reference_kinds: ClassVar[ModelReference] = {
-        "predecessors": {"default": ["azure_load_balancer"]},
+        "predecessors": {"default": ["azure_load_balancer", "azure_subnet"]},
+        "successors": {"default": ["azure_virtual_machine_scale_set_instance"]},
     }
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("id"),
@@ -3287,6 +3293,30 @@ class AzureVirtualMachineScaleSet(AzureResource, BaseAutoScalingGroup):
     virtual_machine_profile: Optional[AzureVirtualMachineScaleSetVMProfile] = field(default=None, metadata={'description': 'Describes a virtual machine scale set virtual machine profile.'})  # fmt: skip
     zone_balance: Optional[bool] = field(default=None, metadata={'description': 'Whether to force strictly even virtual machine distribution cross x-zones in case there is zone outage. Zonebalance property can only be set if the zones property of the scale set contains more than one zone. If there are no zones or only one zone specified, then zonebalance property should not be set.'})  # fmt: skip
 
+    def post_process(self, graph_builder: GraphBuilder, source: Json) -> None:
+        def collect_vmss_instances() -> None:
+            api_spec = AzureApiSpec(
+                service="compute",
+                version="2023-09-01",
+                path=f"{self.id}/virtualMachines",
+                path_parameters=[],
+                query_parameters=["api-version"],
+                access_path="value",
+                expect_array=True,
+            )
+
+            items = graph_builder.client.list(api_spec)
+            vmss_instance_ids = [str(item.get("id")) for item in items if item.get("id") is not None]
+
+            AzureVirtualMachineScaleSetInstance.collect(items, graph_builder)
+
+            for vmss_instance_id in vmss_instance_ids:
+                graph_builder.add_edge(
+                    self, edge_type=EdgeType.default, clazz=AzureVirtualMachineScaleSetInstance, id=vmss_instance_id
+                )
+
+        graph_builder.submit_work("azure_all", collect_vmss_instances)
+
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if (
             (vm_profile := self.virtual_machine_profile)
@@ -3298,14 +3328,22 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
                     for ip_config in ip_configs:
                         if baps := ip_config.load_balancer_backend_address_pools:
                             for bap in baps:
-                                if bap_id := bap:
-                                    builder.add_edge(
-                                        self,
-                                        edge_type=EdgeType.default,
-                                        reverse=True,
-                                        clazz=AzureLoadBalancer,
-                                        id=bap_id,
-                                    )
+                                bap_id = "/".join(bap.split("/")[:-2])
+                                builder.add_edge(
+                                    self,
+                                    edge_type=EdgeType.default,
+                                    reverse=True,
+                                    clazz=AzureLoadBalancer,
+                                    id=bap_id,
+                                )
+                        if subnet_id := ip_config.subnet:
+                            builder.add_edge(
+                                self,
+                                edge_type=EdgeType.default,
+                                reverse=True,
+                                clazz=AzureSubnet,
+                                id=subnet_id,
+                            )
 
 
 @define(eq=False, slots=False)
@@ -3344,6 +3382,13 @@ def pre_process(self, graph_builder: GraphBuilder, source: Json) -> None:
         self.location = graph_builder.location.name if graph_builder.location else ""
 
 
+@define(eq=False, slots=False)
+class AzureVirtualMachineScaleSetInstance(AzureVirtualMachineBase):
+    # note: instances are collected as part of collecting AzureVirtualMachineScaleSets
+
+    kind: ClassVar[str] = "azure_virtual_machine_scale_set_instance"
+
+
 resources: List[Type[AzureResource]] = [
     AzureAvailabilitySet,
     AzureCapacityReservationGroup,

plugins/azure/fix_plugin_azure/resource/containerservice.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import ClassVar, Dict, Optional, List, Type
+from typing import ClassVar, Dict, Optional, List, Tuple, Type
 
 from attr import define, field
 
@@ -814,7 +814,7 @@ class AzureManagedCluster(AzureResource):
         expect_array=True,
     )
     reference_kinds: ClassVar[ModelReference] = {
-        "successors": {"default": ["azure_disk_encryption_set"]},
+        "successors": {"default": ["azure_disk_encryption_set", "azure_virtual_machine_scale_set"]},
     }
     mapping: ClassVar[Dict[str, Bender]] = AzureTrackedResource.mapping | {
         "id": S("id"),
@@ -875,9 +875,7 @@ class AzureManagedCluster(AzureResource):
     }
     aad_profile: Optional[AzureManagedClusterAADProfile] = field(default=None, metadata={'description': 'For more details see [managed AAD on AKS](https://docs.microsoft.com/azure/aks/managed-aad).'})  # fmt: skip
     addon_profiles: Optional[Dict[str, AzureManagedClusterAddonProfile]] = field(default=None, metadata={'description': 'The profile of managed cluster add-on.'})  # fmt: skip
-    agent_pool_profiles: Optional[List[str]] = field(
-        default=None, metadata={"description": "The agent pool properties."}
-    )
+    agent_pool_profiles: Optional[List[str]] = field(default=None, metadata={"description": "The agent pool properties."})  # fmt: skip
     api_server_access_profile: Optional[AzureManagedClusterAPIServerAccessProfile] = field(default=None, metadata={'description': 'Access profile for managed cluster API server.'})  # fmt: skip
     auto_scaler_profile: Optional[AutoScalerProfile] = field(default=None, metadata={'description': 'Parameters to be applied to the cluster-autoscaler when enabled'})  # fmt: skip
     auto_upgrade_profile: Optional[AzureManagedClusterAutoUpgradeProfile] = field(default=None, metadata={'description': 'Auto upgrade profile for a managed cluster.'})  # fmt: skip
@@ -919,11 +917,30 @@ class AzureManagedCluster(AzureResource):
     workload_auto_scaler_profile: Optional[AzureManagedClusterWorkloadAutoScalerProfile] = field(default=None, metadata={'description': 'Workload Auto-scaler profile for the managed cluster.'})  # fmt: skip
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
-        from fix_plugin_azure.resource.compute import AzureDiskEncryptionSet
+        from fix_plugin_azure.resource.compute import AzureDiskEncryptionSet, AzureVirtualMachineScaleSet
 
         if disk_id := self.disk_encryption_set_id:
             builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureDiskEncryptionSet, id=disk_id)
 
+        if agent_pool_profiles := self.agent_pool_profiles:
+            vmss_agent_pool_names_and_ids = self._get_poolnames_and_vmss_ids(builder)
+            for agent_pool_profile_name in agent_pool_profiles:
+                for info in vmss_agent_pool_names_and_ids:
+                    pool_name, vmss_id = info
+                    if agent_pool_profile_name == pool_name:
+                        builder.add_edge(
+                            self, edge_type=EdgeType.default, clazz=AzureVirtualMachineScaleSet, id=vmss_id
+                        )
+
+    def _get_poolnames_and_vmss_ids(self, builder: GraphBuilder) -> List[Tuple[str, str]]:
+        from fix_plugin_azure.resource.compute import AzureVirtualMachineScaleSet
+
+        return [
+            (poolname, vmss_id)
+            for vmss in builder.nodes(clazz=AzureVirtualMachineScaleSet)
+            if (poolname := vmss.tags.get("aks-managed-poolName")) and (vmss_id := vmss.id)
+        ]
+
 
 @define(eq=False, slots=False)
 class AzureProxyResource:

plugins/azure/fix_plugin_azure/resource/metrics.py

@@ -219,6 +219,7 @@ def query_for(
                 "aggregation",
                 "interval",
                 "AutoAdjustTimegrain",
+                "ValidateDimensions",
             ],
             access_path=None,
             expect_array=False,
@@ -272,6 +273,7 @@ def _query_for_single(
             aggregation=aggregation,
             interval=interval,
             AutoAdjustTimegrain=True,
+            ValidateDimensions=False,
         )
         # Iterate over the retrieved data and map it to AzureMetricData objects
         for single in part: