Fixed azure_compute_availability_set always returns null value for th…

…e column virtual_machines Closes #845
turbot · Nov 27, 2024 · 14ed613 · 14ed613
1 parent 405604f
commit 14ed613
Showing 1 changed file with 39 additions and 4 deletions.
diff --git a/azure/table_azure_compute_availability_set.go b/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"
@@ -88,7 +89,8 @@ func tableAzureComputeAvailabilitySet(_ context.Context) *plugin.Table {
 				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(extractVirtualMechinesForScaleset),
 			},
 
 			// Steampipe standard columns
@@ -175,10 +177,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()
+
+	if h.Item != nil {
+		availabilitySet := h.Item.(compute.AvailabilitySet)
+		id := availabilitySet.ID
+		resourceGroup = strings.Split(*id, "/")[4]
+		name = *availabilitySet.Name
+	}
 
-	name := d.EqualsQuals["name"].GetStringValue()
-	resourceGroup := d.EqualsQuals["resource_group"].GetStringValue()
+	// Empty check
+	if name == "" || resourceGroup == "" {
+		return nil, nil
+	}
 
 	session, err := GetNewSession(ctx, d, "MANAGEMENT")
 	if err != nil {
@@ -201,3 +215,24 @@ func getAzureComputeAvailabilitySet(ctx context.Context, d *plugin.QueryData, h
 
 	return nil, nil
 }
+
+//// UTILITY FUNCTION
+
+func extractVirtualMechinesForScaleset(ctx context.Context, d *transform.TransformData) (interface{}, error) {
+	scalSet := d.HydrateItem.(compute.AvailabilitySet)
+	var properties []map[string]interface{}
+	// AvailabilitySetProperties.VirtualMachines
+
+	if scalSet.VirtualMachines != nil {
+		vmProperies := scalSet.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
+}