res mode: Fix a bug that some resource types are not exported when --include-role-assignment is specified

go.mod

@@ -24,7 +24,7 @@ require (
 	github.com/hashicorp/terraform-json v0.24.0
 	github.com/hexops/gotextdiff v1.0.3
 	github.com/magodo/armid v0.0.0-20240524082432-7ce06ae46c33
-	github.com/magodo/azlist v0.0.0-20250429005249-2b51aca0f1e1
+	github.com/magodo/azlist v0.0.0-20250807062828-b842fde1b538
 	github.com/magodo/aztft v0.3.1-0.20250606003245-7222b835cdcc
 	github.com/magodo/slog2hclog v0.0.0-20240614031327-090ebd72a033
 	github.com/magodo/spinner v0.0.0-20240524082745-3a2305db1bdc

go.sum

@@ -256,8 +256,8 @@ github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69
 github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/magodo/armid v0.0.0-20240524082432-7ce06ae46c33 h1:KmQ16pNsI7DaELU+CbqZKXVdvkE/YXqMH6LLkp6rw/Y=
 github.com/magodo/armid v0.0.0-20240524082432-7ce06ae46c33/go.mod h1:rR8E7zfGMbmfnSQvrkFiWYdhrfTqsVSltelnZB09BwA=
-github.com/magodo/azlist v0.0.0-20250429005249-2b51aca0f1e1 h1:fZiuTRpB70qS2b7QPtXPgWi+tTTaFQYiLPmjxGyWS7M=
-github.com/magodo/azlist v0.0.0-20250429005249-2b51aca0f1e1/go.mod h1:uRr4kq5cpk4NIUuT1q6ZYYfYSIHO5Z7XlHeN9c2mlKs=
+github.com/magodo/azlist v0.0.0-20250807062828-b842fde1b538 h1:+OY84fSoUN5cdamJQ2TOWjCYCK6d9AktZefEz57ZxXQ=
+github.com/magodo/azlist v0.0.0-20250807062828-b842fde1b538/go.mod h1:uRr4kq5cpk4NIUuT1q6ZYYfYSIHO5Z7XlHeN9c2mlKs=
 github.com/magodo/aztft v0.3.1-0.20250606003245-7222b835cdcc h1:bLGK1TzoWNdEgal5aV+68LD0PKUyrRXlkTz9quVIn+E=
 github.com/magodo/aztft v0.3.1-0.20250606003245-7222b835cdcc/go.mod h1:t83k9NbdobHHhjtz6MFetOSRqiKfaNfbBpowxZyzYtw=
 github.com/magodo/slog2hclog v0.0.0-20240614031327-090ebd72a033 h1:K2seYsMAzoICCLdDe7uU2WyaACLW+tvdTWG3QB+pyec=

internal/meta/config_info_test.go

@@ -7,9 +7,18 @@ import (
 	"github.com/Azure/aztfexport/internal/tfaddr"
 	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/hclwrite"
+	"github.com/magodo/armid"
 	"github.com/stretchr/testify/assert"
 )
 
+func mustParseResourceId(id string) armid.ResourceId {
+	parsed, err := armid.ParseResourceId(id)
+	if err != nil {
+		panic(err)
+	}
+	return parsed
+}
+
 func newConfigInfo(azid, tfid, tfaddr, hcl string, deps *Dependencies) ConfigInfo {
 	cinfo := ConfigInfo{
 		ImportItem: ImportItem{

internal/meta/meta_query.go

@@ -122,7 +122,7 @@ func (meta MetaQuery) queryResourceSet(ctx context.Context, predicate string, re
 	if err != nil {
 		return nil, fmt.Errorf("building azlister: %v", err)
 	}
-	result, err := lister.List(ctx, predicate)
+	result, err := lister.ListByQuery(ctx, predicate)
 	if err != nil {
 		return nil, fmt.Errorf("listing resource set: %w", err)
 	}

internal/meta/meta_res.go

@@ -3,7 +3,6 @@ package meta
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	"github.com/Azure/aztfexport/internal/resourceset"
 	"github.com/Azure/aztfexport/internal/tfaddr"
@@ -62,38 +61,47 @@ func (meta MetaResource) ScopeName() string {
 }
 
 func (meta *MetaResource) ListResource(ctx context.Context) (ImportList, error) {
-	var resources []resourceset.AzureResource
+	var rl []resourceset.AzureResource
 	for _, id := range meta.AzureIds {
-		resources = append(resources, resourceset.AzureResource{Id: id})
+		rl = append(rl, resourceset.AzureResource{Id: id})
 	}
 
-	rset := &resourceset.AzureResourceSet{
-		Resources: resources,
-	}
 	if meta.includeRoleAssignment {
 		var err error
-		rset, err = meta.queryResourceSet(ctx, resources)
+		rl, err = meta.listByIds(ctx, rl)
 		if err != nil {
-			return nil, fmt.Errorf("querying resource set: %v", err)
+			return nil, fmt.Errorf("querying extension resources: %v", err)
 		}
 	}
 
 	meta.Logger().Debug("Azure Resource set map to TF resource set")
-
-	var rl []resourceset.TFResource
+	rset := &resourceset.AzureResourceSet{Resources: rl}
+	var tfl []resourceset.TFResource
 	if meta.useAzAPI() {
-		rl = rset.ToTFAzAPIResources()
+		tfl = rset.ToTFAzAPIResources()
 	} else {
-		rl = rset.ToTFAzureRMResources(meta.Logger(), meta.parallelism, meta.azureSDKCred, meta.azureSDKClientOpt)
+		tfl = rset.ToTFAzureRMResources(meta.Logger(), meta.parallelism, meta.azureSDKCred, meta.azureSDKClientOpt)
 	}
 
-	var l ImportList
+	// Split the specified resources and the extension resources
+	var tfrl, tfel []resourceset.TFResource
+	for _, tfres := range tfl {
+		rmap := map[string]bool{}
+		for _, r := range rl {
+			rmap[r.Id.String()] = true
+		}
+		if rmap[tfres.AzureId.String()] {
+			tfrl = append(tfrl, tfres)
+		} else {
+			tfel = append(tfel, tfres)
+		}
+	}
 
-	originalRl, extensionRl := splitOriginalAndExtension(rl, resources)
+	var l ImportList
 
 	// The ResourceName and ResourceType are only honored for single non-role-assignment-resource
-	if len(originalRl) == 1 {
-		res := originalRl[0]
+	if len(tfrl) == 1 {
+		res := tfrl[0]
 
 		// Honor the ResourceName
 		name := meta.ResourceName
@@ -134,21 +142,20 @@ func (meta *MetaResource) ListResource(ctx context.Context) (ImportList, error)
 		}
 		l = append(l, item)
 	} else {
-		meta.appendRlToImportList(originalRl, &l)
+		l = append(l, meta.toImportList(tfrl, 0)...)
 	}
-
-	meta.appendRlToImportList(extensionRl, &l)
+	l = append(l, meta.toImportList(tfel, len(tfrl))...)
 
 	l = meta.excludeImportList(l)
-
 	return l, nil
 }
 
-func (meta MetaResource) appendRlToImportList(rl []resourceset.TFResource, l *ImportList) {
-	for _, res := range rl {
+func (meta MetaResource) toImportList(rl []resourceset.TFResource, fromIdx int) ImportList {
+	var l ImportList
+	for idx, res := range rl {
 		tfAddr := tfaddr.TFAddr{
 			Type: "",
-			Name: fmt.Sprintf("%s%d%s", meta.resourceNamePrefix, len(*l), meta.resourceNameSuffix),
+			Name: fmt.Sprintf("%s%d%s", meta.resourceNamePrefix, idx+fromIdx, meta.resourceNameSuffix),
 		}
 		item := ImportItem{
 			AzureResourceID: res.AzureId,
@@ -163,14 +170,12 @@ func (meta MetaResource) appendRlToImportList(rl []resourceset.TFResource, l *Im
 			item.IsRecommended = true
 		}
 
-		*l = append(*l, item)
+		l = append(l, item)
 	}
+	return l
 }
 
-func (meta MetaResource) queryResourceSet(ctx context.Context, resources []resourceset.AzureResource) (*resourceset.AzureResourceSet, error) {
-	var rl []resourceset.AzureResource
-	var quotedIds []string
-
+func (meta MetaResource) listByIds(ctx context.Context, resources []resourceset.AzureResource) ([]resourceset.AzureResource, error) {
 	opt := azlist.Option{
 		Logger:                 meta.logger.WithGroup("azlist"),
 		SubscriptionId:         meta.subscriptionId,
@@ -185,15 +190,17 @@ func (meta MetaResource) queryResourceSet(ctx context.Context, resources []resou
 		return nil, fmt.Errorf("building azlister for listing resources: %v", err)
 	}
 
-	for _, res := range resources {
-		quotedIds = append(quotedIds, fmt.Sprintf("%q", res.Id.String()))
+	var ids []string
+	for _, r := range resources {
+		ids = append(ids, r.Id.String())
 	}
 
-	result, err := lister.List(ctx, fmt.Sprintf("id in (%s)", strings.Join(quotedIds, ", ")))
+	result, err := lister.ListByIds(ctx, ids)
 	if err != nil {
-		return nil, fmt.Errorf("listing resources: %w", err)
+		return nil, fmt.Errorf("listing extension resources: %w", err)
 	}
 
+	var rl []resourceset.AzureResource
 	for _, res := range result.Resources {
 		res := resourceset.AzureResource{
 			Id:         res.Id,
@@ -202,24 +209,5 @@ func (meta MetaResource) queryResourceSet(ctx context.Context, resources []resou
 		rl = append(rl, res)
 	}
 
-	return &resourceset.AzureResourceSet{
-		Resources: rl,
-	}, nil
-}
-
-func splitOriginalAndExtension(combined []resourceset.TFResource, requested []resourceset.AzureResource) (original []resourceset.TFResource, extension []resourceset.TFResource) {
-	idRequested := make(map[string]bool, len(requested))
-	for _, res := range requested {
-		idRequested[res.Id.String()] = true
-	}
-
-	for _, res := range combined {
-		if idRequested[res.AzureId.String()] {
-			original = append(original, res)
-		} else {
-			extension = append(extension, res)
-		}
-	}
-
-	return
+	return rl, nil
 }

internal/meta/meta_rg.go

@@ -107,7 +107,7 @@ func (meta MetaResourceGroup) queryResourceSet(ctx context.Context, rg string) (
 	if err != nil {
 		return nil, fmt.Errorf("building azlister for listing resource group only: %v", err)
 	}
-	result, err := lister.List(ctx, fmt.Sprintf("name == %q", rg))
+	result, err := lister.ListByQuery(ctx, fmt.Sprintf("name == %q", rg))
 	if err != nil {
 		return nil, fmt.Errorf("listing resource group only: %w", err)
 	}
@@ -138,7 +138,7 @@ func (meta MetaResourceGroup) queryResourceSet(ctx context.Context, rg string) (
 	if err != nil {
 		return nil, fmt.Errorf("building azlister for listing resource group: %v", err)
 	}
-	result, err = lister.List(ctx, fmt.Sprintf("resourceGroup =~ %q", rg))
+	result, err = lister.ListByQuery(ctx, fmt.Sprintf("resourceGroup =~ %q", rg))
 	if err != nil {
 		return nil, fmt.Errorf("listing resource group: %w", err)
 	}