[azure][feat] Add postgres service collection (#2155)

Co-authored-by: Matthias Veit <matthias_veit@yahoo.de>

Author: 1101-1

plugins/azure/fix_plugin_azure/collector.py

@@ -34,8 +34,12 @@
     AzureNetworkUsage,
     resources as network_resources,
 )
-from fix_plugin_azure.resource.mysql import AzureMysqlCapability, AzureMysqlServerType, resources as mysql_resources
+from fix_plugin_azure.resource.mysql import AzureMysqlServerType, resources as mysql_resources
 from fix_plugin_azure.resource.sql_server import resources as sql_resources
+from fix_plugin_azure.resource.postgresql import (
+    AzurePostgresqlServerType,
+    resources as postgresql_resources,
+)
 from fix_plugin_azure.resource.storage import AzureStorageAccountUsage, AzureStorageSku, resources as storage_resources
 from fixlib.baseresources import Cloud, GraphRoot, BaseAccount, BaseRegion
 from fixlib.core.actions import CoreFeedback, ErrorAccumulator
@@ -63,6 +67,7 @@ def resource_with_params(clazz: Type[MicrosoftResource], param: str) -> bool:
     + storage_resources
     + sql_resources
     + mysql_resources
+    + postgresql_resources
 )
 all_resources = subscription_resources + graph_resources  # defines all resource kinds. used in model check
 
@@ -233,18 +238,16 @@ def remove_usage_zero_value() -> None:
         rm_nodes(AzureVirtualMachineSize, AzureLocation)
         rm_nodes(AzureExpressRoutePortsLocation, AzureSubscription)
         rm_nodes(AzureNetworkVirtualApplianceSku, AzureSubscription)
-        rm_nodes(AzureDiskType, AzureLocation)
+        rm_nodes(AzureDiskType, AzureSubscription)
         rm_nodes(AzureStorageSku, AzureLocation)
-        rm_nodes(AzureMysqlServerType, AzureLocation)
+        rm_nodes(AzureMysqlServerType, AzureSubscription)
+        rm_nodes(AzurePostgresqlServerType, AzureSubscription)
         remove_usage_zero_value()
 
     def after_collect(self) -> None:
         # Filter unnecessary nodes such as AzureDiskTypePricing
         nodes_to_remove = []
-        node_types = (
-            AzureDiskTypePricing,
-            AzureMysqlCapability,
-        )
+        node_types = (AzureDiskTypePricing,)
 
         for node in self.graph.nodes:
             if not isinstance(node, node_types):

plugins/azure/fix_plugin_azure/resource/base.py

@@ -473,6 +473,29 @@ class AzureUserAssignedIdentity:
     resource_id: Optional[str] = field(default=None, metadata={'description': 'The resource ID of the user assigned identity.'})  # fmt: skip
 
 
+@define(eq=False, slots=False)
+class AzureUserIdentity:
+    kind: ClassVar[str] = "azure_user_identity"
+    mapping: ClassVar[Dict[str, Bender]] = {"client_id": S("clientId"), "principal_id": S("principalId")}
+    client_id: Optional[str] = field(default=None, metadata={'description': 'the client identifier of the Service Principal which this identity represents.'})  # fmt: skip
+    principal_id: Optional[str] = field(default=None, metadata={'description': 'the object identifier of the Service Principal which this identity represents.'})  # fmt: skip
+
+
+@define(eq=False, slots=False)
+class AzureResourceIdentity:
+    kind: ClassVar[str] = "azure_resource_identity"
+    mapping: ClassVar[Dict[str, Bender]] = {
+        "principal_id": S("principalId"),
+        "tenant_id": S("tenantId"),
+        "type": S("type"),
+        "user_assigned_identities": S("userAssignedIdentities"),
+    }
+    principal_id: Optional[str] = field(default=None, metadata={'description': 'The Azure Active Directory principal id.'})  # fmt: skip
+    tenant_id: Optional[str] = field(default=None, metadata={"description": "The Azure Active Directory tenant id."})
+    type: Optional[str] = field(default=None, metadata={'description': 'The identity type. Set this to SystemAssigned in order to automatically create and assign an Azure Active Directory principal for the resource.'})  # fmt: skip
+    user_assigned_identities: Optional[Dict[str, AzureUserIdentity]] = field(default=None, metadata={'description': 'The resource ids of the user assigned identities to use'})  # fmt: skip
+
+
 @define(eq=False, slots=False)
 class AzurePrincipalClient:
     kind: ClassVar[str] = "azure_principal_client"
