Fixed table azure_compute_availability_set always returns null value for the column virtual_machines Closes #845 (#856)

Author: ParthaI

azure/table_azure_compute_availability_set.go

@@ -2,6 +2,7 @@ package azure
 
 import (
 	"context"
+	"strings"
 
 	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
 	"github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto"
@@ -82,13 +83,15 @@ func tableAzureComputeAvailabilitySet(_ context.Context) *plugin.Table {
 				Name:        "status",
 				Description: "The resource status information",
 				Type:        proto.ColumnType_JSON,
-				Transform:   transform.FromField("AvailabilitySetProperties.Statuses"),
+				Hydrate:     getAzureComputeAvailabilitySet,
+				Transform:   transform.From(extractStatusForAvailabilitySet),
 			},
 			{
 				Name:        "virtual_machines",
 				Description: "A list of references to all virtual machines in the availability set",
 				Type:        proto.ColumnType_JSON,
-				Transform:   transform.FromField("AvailabilitySetProperties.VirtualMachines"),
+				Hydrate:     getAzureComputeAvailabilitySet,
+				Transform:   transform.From(extractVirtualMachinesForAvailabilitySet),
 			},
 
 			// Steampipe standard columns
@@ -175,10 +178,22 @@ func listAzureComputeAvailabilitySets(ctx context.Context, d *plugin.QueryData,
 //// HYDRATE FUNCTIONS ////
 
 func getAzureComputeAvailabilitySet(ctx context.Context, d *plugin.QueryData, h *plugin.HydrateData) (interface{}, error) {
-	plugin.Logger(ctx).Trace("getAzureComputeAvailabilitySet")
+	resourceGroup, name := "", ""
 
-	name := d.EqualsQuals["name"].GetStringValue()
-	resourceGroup := d.EqualsQuals["resource_group"].GetStringValue()
+	name = d.EqualsQuals["name"].GetStringValue()
+	resourceGroup = d.EqualsQuals["resource_group"].GetStringValue()
+
+	if h.Item != nil {
+		availabilitySet := h.Item.(compute.AvailabilitySet)
+		id := availabilitySet.ID
+		resourceGroup = strings.Split(*id, "/")[4]
+		name = *availabilitySet.Name
+	}
+
+	// Empty check
+	if name == "" || resourceGroup == "" {
+		return nil, nil
+	}
 
 	session, err := GetNewSession(ctx, d, "MANAGEMENT")
 	if err != nil {
@@ -201,3 +216,51 @@ func getAzureComputeAvailabilitySet(ctx context.Context, d *plugin.QueryData, h
 
 	return nil, nil
 }
+
+//// UTILITY FUNCTION
+
+func extractVirtualMachinesForAvailabilitySet(_ context.Context, d *transform.TransformData) (interface{}, error) {
+	availabilitySet := d.HydrateItem.(compute.AvailabilitySet)
+	var properties []map[string]interface{}
+
+	if availabilitySet.AvailabilitySetProperties != nil && availabilitySet.AvailabilitySetProperties.VirtualMachines != nil {
+		vmProperies := availabilitySet.AvailabilitySetProperties
+		for _, i := range *vmProperies.VirtualMachines {
+			objectMap := make(map[string]interface{})
+			if i.ID != nil {
+				objectMap["id"] = i.ID
+			}
+			properties = append(properties, objectMap)
+		}
+	}
+
+	return properties, nil
+}
+
+func extractStatusForAvailabilitySet(_ context.Context, d *transform.TransformData) (interface{}, error) {
+	availabilitySet := d.HydrateItem.(compute.AvailabilitySet)
+	var properties []map[string]interface{}
+
+	if availabilitySet.AvailabilitySetProperties != nil && availabilitySet.AvailabilitySetProperties.Statuses != nil {
+		properies := availabilitySet.AvailabilitySetProperties
+		for _, i := range *properies.Statuses {
+			objectMap := make(map[string]interface{})
+			if i.Code != nil {
+				objectMap["code"] = i.Code
+			}
+			if i.DisplayStatus != nil {
+				objectMap["displayStatus"] = i.DisplayStatus
+			}
+			if i.Level != "" {
+				objectMap["level"] = i.Level
+			}
+			if i.Message != nil {
+				objectMap["message"] = i.Message
+			}
+
+			properties = append(properties, objectMap)
+		}
+	}
+
+	return properties, nil
+}