Replace --include-role-assignment with --include-extension to support more extension resources

command_before_func.go

@@ -6,6 +6,7 @@ import (
 	"slices"
 	"strings"
 
+	"github.com/Azure/aztfexport/internal/meta"
 	"github.com/Azure/aztfexport/internal/utils"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resourcegraph/armresourcegraph"
 	"github.com/hashicorp/terraform-config-inspect/tfconfig"
@@ -52,6 +53,12 @@ func commandBeforeFunc(fset *FlagSet, mode Mode) func(ctx *cli.Context) error {
 			}
 		}
 
+		for _, ext := range fset.flagIncludeExtension.Value() {
+			if !slices.Contains(meta.SupportedExtensionResourceTypes, ext) {
+				return fmt.Errorf("invalid value of `--include-extension`: %v is not supported", ext)
+			}
+		}
+
 		if err := conflictArgs([]argDesc{
 			{
 				name:  "--client-id",

flag.go

@@ -52,7 +52,7 @@ type FlagSet struct {
 	flagExcludeTerraformResource     cli.StringSlice
 	flagExcludeTerraformResourceFile string
 	flagIncludeManagedResource       bool
-	flagIncludeRoleAssignment        bool
+	flagIncludeExtension             cli.StringSlice
 
 	// common flags (auth)
 	flagEnv                       string
@@ -229,8 +229,10 @@ func (flag FlagSet) DescribeCLI(mode Mode) string {
 	if flag.hflagTFClientPluginPath != "" {
 		args = append(args, "--tfclient-plugin-path="+flag.hflagTFClientPluginPath)
 	}
-	if flag.flagIncludeRoleAssignment {
-		args = append(args, "--include-role-assignment=true")
+	if v := flag.flagIncludeExtension.Value(); len(v) != 0 {
+		for _, e := range v {
+			args = append(args, "--include-extension="+e)
+		}
 	}
 	if flag.flagIncludeManagedResource {
 		args = append(args, "--include-managed-resource=true")

internal/meta/extension_resource.go

@@ -6,35 +6,84 @@ import (
 	"github.com/magodo/azlist/azlist"
 )
 
+// Supported extension resource types
+const (
+	ExtensionResourceTypeRoleAssignments    = "role-assignments"
+	ExtensionResourceTypeLocks              = "locks"
+	ExtensionResourceTypeDiagnosticSettings = "diagnostics-settings"
+)
+
+// SupportedExtensionResourceTypes is the list of supported extension resource types.
+var SupportedExtensionResourceTypes = []string{
+	ExtensionResourceTypeRoleAssignments,
+	ExtensionResourceTypeLocks,
+	ExtensionResourceTypeDiagnosticSettings,
+}
+
 type extBuilder struct {
-	includeRoleAssignment bool
+	includeExtensions []string
 }
 
 func (b extBuilder) Build() []azlist.ExtensionResource {
 	var el []azlist.ExtensionResource
-	if b.includeRoleAssignment {
-		el = append(el, azlist.ExtensionResource{
-			Type: "Microsoft.Authorization/roleAssignments",
-			Filter: func(res, extensionRes map[string]interface{}) bool {
-				idRaw, ok := res["id"]
-				if !ok {
-					return false
-				}
-				id := idRaw.(string)
-
-				propsRaw, ok := extensionRes["properties"]
-				if !ok {
-					return false
-				}
-				scopeRaw, ok := propsRaw.(map[string]interface{})["scope"]
-				if !ok {
-					return false
-				}
-				scope := scopeRaw.(string)
-
-				return strings.EqualFold(id, scope)
-			},
-		})
+	for _, ext := range b.includeExtensions {
+		switch ext {
+		case ExtensionResourceTypeRoleAssignments:
+			el = append(el, azlist.ExtensionResource{
+				Type: "Microsoft.Authorization/roleAssignments",
+				Filter: func(res, extensionRes map[string]any) bool {
+					idRaw, ok := res["id"]
+					if !ok {
+						return false
+					}
+					id := idRaw.(string)
+
+					propsRaw, ok := extensionRes["properties"]
+					if !ok {
+						return false
+					}
+					scopeRaw, ok := propsRaw.(map[string]any)["scope"]
+					if !ok {
+						return false
+					}
+					scope := scopeRaw.(string)
+
+					return strings.EqualFold(id, scope)
+				},
+			})
+		case ExtensionResourceTypeLocks:
+			el = append(el, azlist.ExtensionResource{
+				Type: "Microsoft.Authorization/locks",
+				Filter: func(res, extensionRes map[string]any) bool {
+					idRaw, ok := res["id"]
+					if !ok {
+						return false
+					}
+					id := idRaw.(string)
+
+					lockIdRaw, ok := extensionRes["id"]
+					if !ok {
+						return false
+					}
+					lockId := lockIdRaw.(string)
+
+					// A management lock's id is in form of "<scope>/providers/Microsoft.Authorization/locks/<name>".
+					// Derive the scope that the lock applies to and only keep the lock if that scope is
+					// one of the exported resources (i.e. matches the parent resource being listed).
+					idx := strings.Index(strings.ToLower(lockId), "/providers/microsoft.authorization/locks/")
+					if idx == -1 {
+						return false
+					}
+					scope := lockId[:idx]
+
+					return strings.EqualFold(id, scope)
+				},
+			})
+		case ExtensionResourceTypeDiagnosticSettings:
+			el = append(el, azlist.ExtensionResource{
+				Type: "Microsoft.Insights/diagnosticSettings",
+			})
+		}
 	}
 
 	return el

internal/meta/meta_query.go

@@ -16,7 +16,7 @@ type MetaQuery struct {
 	argPredicate                 string
 	recursiveQuery               bool
 	resourceNameExpander         *nameExpander
-	includeRoleAssignment        bool
+	includeExtensions            []string
 	includeManagedResource       bool
 	includeResourceGroup         bool
 	argTable                     string
@@ -34,7 +34,7 @@ func NewMetaQuery(cfg config.Config) (*MetaQuery, error) {
 		baseMeta:                     *baseMeta,
 		argPredicate:                 cfg.ARGPredicate,
 		recursiveQuery:               cfg.RecursiveQuery,
-		includeRoleAssignment:        cfg.IncludeRoleAssignment,
+		includeExtensions:            cfg.IncludeExtensions,
 		includeManagedResource:       cfg.IncludeManagedResource,
 		includeResourceGroup:         cfg.IncludeResourceGroup,
 		argTable:                     cfg.ARGTable,
@@ -112,7 +112,7 @@ func (meta MetaQuery) queryResourceSet(ctx context.Context, predicate string, re
 		Parallelism:                 meta.parallelism,
 		Recursive:                   recursive,
 		IncludeResourceGroup:        meta.includeResourceGroup,
-		ExtensionResourceTypes:      extBuilder{includeRoleAssignment: meta.includeRoleAssignment}.Build(),
+		ExtensionResourceTypes:      extBuilder{includeExtensions: meta.includeExtensions}.Build(),
 		IncludeManaged:              meta.includeManagedResource,
 		ARGTable:                    meta.argTable,
 		ARGAuthorizationScopeFilter: meta.argAuthenticationScopeFilter,

internal/meta/meta_res.go

@@ -19,7 +19,7 @@ type MetaResource struct {
 	ResourceType         string
 	resourceNameExpander *nameExpander
 
-	includeRoleAssignment  bool
+	includeExtensions      []string
 	includeManagedResource bool
 	includeResourceGroup   bool
 	recursive              bool
@@ -48,7 +48,7 @@ func NewMetaResource(cfg config.Config) (*MetaResource, error) {
 		ResourceName:           cfg.TFResourceName,
 		ResourceType:           cfg.TFResourceType,
 		recursive:              cfg.RecursiveQuery,
-		includeRoleAssignment:  cfg.IncludeRoleAssignment,
+		includeExtensions:      cfg.IncludeExtensions,
 		includeManagedResource: cfg.IncludeManagedResource,
 		includeResourceGroup:   cfg.IncludeResourceGroup,
 	}
@@ -181,7 +181,7 @@ func (meta MetaResource) toImportList(rl []resourceset.TFResource) ImportList {
 func (meta MetaResource) listByIds(ctx context.Context, resources []resourceset.AzureResource) ([]resourceset.AzureResource, error) {
 	// In case only the specified resource ids are to list, no need to call azlist as the resource ids are already provided.
 	// This avoids the case that the schema in azlist is a bit behind, causing failed to find the API version for reading each resource.
-	if !meta.recursive && !meta.includeRoleAssignment && !meta.includeManagedResource && !meta.includeResourceGroup {
+	if !meta.recursive && len(meta.includeExtensions) == 0 && !meta.includeManagedResource && !meta.includeResourceGroup {
 		var rl []resourceset.AzureResource
 		for _, r := range resources {
 			res := resourceset.AzureResource{
@@ -198,7 +198,7 @@ func (meta MetaResource) listByIds(ctx context.Context, resources []resourceset.
 		Cred:                   meta.azureSDKCred,
 		ClientOpt:              meta.azureSDKClientOpt,
 		Parallelism:            meta.parallelism,
-		ExtensionResourceTypes: extBuilder{includeRoleAssignment: meta.includeRoleAssignment}.Build(),
+		ExtensionResourceTypes: extBuilder{includeExtensions: meta.includeExtensions}.Build(),
 		IncludeManaged:         meta.includeManagedResource,
 		IncludeResourceGroup:   meta.includeResourceGroup,
 		Recursive:              meta.recursive,

internal/meta/meta_rg.go

@@ -14,7 +14,7 @@ type MetaResourceGroup struct {
 	baseMeta
 	resourceGroup          string
 	resourceNameExpander   *nameExpander
-	includeRoleAssignment  bool
+	includeExtensions      []string
 	includeManagedResource bool
 }
 
@@ -28,7 +28,7 @@ func NewMetaResourceGroup(cfg config.Config) (*MetaResourceGroup, error) {
 	meta := &MetaResourceGroup{
 		baseMeta:               *baseMeta,
 		resourceGroup:          cfg.ResourceGroupName,
-		includeRoleAssignment:  cfg.IncludeRoleAssignment,
+		includeExtensions:      cfg.IncludeExtensions,
 		includeManagedResource: cfg.IncludeManagedResource,
 	}
 	meta.resourceNameExpander = newNameExpander(cfg.ResourceNamePattern)
@@ -97,7 +97,7 @@ func (meta MetaResourceGroup) queryResourceSet(ctx context.Context, rg string) (
 		Cred:                   meta.azureSDKCred,
 		ClientOpt:              meta.azureSDKClientOpt,
 		Parallelism:            meta.parallelism,
-		ExtensionResourceTypes: extBuilder{includeRoleAssignment: meta.includeRoleAssignment}.Build(),
+		ExtensionResourceTypes: extBuilder{includeExtensions: meta.includeExtensions}.Build(),
 		IncludeManaged:         meta.includeManagedResource,
 		ARGTable:               "ResourceContainers",
 	}
@@ -128,7 +128,7 @@ func (meta MetaResourceGroup) queryResourceSet(ctx context.Context, rg string) (
 		Cred:                   meta.azureSDKCred,
 		ClientOpt:              meta.azureSDKClientOpt,
 		Parallelism:            meta.parallelism,
-		ExtensionResourceTypes: extBuilder{includeRoleAssignment: meta.includeRoleAssignment}.Build(),
+		ExtensionResourceTypes: extBuilder{includeExtensions: meta.includeExtensions}.Build(),
 		IncludeManaged:         meta.includeManagedResource,
 		Recursive:              true,
 	}

main.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/Azure/aztfexport/internal/cfgfile"
 	internalconfig "github.com/Azure/aztfexport/internal/config"
+	"github.com/Azure/aztfexport/internal/meta"
 	"github.com/Azure/aztfexport/pkg/telemetry"
 	"github.com/gofrs/uuid"
 	"github.com/pkg/profile"
@@ -258,11 +259,11 @@ func main() {
 			Usage:       "Path to a file recording the excluded resources from being exported based on the Terraform resource type, one per line",
 			Destination: &flagset.flagExcludeTerraformResourceFile,
 		},
-		&cli.BoolFlag{
-			Name:        "include-role-assignment",
-			EnvVars:     []string{"AZTFEXPORT_INCLUDE_ROLE_ASSIGNMENT"},
-			Usage:       `Whether to include role assignments assigned to the resources exported`,
-			Destination: &flagset.flagIncludeRoleAssignment,
+		&cli.StringSliceFlag{
+			Name:        "include-extension",
+			EnvVars:     []string{"AZTFEXPORT_INCLUDE_EXTENSION"},
+			Usage:       fmt.Sprintf(`Include extension resource types associated to the resources exported. Can be specified multiple times. Supported values: %v`, meta.SupportedExtensionResourceTypes),
+			Destination: &flagset.flagIncludeExtension,
 		},
 		&cli.BoolFlag{
 			Name:        "include-managed-resource",
@@ -611,7 +612,7 @@ func main() {
 						ResourceNamePattern:    flagset.flagPattern,
 						RecursiveQuery:         flagset.flagRecursive,
 						IncludeResourceGroup:   flagset.flagIncludeResourceGroup,
-						IncludeRoleAssignment:  flagset.flagIncludeRoleAssignment,
+						IncludeExtensions:      flagset.flagIncludeExtension.Value(),
 						IncludeManagedResource: flagset.flagIncludeManagedResource,
 					}
 
@@ -646,7 +647,7 @@ func main() {
 						ResourceGroupName:      rg,
 						ResourceNamePattern:    flagset.flagPattern,
 						RecursiveQuery:         true,
-						IncludeRoleAssignment:  flagset.flagIncludeRoleAssignment,
+						IncludeExtensions:      flagset.flagIncludeExtension.Value(),
 						IncludeManagedResource: flagset.flagIncludeManagedResource,
 					}
 
@@ -680,7 +681,7 @@ func main() {
 						ARGPredicate:                predicate,
 						ResourceNamePattern:         flagset.flagPattern,
 						RecursiveQuery:              flagset.flagRecursive,
-						IncludeRoleAssignment:       flagset.flagIncludeRoleAssignment,
+						IncludeExtensions:           flagset.flagIncludeExtension.Value(),
 						IncludeManagedResource:      flagset.flagIncludeManagedResource,
 						IncludeResourceGroup:        flagset.flagIncludeResourceGroup,
 						ARGTable:                    flagset.flagARGTable,

pkg/config/config.go

@@ -154,8 +154,9 @@ type Config struct {
 	// Each expanded value is sanitized to be a valid Terraform identifier.
 	ResourceNamePattern string
 
-	// IncludeRoleAssignment specifies whether to include the role assignments assigned to the exported resources
-	IncludeRoleAssignment bool
+	// IncludeExtensions specifies the set of extension resource types to include for the exported resources.
+	// Supported values are defined in the meta package (e.g. "role-assignment").
+	IncludeExtensions []string
 
 	// IncludeManagedResource specifies whether to allow service team/3rd party managed resources to be exported
 	IncludeManagedResource bool