[azure] [feat] Filter virtual machine sizes (#1852)

plugins/azure/resoto_plugin_azure/azure_client.py

@@ -153,7 +153,7 @@ def _call(self, spec: AzureApiSpec, **kwargs: Any) -> List[Json]:
             if lookup_map.get(param, None) is not None:
                 path_map[param] = lookup_map[param]
             else:
-                raise ValueError(f"Param {param} in lookup_map does not found")
+                raise KeyError(f"Path parameter {param} was not provided as argument")
 
         # Construct parameters
         params = case_insensitive_dict()
@@ -163,7 +163,7 @@ def _call(self, spec: AzureApiSpec, **kwargs: Any) -> List[Json]:
                 if lookup_map.get(param, None) is not None:
                     params[param] = ser.query(param, lookup_map[param], "str")  # type: ignore # noqa: E501
                 else:
-                    raise ValueError(f"Param {param} in lookup_map does not found")
+                    raise KeyError(f"Query parameter {param} was not provided as argument")
 
         # Construct url
         path = spec.path.format_map(path_map)

plugins/azure/resoto_plugin_azure/collector.py

@@ -1,12 +1,15 @@
 import logging
 from concurrent.futures import ThreadPoolExecutor, Future
-from typing import Type, Set, List
-
+from typing import Any, Optional, Type, Set, List
 
 from resoto_plugin_azure.config import AzureConfig, AzureCredentials
 from resoto_plugin_azure.azure_client import AzureClient
-from resoto_plugin_azure.resource.compute import resources as compute_resources
+from resoto_plugin_azure.resource.compute import (
+    AzureVirtualMachineSize,
+    resources as compute_resources,
+)
 from resoto_plugin_azure.resource.base import (
+    AzureLocation,
     AzureSubscription,
     GraphBuilder,
     AzureResource,
@@ -78,6 +81,8 @@ def collect(self) -> None:
                     raise Exception(f"Only Azure resources expected, but got {node}")
             # wait for all work to finish
             queue.wait_for_submitted_work()
+            # filter nodes
+            self.filter_nodes()
             self.core_feedback.progress_done(self.subscription.subscription_id, 1, 1, context=[self.cloud.id])
             log.info(f"[Azure:{self.subscription.safe_name}] Collecting resources done.")
 
@@ -99,3 +104,26 @@ def work_done(_: Future[None]) -> None:
         all_done = GatherFutures.all(group_futures)
         all_done.add_done_callback(work_done)
         return all_done
+
+    def filter_nodes(self) -> None:
+        remove_nodes = []
+
+        def rm_nodes(cls, ignore_kinds: Optional[Type[Any]] = None) -> None:  # type: ignore
+            for node in self.graph.nodes:
+                if not isinstance(node, cls):
+                    continue
+                pred = list(self.graph.predecessors(node))
+                if ignore_kinds is not None:
+                    pred = [p for p in pred if not isinstance(p, ignore_kinds)]
+                if not pred:
+                    remove_nodes.append(node)
+            removed = set()
+            for node in remove_nodes:
+                if node in removed:
+                    continue
+                removed.add(node)
+                self.graph.remove_node(node)
+            log.debug(f"Removing {len(remove_nodes)} unreferenced nodes of type {cls}")
+            remove_nodes.clear()
+
+        rm_nodes(AzureVirtualMachineSize, AzureLocation)

plugins/azure/resoto_plugin_azure/resource/base.py

@@ -4,7 +4,7 @@
 from concurrent.futures import Future
 from datetime import datetime
 from threading import Lock
-from typing import Any, ClassVar, Dict, Optional, TypeVar, List, Tuple, Type, Callable, Union, cast
+from typing import Any, ClassVar, Dict, Optional, TypeVar, List, Type, Callable, cast
 
 from attr import define, field
 from azure.core.utils import CaseInsensitiveDict
@@ -134,35 +134,6 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         # Default behavior: add resource to the namespace
         pass
 
-    def fetch_resources(
-        self,
-        builder: GraphBuilder,
-        service: str,
-        api_version: str,
-        path: str,
-        path_parameters: List[str],
-        query_parameters: List[str],
-        compared_property: Callable[[Json], Union[List[str], str]],
-        binding_property: Callable[[Json], Union[List[str], str]],
-    ) -> List[Tuple[Union[str, List[str]], Union[str, List[str]]]]:
-        """
-        Fetch additional resources from the Azure API for further connection using the connect_in_graph method.
-
-        Returns:
-        List[Tuple[Union[str, List[str]], str]]: A list of tuples containing information to compare and connect the retrieved resources.
-        """
-        resources_api_spec = AzureApiSpec(
-            service=service,
-            version=api_version,
-            path=path,
-            path_parameters=path_parameters,
-            query_parameters=query_parameters,
-            access_path="value",
-            expect_array=True,
-        )
-
-        return [(compared_property(r), binding_property(r)) for r in builder.client.list(resources_api_spec)]
-
     @classmethod
     def collect_resources(
         cls: Type[AzureResourceType], builder: GraphBuilder, **kwargs: Any
@@ -176,7 +147,11 @@ def collect_resources(
         return []
 
     @classmethod
-    def collect(cls: Type[AzureResourceType], raw: List[Json], builder: GraphBuilder) -> List[AzureResourceType]:
+    def collect(
+        cls: Type[AzureResourceType],
+        raw: List[Json],
+        builder: GraphBuilder,
+    ) -> List[AzureResourceType]:
         # Default behavior: iterate over json snippets and for each:
         # - bend the json
         # - transform the result into a resource

plugins/azure/resoto_plugin_azure/resource/compute.py

@@ -2529,11 +2529,7 @@ class AzureVirtualMachine(AzureResource, BaseInstance):
             ]
         },
         "successors": {
-            "default": [
-                "azure_image",
-                "azure_disk",
-                "azure_network_interface",
-            ]
+            "default": ["azure_image", "azure_disk", "azure_network_interface", "azure_virtual_machine_size"]
         },
     }
     mapping: ClassVar[Dict[str, Bender]] = {
@@ -2578,6 +2574,7 @@ class AzureVirtualMachine(AzureResource, BaseInstance):
         "user_data": S("properties", "userData"),
         "virtual_machine_scale_set": S("properties", "virtualMachineScaleSet", "id"),
         "vm_id": S("properties", "vmId"),
+        "location": S("location"),
         "instance_type": S("properties", "hardwareProfile", "vmSize"),
         "instance_status": S("properties", "provisioningState")
         >> MapEnum(InstanceStatusMapping, default=InstanceStatus.UNKNOWN),
@@ -2612,6 +2609,7 @@ class AzureVirtualMachine(AzureResource, BaseInstance):
     user_data: Optional[str] = field(default=None, metadata={'description': 'Userdata for the vm, which must be base-64 encoded. Customer should not pass any secrets in here. Minimum api-version: 2021-03-01.'})  # fmt: skip
     virtual_machine_scale_set: Optional[str] = field(default=None, metadata={"description": ""})
     vm_id: Optional[str] = field(default=None, metadata={'description': 'Specifies the vm unique id which is a 128-bits identifier that is encoded and stored in all azure iaas vms smbios and can be read using platform bios commands.'})  # fmt: skip
+    location: Optional[str] = field(default=None, metadata={"description": "Resource location."})
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if placement_group_id := self.proximity_placement_group:
@@ -2666,6 +2664,10 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
             for network_interface in network_interfaces:
                 if ni_id := network_interface.id:
                     builder.add_edge(self, edge_type=EdgeType.default, clazz=AzureNetworkInterface, id=ni_id)
+        if (vms_type := self.instance_type) and (vm_location := self.location):
+            builder.add_edge(
+                self, edge_type=EdgeType.default, clazz=AzureVirtualMachineSize, name=vms_type, location=vm_location
+            )
 
 
 @define(eq=False, slots=False)
@@ -3205,7 +3207,7 @@ class AzureVirtualMachineSize(AzureResource, BaseInstanceType):
         service="compute",
         version="2023-03-01",
         path="/subscriptions/{subscriptionId}/providers/Microsoft.Compute/locations/{location}/vmSizes",
-        path_parameters=["location", "subscriptionId"],
+        path_parameters=["subscriptionId", "location"],
         query_parameters=["api-version"],
         access_path="value",
         expect_array=True,
@@ -3231,6 +3233,10 @@ class AzureVirtualMachineSize(AzureResource, BaseInstanceType):
     number_of_cores: Optional[int] = field(default=None, metadata={'description': 'The number of cores supported by the virtual machine size. For constrained vcpu capable vm sizes, this number represents the total vcpus of quota that the vm uses. For accurate vcpu count, please refer to https://docs. Microsoft. Com/azure/virtual-machines/constrained-vcpu or https://docs. Microsoft. Com/rest/api/compute/resourceskus/list.'})  # fmt: skip
     os_disk_size_in_mb: Optional[int] = field(default=None, metadata={'description': 'The os disk size, in mb, allowed by the virtual machine size.'})  # fmt: skip
     resource_disk_size_in_mb: Optional[int] = field(default=None, metadata={'description': 'The resource disk size, in mb, allowed by the virtual machine size.'})  # fmt: skip
+    location: Optional[str] = field(default=None, metadata={"description": "Resource location."})
+
+    def pre_process(self, graph_builder: GraphBuilder, source: Json) -> None:
+        self.location = graph_builder.location.name if graph_builder.location else ""
 
 
 resources: List[Type[AzureResource]] = [