Jobs GoSDK schema migration (#3532)

* update

* update

* address comments

Author: edwardfeng-db

clusters/resource_cluster.go

@@ -163,6 +163,34 @@ type ClusterSpec struct {
 	Libraries []compute.Library `json:"libraries,omitempty" tf:"slice_set,alias:library"`
 }
 
+func (ClusterSpec) CustomizeSchemaResourceSpecific(s *common.CustomizableSchema) *common.CustomizableSchema {
+	s.SchemaPath("autotermination_minutes").SetDefault(60)
+	s.AddNewField("default_tags", &schema.Schema{
+		Type:     schema.TypeMap,
+		Computed: true,
+	})
+	s.AddNewField("is_pinned", &schema.Schema{
+		Type:     schema.TypeBool,
+		Optional: true,
+		Default:  false,
+		DiffSuppressFunc: func(k, old, new string, d *schema.ResourceData) bool {
+			if old == "" && new == "false" {
+				return true
+			}
+			return old == new
+		},
+	})
+	s.AddNewField("state", &schema.Schema{
+		Type:     schema.TypeString,
+		Computed: true,
+	})
+	s.AddNewField("url", &schema.Schema{
+		Type:     schema.TypeString,
+		Computed: true,
+	})
+	return s
+}
+
 func (ClusterSpec) CustomizeSchema(s *common.CustomizableSchema) *common.CustomizableSchema {
 	s.SchemaPath("cluster_source").SetReadOnly()
 	s.SchemaPath("enable_elastic_disk").SetComputed()
@@ -198,7 +226,6 @@ func (ClusterSpec) CustomizeSchema(s *common.CustomizableSchema) *common.Customi
 	s.SchemaPath("aws_attributes", "zone_id").SetCustomSuppressDiff(ZoneDiffSuppress)
 	s.SchemaPath("azure_attributes").SetSuppressDiff().SetConflictsWith([]string{"aws_attributes", "gcp_attributes"})
 	s.SchemaPath("gcp_attributes").SetSuppressDiff().SetConflictsWith([]string{"aws_attributes", "azure_attributes"})
-	s.SchemaPath("autotermination_minutes").SetDefault(60)
 	s.SchemaPath("autoscale", "max_workers").SetOptional()
 	s.SchemaPath("autoscale", "min_workers").SetOptional()
 	s.SchemaPath("cluster_log_conf", "dbfs", "destination").SetRequired()
@@ -208,31 +235,8 @@ func (ClusterSpec) CustomizeSchema(s *common.CustomizableSchema) *common.Customi
 		Type:     schema.TypeString,
 		Computed: true,
 	})
-	s.AddNewField("default_tags", &schema.Schema{
-		Type:     schema.TypeMap,
-		Computed: true,
-	})
-	s.AddNewField("state", &schema.Schema{
-		Type:     schema.TypeString,
-		Computed: true,
-	})
 	s.SchemaPath("instance_pool_id").SetConflictsWith([]string{"driver_node_type_id", "node_type_id"})
 	s.SchemaPath("runtime_engine").SetValidateFunc(validation.StringInSlice([]string{"PHOTON", "STANDARD"}, false))