@@ -589,6 +612,30 @@ class AzureSystemData:
     last_modified_by_type: Optional[str] = field(default=None, metadata={'description': 'The type of identity that last modified the resource.'})  # fmt: skip
 
 
+@define(eq=False, slots=False)
+class AzureTrackedResource:
+    kind: ClassVar[str] = "azure_tracked_resource"
+    mapping: ClassVar[Dict[str, Bender]] = {
+        "location": S("location"),
+        "system_data": S("systemData") >> Bend(AzureSystemData.mapping),
+        "type": S("type"),
+    }
+    location: Optional[str] = field(default=None, metadata={'description': 'The geo-location where the resource lives'})  # fmt: skip
+    system_data: Optional[AzureSystemData] = field(default=None, metadata={'description': 'Metadata pertaining to creation and last modification of the resource.'})  # fmt: skip
+    type: Optional[str] = field(default=None, metadata={'description': 'The type of the resource. E.g. Microsoft.Compute/virtualMachines or Microsoft.Storage/storageAccounts '})  # fmt: skip
+
+
+@define(eq=False, slots=False)
+class AzureProxyResource:
+    kind: ClassVar[str] = "azure_proxy_resource"
+    mapping: ClassVar[Dict[str, Bender]] = {
+        "system_data": S("systemData") >> Bend(AzureSystemData.mapping),
+        "type": S("type"),
+    }
+    system_data: Optional[AzureSystemData] = field(default=None, metadata={'description': 'Metadata pertaining to creation and last modification of the resource.'})  # fmt: skip
+    type: Optional[str] = field(default=None, metadata={'description': 'The type of the resource. E.g. Microsoft.Compute/virtualMachines or Microsoft.Storage/storageAccounts '})  # fmt: skip
+
+
 @define(eq=False, slots=False)
 class AzureSku:
     kind: ClassVar[str] = "azure_sku"

plugins/azure/fix_plugin_azure/resource/compute.py

