From f67b637dd9107f7aee466417eb3dd84f11d029d9 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Tue, 3 Dec 2024 14:06:39 +0100
Subject: [PATCH] address comment

---
 .../pluginfw/tfschema/attribute_converter.go  |  6 ++++
 .../pluginfw/tfschema/customizable_schema.go  | 30 ++-----------------
 .../pluginfw/tfschema/list_nested_block.go    | 18 +++++++++++
 .../pluginfw/tfschema/single_nested_block.go  | 16 ++++++++++
 4 files changed, 42 insertions(+), 28 deletions(-)
 create mode 100644 internal/providers/pluginfw/tfschema/attribute_converter.go

diff --git a/internal/providers/pluginfw/tfschema/attribute_converter.go b/internal/providers/pluginfw/tfschema/attribute_converter.go
new file mode 100644
index 000000000..db2b1bda5
--- /dev/null
+++ b/internal/providers/pluginfw/tfschema/attribute_converter.go
@@ -0,0 +1,6 @@
+package tfschema
+
+type BlockToAttributeConverter interface {
+	// ConvertBlockToAttribute converts a contained block to its corresponding attribute type.
+	ConvertBlockToAttribute(string) BaseSchemaBuilder
+}
diff --git a/internal/providers/pluginfw/tfschema/customizable_schema.go b/internal/providers/pluginfw/tfschema/customizable_schema.go
index 136482673..b7051b450 100644
--- a/internal/providers/pluginfw/tfschema/customizable_schema.go
+++ b/internal/providers/pluginfw/tfschema/customizable_schema.go
@@ -211,34 +211,8 @@ func (s *CustomizableSchema) ConvertToAttribute(path ...string) *CustomizableSch
 
 	cb := func(attr BaseSchemaBuilder) BaseSchemaBuilder {
 		switch a := attr.(type) {
-		case ListNestedBlockBuilder:
-			elem, ok := a.NestedObject.Blocks[field]
-			if !ok {
-				panic(fmt.Errorf("field %s does not exist in nested block", field))
-			}
-			if a.NestedObject.Attributes == nil {
-				a.NestedObject.Attributes = make(map[string]AttributeBuilder)
-			}
-			a.NestedObject.Attributes[field] = elem.ToAttribute()
-			delete(a.NestedObject.Blocks, field)
-			if len(a.NestedObject.Blocks) == 0 {
-				a.NestedObject.Blocks = nil
-			}
-			return a
-		case SingleNestedBlockBuilder:
-			elem, ok := a.NestedObject.Blocks[field]
-			if !ok {
-				panic(fmt.Errorf("field %s does not exist in nested block", field))
-			}
-			if a.NestedObject.Attributes == nil {
-				a.NestedObject.Attributes = make(map[string]AttributeBuilder)
-			}
-			a.NestedObject.Attributes[field] = elem.ToAttribute()
-			delete(a.NestedObject.Blocks, field)
-			if len(a.NestedObject.Blocks) == 0 {
-				a.NestedObject.Blocks = nil
-			}
-			return a
+		case BlockToAttributeConverter:
+			return a.ConvertBlockToAttribute(field)
 		default:
 			panic(fmt.Errorf("ConvertToAttribute called on invalid attribute type: %s. %s", reflect.TypeOf(attr).String(), common.TerraformBugErrorMessage))
 		}
diff --git a/internal/providers/pluginfw/tfschema/list_nested_block.go b/internal/providers/pluginfw/tfschema/list_nested_block.go
index 03d9a5128..91c7af036 100644
--- a/internal/providers/pluginfw/tfschema/list_nested_block.go
+++ b/internal/providers/pluginfw/tfschema/list_nested_block.go
@@ -1,6 +1,8 @@
 package tfschema
 
 import (
+	"fmt"
+
 	dataschema "github.com/hashicorp/terraform-plugin-framework/datasource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/resource/schema"
 	"github.com/hashicorp/terraform-plugin-framework/resource/schema/planmodifier"
@@ -56,3 +58,19 @@ func (a ListNestedBlockBuilder) AddPlanModifier(v planmodifier.List) BaseSchemaB
 	a.PlanModifiers = append(a.PlanModifiers, v)
 	return a
 }
+
+func (a ListNestedBlockBuilder) ConvertBlockToAttribute(field string) BaseSchemaBuilder {
+	elem, ok := a.NestedObject.Blocks[field]
+	if !ok {
+		panic(fmt.Errorf("field %s does not exist in nested block", field))
+	}
+	if a.NestedObject.Attributes == nil {
+		a.NestedObject.Attributes = make(map[string]AttributeBuilder)
+	}
+	a.NestedObject.Attributes[field] = elem.ToAttribute()
+	delete(a.NestedObject.Blocks, field)
+	if len(a.NestedObject.Blocks) == 0 {
+		a.NestedObject.Blocks = nil
+	}
+	return a
+}
diff --git a/internal/providers/pluginfw/tfschema/single_nested_block.go b/internal/providers/pluginfw/tfschema/single_nested_block.go
index 1b88c9ae0..59508235b 100644
--- a/internal/providers/pluginfw/tfschema/single_nested_block.go
+++ b/internal/providers/pluginfw/tfschema/single_nested_block.go
@@ -68,3 +68,19 @@ func (a SingleNestedBlockBuilder) AddPlanModifier(v planmodifier.Object) BaseSch
 	a.PlanModifiers = append(a.PlanModifiers, v)
 	return a
 }
+
+func (a SingleNestedBlockBuilder) ConvertBlockToAttribute(field string) BaseSchemaBuilder {
+	elem, ok := a.NestedObject.Blocks[field]
+	if !ok {
+		panic(fmt.Errorf("field %s does not exist in nested block", field))
+	}
+	if a.NestedObject.Attributes == nil {
+		a.NestedObject.Attributes = make(map[string]AttributeBuilder)
+	}
+	a.NestedObject.Attributes[field] = elem.ToAttribute()
+	delete(a.NestedObject.Blocks, field)
+	if len(a.NestedObject.Blocks) == 0 {
+		a.NestedObject.Blocks = nil
+	}
+	return a
+}