-	s.AddNewField("is_pinned", &schema.Schema{
-		Type:     schema.TypeBool,
-		Optional: true,
-		Default:  false,
-		DiffSuppressFunc: func(k, old, new string, d *schema.ResourceData) bool {
-			if old == "" && new == "false" {
-				return true
-			}
-			return old == new
-		},
-	})
-	s.AddNewField("url", &schema.Schema{
-		Type:     schema.TypeString,
-		Computed: true,
-	})
 	s.SchemaPath("num_workers").SetDefault(0).SetValidateDiagFunc(validation.ToDiagFunc(validation.IntAtLeast(0)))
 	s.AddNewField("cluster_mount_info", &schema.Schema{
 		Type:     schema.TypeList,

clusters/resource_library.go

@@ -16,14 +16,18 @@ type LibraryResource struct {
 	compute.Library
 }
 
-func (LibraryResource) CustomizeSchema(s *common.CustomizableSchema) *common.CustomizableSchema {
+func (LibraryResource) CustomizeSchemaResourceSpecific(s *common.CustomizableSchema) *common.CustomizableSchema {
 	s.AddNewField("cluster_id", &schema.Schema{
 		Type:     schema.TypeString,
 		Required: true,
 	})
 	return s
 }
 
+func (LibraryResource) CustomizeSchema(s *common.CustomizableSchema) *common.CustomizableSchema {
+	return s
+}
+
 func ResourceLibrary() common.Resource {
 	libraySdkSchema := common.StructToSchema(LibraryResource{}, nil)
 	parseId := func(id string) (string, string) {

common/customizable_schema.go

@@ -253,17 +253,24 @@ func (s *CustomizableSchema) SetValidateDiagFunc(validate func(interface{}, cty.
 }
 
 func (s *CustomizableSchema) AddNewField(key string, newField *schema.Schema) *CustomizableSchema {
-	cv, ok := s.Schema.Elem.(*schema.Resource)
-	if !ok {
-		panic("Cannot add new field, target is not nested resource")
-	}
-	_, exists := cv.Schema[key]
+	scm := s.GetSchemaMap()
+	_, exists := scm[key]
 	if exists {
 		panic("Cannot add new field, " + key + " already exists in the schema")
 	}
-	cv.Schema[key] = newField
+	scm[key] = newField
 	if s.isSuppressDiff {
 		newField.DiffSuppressFunc = diffSuppressor(key, newField)
 	}
 	return s
 }
+
+func (s *CustomizableSchema) RemoveField(key string) *CustomizableSchema {
+	scm := s.GetSchemaMap()
+	_, exists := scm[key]
+	if !exists {
+		panic("Cannot remove new field, " + key + " does not exist in the schema")
+	}
+	delete(scm, key)
+	return s
+}

common/reflect_resource.go

@@ -42,7 +42,7 @@ var kindMap = map[reflect.Kind]string{
 // Global registry for ResoureProvider, the goal is to make StructToSchema able to find pre-registered ResourceProvider for a
 // given struct so that we don't have to specify aliases and customizations redundantly if one schema references the another
 // schema.
-var resourceProviderRegistry map[string]ResourceProvider
+var resourceProviderRegistry map[reflect.Type]ResourceProvider
 
 // Pre-registered ResourceProvider for a given struct into resourceProviderRegistry.
 // This function should be called in the init() function in packages with ResourceProvider.
@@ -53,14 +53,27 @@ var resourceProviderRegistry map[string]ResourceProvider
 //	}
 func RegisterResourceProvider(v any, r ResourceProvider) {
 	if resourceProviderRegistry == nil {
-		resourceProviderRegistry = make(map[string]ResourceProvider)
+		resourceProviderRegistry = make(map[reflect.Type]ResourceProvider)
 	}
-	typeName := getNameForType(reflect.ValueOf(v).Type())
-	if _, ok := resourceProviderRegistry[typeName]; ok {
-		errMsg := fmt.Sprintf("%s has already been registered, please avoid registering the same type repeatedly.", typeName)
+	t := getResourceProviderRegistryKey(v)
+	if _, ok := resourceProviderRegistry[t]; ok {
+		errMsg := fmt.Sprintf("%s has already been registered, please avoid registering the same type repeatedly.", getNameForType(t))
 		panic(errMsg)
 	}
-	resourceProviderRegistry[typeName] = r
+	resourceProviderRegistry[t] = r
+}
+
+func getResourceProviderRegistryKey(v any) reflect.Type {
+	t := reflect.ValueOf(v).Type()
+	return getNonPointerType(t)
+}
+
+func getNonPointerType(t reflect.Type) reflect.Type {
+	if t.Kind() == reflect.Ptr {
+		return t.Elem()
+	} else {
+		return t
+	}
 }
 
 // Generic interface for ResourceProvider. Using CustomizeSchema function to keep track of additional information
@@ -69,6 +82,19 @@ type ResourceProvider interface {
 	CustomizeSchema(*CustomizableSchema) *CustomizableSchema
 }
 
+// Interface for ResourcePrivider that needs certain customizaitons that shouldn't be shared.
+type ResourceProviderWithResourceSpecificCustomization interface {
+	ResourceProvider
+	// Customizations in this function will not be applied when this type is referenced from another resource by resourceProviderRegistry.
+	// Example use case:
+	//    // autotermination_minutes should only have a default value of 60 for the clusters resource, but not in the nested cluster schema in jobs
+	//    func (ClusterSpec) CustomizeSchemaResourceSpecific(s *common.CustomizableSchema) *common.CustomizableSchema {
+	// 	      s.SchemaPath("autotermination_minutes").SetDefault(60)
+	// 	      return s
+	//    }
+	CustomizeSchemaResourceSpecific(*CustomizableSchema) *CustomizableSchema
+}
+
 // Interface for ResourceProvider instances that need aliases for fields.
 type ResourceProviderWithAlias interface {
 	ResourceProvider
@@ -100,8 +126,19 @@ type RecursiveResourceProvider interface {
 	MaxDepthForTypes() map[string]int
 }
 
-// Takes in a ResourceProvider and converts that into a map from string to schema.
+// Used by StructToSchema, after executing typeToSchema and applying the standard customizations,
+// it also applies the resource specific customizations, if applicable.
 func resourceProviderStructToSchema(v ResourceProvider, scp schemaPathContext) map[string]*schema.Schema {
+	customizedScm := resourceProviderStructToSchemaInternal(v, scp)
+	if rps, ok := v.(ResourceProviderWithResourceSpecificCustomization); ok {
+		// Apply resource specific customizations, if applicable.
+		return rps.CustomizeSchemaResourceSpecific(customizedScm).GetSchemaMap()
+	} else {
+		return customizedScm.GetSchemaMap()
+	}
+}
+
+func resourceProviderStructToSchemaInternal(v ResourceProvider, scp schemaPathContext) *CustomizableSchema {
 	rv := reflect.ValueOf(v)
 	var scm map[string]*schema.Schema
 	aliases := map[string]map[string]string{}
@@ -115,8 +152,7 @@ func resourceProviderStructToSchema(v ResourceProvider, scp schemaPathContext) m
 	}
 	cs := CustomizeSchemaPath(scm)
 	cs.context = scp
-	scm = v.CustomizeSchema(cs).GetSchemaMap()
-	return scm
+	return v.CustomizeSchema(cs)
 }
 
 func reflectKind(k reflect.Kind) string {
@@ -372,6 +408,9 @@ func listAllFields(v reflect.Value) []field {
 }
 
 func typeToSchema(v reflect.Value, aliases map[string]map[string]string, tc trackingContext) map[string]*schema.Schema {
+	if rpStruct, ok := resourceProviderRegistry[getNonPointerType(v.Type())]; ok {
+		return resourceProviderStructToSchemaInternal(rpStruct, tc.pathCtx).GetSchemaMap()
+	}
 	scm := map[string]*schema.Schema{}
 	rk := v.Kind()
 	if rk == reflect.Ptr {