@@ -6,6 +6,7 @@
 
 from fix_plugin_azure.azure_client import AzureResourceSpec
 from fix_plugin_azure.resource.base import (
+    AzureProxyResource,
     MicrosoftResource,
     MicrosoftResourceType,
     GraphBuilder,
@@ -779,16 +780,7 @@ class AzureDiskTypePricing(MicrosoftResource):
 @define(eq=False, slots=False)
 class AzureDiskType(MicrosoftResource, BaseVolumeType):
     kind: ClassVar[str] = "azure_disk_type"
-    # Define api spec to collect as regional resources
-    api_spec: ClassVar[AzureResourceSpec] = AzureResourceSpec(
-        service="compute",
-        version="2023-01-01-preview",
-        path="",
-        path_parameters=["subscriptionId", "location"],
-        query_parameters=[],
-        access_path="Items",
-        expect_array=True,
-    )
+    # Collect via AzureDisk()
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("skuName"),
         "name": S("skuName"),
@@ -951,29 +943,6 @@ def create_unique_disk_sizes(collected_disks: List[MicrosoftResourceType], build
                     seen_hashes.add(hash_value)
         AzureDiskType.collect(disk_sizes, builder)
 
-    @classmethod
-    def collect_resources(
-        cls: Type[MicrosoftResourceType], builder: GraphBuilder, **kwargs: Any
-    ) -> List[MicrosoftResourceType]:
-        log.debug(f"[Azure:{builder.account.id}] Collecting {cls.__name__} with ({kwargs})")
-        product_names = {"Standard SSD Managed Disks", "Premium SSD Managed Disks", "Standard HDD Managed Disks"}
-        sku_items = []
-        for product_name in product_names:
-            api_spec = AzureResourceSpec(
-                service="compute",
-                version="2023-01-01-preview",
-                path=f"https://prices.azure.com/api/retail/prices?$filter=productName eq '{product_name}' and armRegionName eq "
-                + "'{location}' and unitOfMeasure eq '1/Month' and serviceFamily eq 'Storage' and type eq 'Consumption' and isPrimaryMeterRegion eq true",
-                path_parameters=["location"],
-                query_parameters=["api-version"],
-                access_path="Items",
-                expect_array=True,
-            )
-
-            items = builder.client.list(api_spec, **kwargs)
-            sku_items.extend(items)
-        return cls.collect(sku_items, builder)
-
 
 VolumeStatusMapping = {
     "ActiveSAS": VolumeStatus.IN_USE,
@@ -1109,6 +1078,34 @@ def collect_resources(
             return collected
         return []
 
+    def post_process(self, graph_builder: GraphBuilder, source: Json) -> None:
+        if location := self.location:
+
+            def collect_disk_types() -> None:
+                log.debug(f"[Azure:{graph_builder.account.id}] Collecting AzureDiskType")
+                product_names = {
+                    "Standard SSD Managed Disks",
+                    "Premium SSD Managed Disks",
+                    "Standard HDD Managed Disks",
+                }
+                sku_items = []
+                for product_name in product_names:
+                    api_spec = AzureResourceSpec(
+                        service="compute",
+                        version="2023-01-01-preview",
+                        path=f"https://prices.azure.com/api/retail/prices?$filter=productName eq '{product_name}' and armRegionName eq '{location}' and unitOfMeasure eq '1/Month' and serviceFamily eq 'Storage' and type eq 'Consumption' and isPrimaryMeterRegion eq true",
+                        path_parameters=[],
+                        query_parameters=["api-version"],
+                        access_path="Items",
+                        expect_array=True,
+                    )
+
+                    items = graph_builder.client.list(api_spec)
+                    sku_items.extend(items)
+                AzureDiskType.collect(sku_items, graph_builder)
+
+            graph_builder.submit_work(service_name, collect_disk_types)
+
     @classmethod
     def collect_usage_metrics(
         cls: Type[MicrosoftResource], builder: GraphBuilder, collected_resources: List[MicrosoftResourceType]
@@ -1584,15 +1581,6 @@ class AzureRestorePointCollectionSourceProperties:
     location: Optional[str] = field(default=None, metadata={'description': 'Location of the source resource used to create this restore point collection.'})  # fmt: skip
 
 
-@define(eq=False, slots=False)
-class AzureProxyResource:
-    kind: ClassVar[str] = "azure_proxy_resource"
-    mapping: ClassVar[Dict[str, Bender]] = {"id": S("id"), "name": S("name"), "type": S("type")}
-    id: Optional[str] = field(default=None, metadata={"description": "Resource id."})
-    name: Optional[str] = field(default=None, metadata={"description": "Resource name."})
-    type: Optional[str] = field(default=None, metadata={"description": "Resource type."})
-
-
 @define(eq=False, slots=False)
 class AzureVMSizeProperties:
     kind: ClassVar[str] = "azure_vm_size_properties"
@@ -2947,6 +2935,29 @@ def collect_instance_status() -> None:
                 if not instance_status_set:
                     self.instance_status = InstanceStatus.UNKNOWN
 
+        if location := self.location:
+
+            def collect_vm_sizes() -> None:
+                api_spec = AzureResourceSpec(
+                    service="compute",
+                    version="2023-03-01",
+                    path="/subscriptions/{subscriptionId}/providers/Microsoft.Compute/locations/"
+                    + f"{location}/vmSizes",
+                    path_parameters=["subscriptionId"],
+                    query_parameters=["api-version"],
+                    access_path="value",
+                    expect_array=True,
+                )
+                items = graph_builder.client.list(api_spec)
+                if not items:
+                    return
+                # Set location for further connect_in_graph method
+                for item in items:
+                    item["location"] = location
+                AzureVirtualMachineSize.collect(items, graph_builder)
+
+            graph_builder.submit_work(service_name, collect_vm_sizes)
+
         graph_builder.submit_work(service_name, collect_instance_status)
 
     @classmethod
@@ -3663,15 +3674,6 @@ def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
 @define(eq=False, slots=False)
 class AzureVirtualMachineSize(MicrosoftResource, BaseInstanceType):
     kind: ClassVar[str] = "azure_virtual_machine_size"
-    api_spec: ClassVar[AzureResourceSpec] = AzureResourceSpec(
-        service="compute",
-        version="2023-03-01",
-        path="/subscriptions/{subscriptionId}/providers/Microsoft.Compute/locations/{location}/vmSizes",
-        path_parameters=["subscriptionId", "location"],
-        query_parameters=["api-version"],
-        access_path="value",
-        expect_array=True,
-    )
     mapping: ClassVar[Dict[str, Bender]] = {
         "id": S("name"),
         "tags": S("tags", default={}),
@@ -3684,6 +3686,7 @@ class AzureVirtualMachineSize(MicrosoftResource, BaseInstanceType):
         "instance_type": S("name"),
         "instance_cores": S("numberOfCores"),
         "instance_memory": S("memoryInMB") >> F(lambda x: int(x) / 1024),
+        "location": S("location"),
     }
     _is_provider_link: ClassVar[bool] = False
     max_data_disk_count: Optional[int] = field(default=None, metadata={'description': 'The maximum number of data disks that can be attached to the virtual machine size.'})  # fmt: skip
@@ -3693,9 +3696,6 @@ class AzureVirtualMachineSize(MicrosoftResource, BaseInstanceType):
     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 ""
-
     def after_collect(self, builder: GraphBuilder, source: Json) -> None:
         if (location := self.location) and (sku_name := self.name):
             api_spec = AzureResourceSpec(

plugins/azure/fix_plugin_azure/resource/containerservice.py

@@ -6,8 +6,8 @@
 
 from fix_plugin_azure.azure_client import AzureResourceSpec
 from fix_plugin_azure.resource.base import (
+    AzureTrackedResource,
     MicrosoftResource,
-    AzureSystemData,
     GraphBuilder,
     AzureExtendedLocation,
     AzureUserAssignedIdentity,
@@ -22,25 +22,6 @@
 log = logging.getLogger("fix.plugins.azure")
 
 
-@define(eq=False, slots=False)
-class AzureTrackedResource:
-    kind: ClassVar[str] = "azure_tracked_resource"
-    mapping: ClassVar[Dict[str, Bender]] = {
-        "id": S("id"),
-        "location": S("location"),
-        "name": S("name"),
-        "system_data": S("systemData") >> Bend(AzureSystemData.mapping),
-        "tags": S("tags"),
-        "type": S("type"),
-    }
-    id: Optional[str] = field(default=None, metadata={'description': 'Fully qualified resource ID for the resource. Ex - /subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/{resourceProviderNamespace}/{resourceType}/{resourceName}'})  # fmt: skip
-    location: Optional[str] = field(default=None, metadata={'description': 'The geo-location where the resource lives'})  # fmt: skip
-    name: Optional[str] = field(default=None, metadata={"description": "The name of the resource"})
-    system_data: Optional[AzureSystemData] = field(default=None, metadata={'description': 'Metadata pertaining to creation and last modification of the resource.'})  # fmt: skip
-    tags: Optional[Dict[str, str]] = field(default=None, metadata={"description": "Resource tags."})
-    type: Optional[str] = field(default=None, metadata={'description': 'The type of the resource. E.g. Microsoft.Compute/virtualMachines or Microsoft.Storage/storageAccounts '})  # fmt: skip
-
-
 @define(eq=False, slots=False)
 class AzureAPIServerAccessProfile:
     kind: ClassVar[str] = "azure_api_server_access_profile"
@@ -915,21 +896,6 @@ def _get_poolnames_and_vmss_ids(self, builder: GraphBuilder) -> List[Tuple[str,
         ]
 
 
-@define(eq=False, slots=False)
-class AzureProxyResource:
-    kind: ClassVar[str] = "azure_proxy_resource"
-    mapping: ClassVar[Dict[str, Bender]] = {
-        "id": S("id"),
-        "name": S("name"),
-        "system_data": S("systemData") >> Bend(AzureSystemData.mapping),
-        "type": S("type"),
-    }
-    id: Optional[str] = field(default=None, metadata={'description': 'Fully qualified resource ID for the resource. Ex - /subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/{resourceProviderNamespace}/{resourceType}/{resourceName}'})  # fmt: skip
-    name: Optional[str] = field(default=None, metadata={"description": "The name of the resource"})
-    system_data: Optional[AzureSystemData] = field(default=None, metadata={'description': 'Metadata pertaining to creation and last modification of the resource.'})  # fmt: skip
-    type: Optional[str] = field(default=None, metadata={'description': 'The type of the resource. E.g. Microsoft.Compute/virtualMachines or Microsoft.Storage/storageAccounts '})  # fmt: skip
-
-
 @define(eq=False, slots=False)
 class AzureCompatibleVersions:
     kind: ClassVar[str] = "azure_compatible_versions"