From b35195cd3492c11957bb9fe9004c235068c6fe23 Mon Sep 17 00:00:00 2001
From: Tanmay Rustagi <tanmay.rustagi@databricks.com>
Date: Mon, 29 Jul 2024 02:29:25 +0200
Subject: [PATCH] -

---
 common/reflect_resource_plugin_framework.go     | 13 +++++++------
 pluginframework/resource_quality_monitor.go     | 17 +++++++++++++++--
 .../resource/quality-monitor/quality_monitor.tf |  7 ++++---
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/common/reflect_resource_plugin_framework.go b/common/reflect_resource_plugin_framework.go
index 58917c2e78..ea0a2f3ae5 100644
--- a/common/reflect_resource_plugin_framework.go
+++ b/common/reflect_resource_plugin_framework.go
@@ -403,16 +403,17 @@ func fieldIsOptional(field reflect.StructField) bool {
 }
 
 func PluginFrameworkResourceStructToSchema(v any, customizeSchema func(CustomizableSchemaPluginFramework) CustomizableSchemaPluginFramework) schema.Schema {
+	attributes := PluginFrameworkResourceStructToSchemaMap(v, customizeSchema)
+	return schema.Schema{Attributes: attributes}
+}
+
+func PluginFrameworkResourceStructToSchemaMap(v any, customizeSchema func(CustomizableSchemaPluginFramework) CustomizableSchemaPluginFramework) map[string]schema.Attribute {
 	attributes := pluginFrameworkResourceTypeToSchema(reflect.ValueOf(v))
 
 	if customizeSchema != nil {
 		cs := customizeSchema(*ConstructCustomizableSchema(attributes))
-		return schema.Schema{Attributes: cs.ToAttributeMap()}
+		return cs.ToAttributeMap()
 	} else {
-		return schema.Schema{Attributes: attributes}
+		return attributes
 	}
 }
-
-func PluginFrameworkResourceStructToSchemaMap(v any) map[string]schema.Attribute {
-	return pluginFrameworkResourceTypeToSchema(reflect.ValueOf(v))
-}
diff --git a/pluginframework/resource_quality_monitor.go b/pluginframework/resource_quality_monitor.go
index 01f021faca..5dbc1bf1c2 100644
--- a/pluginframework/resource_quality_monitor.go
+++ b/pluginframework/resource_quality_monitor.go
@@ -52,8 +52,21 @@ func (r *QualityMonitorResource) Metadata(ctx context.Context, req resource.Meta
 func (r *QualityMonitorResource) Schema(ctx context.Context, req resource.SchemaRequest, resp *resource.SchemaResponse) {
 	resp.Schema = schema.Schema{
 		Description: "Terraform schema for Databricks Lakehouse Monitor. MonitorInfo struct is used to create the schema",
-		// TODO: Add CustomizeSchemaPath once it is supported in the plugin framework
-		Attributes: common.PluginFrameworkResourceStructToSchemaMap(catalog_tf.MonitorInfo{}),
+		Attributes: common.PluginFrameworkResourceStructToSchemaMap(catalog_tf.MonitorInfo{}, func(c common.CustomizableSchemaPluginFramework) common.CustomizableSchemaPluginFramework {
+			c.AddNewField("skip_builtin_dashboard", schema.BoolAttribute{Optional: true, Required: false})
+			c.AddNewField("warehouse_id", schema.StringAttribute{Optional: true, Required: false})
+			c.SetRequired("assets_dir")
+			c.SetRequired("output_schema_name")
+			c.SetRequired("table_name")
+			// TODO: Uncomment this once SetReadOnly is supported in the plugin framework
+			// c.SetReadOnly("monitor_version")
+			// c.SetReadOnly("drift_metrics_table_name")
+			// c.SetReadOnly("profile_metrics_table_name")
+			// c.SetReadOnly("status")
+			// c.SetReadOnly("dashboard_id")
+			// c.SetReadOnly("schedule", "pause_status")
+			return c
+		}),
 	}
 }
 
diff --git a/pluginframework/test/resource/quality-monitor/quality_monitor.tf b/pluginframework/test/resource/quality-monitor/quality_monitor.tf
index 665a982d20..e98b118ebb 100644
--- a/pluginframework/test/resource/quality-monitor/quality_monitor.tf
+++ b/pluginframework/test/resource/quality-monitor/quality_monitor.tf
@@ -40,7 +40,8 @@ resource "databricks_sql_table" "testtable" {
   }
 }
 
-resource "databricks_lakehouse_monitor_pluginframework" "testmonitor" { 
-  table_name = "${databricks_catalog.testcatalog.name}.${databricks_schema.testschema.name}.${databricks_sql_table.testtable.name}"
-}
+# TODO: Uncomment once SetReadOnly is supported in the plugin framework and add the necessary fields
+# resource "databricks_lakehouse_monitor_pluginframework" "testmonitor" { 
+#   table_name = "${databricks_catalog.testcatalog.name}.${databricks_schema.testschema.name}.${databricks_sql_table.testtable.name}"
+# }