Allow specifying visualization by its Terraform resource ID

Author: pietern

sqlanalytics/api/widget.go

@@ -52,3 +52,26 @@ type WidgetParameterMapping struct {
 	// This title overrides the title given to this parameter by the query, if specified.
 	Title string `json:"title,omitempty"`
 }
+
+// UnmarshalJSON ...
+func (w *Widget) UnmarshalJSON(b []byte) error {
+	type localWidget Widget
+	err := json.Unmarshal(b, (*localWidget)(w))
+	if err != nil {
+		return err
+	}
+
+	// If a visualization is configured, set the corresponding visualization ID.
+	// The visualization ID is only used when creating or updating a widget
+	// and not set when retrieving an existing widget with visualization.
+	if w.Visualization != nil {
+		var v Visualization
+		err = json.Unmarshal(w.Visualization, &v)
+		if err != nil {
+			return err
+		}
+		w.VisualizationID = &v.ID
+	}
+
+	return nil
+}

sqlanalytics/api/widget_test.go

@@ -49,3 +49,27 @@ func TestWidgetMarshalUnmarshal(t *testing.T) {
 
 	assert.Equal(t, w, wp)
 }
+
+func TestWidgetUnmarshalWithVisualization(t *testing.T) {
+	w := Widget{
+		Visualization: json.RawMessage(`
+			{
+				"id": 12345
+			}
+		`),
+	}
+
+	out, err := json.Marshal(w)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var wp Widget
+	if err := json.Unmarshal(out, &wp); err != nil {
+		t.Fatal(err)
+	}
+
+	if assert.NotNil(t, wp.VisualizationID) {
+		assert.Equal(t, 12345, *wp.VisualizationID)
+	}
+}

sqlanalytics/resource_widget.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"sort"
 	"strconv"
+	"strings"
 
 	"github.com/databrickslabs/terraform-provider-databricks/common"
 	"github.com/databrickslabs/terraform-provider-databricks/sqlanalytics/api"
@@ -58,6 +59,14 @@ func (a sortWidgetParameter) Less(i, j int) bool {
 	return a[i].Name < a[j].Name
 }
 
+// Use second part of ID if it's a composite ID, or the verbatim value.
+// This allows setting the visualization ID as the backend visualization ID
+// or as the visualization's resource composite ID. Both will work.
+func extractVisualizationID(id string) string {
+	parts := strings.SplitN(id, "/", 2)
+	return parts[len(parts)-1]
+}
+
 func (w *WidgetEntity) toAPIObject(schema map[string]*schema.Schema, data *schema.ResourceData) (*api.Widget, error) {
 	var aw api.Widget
 
@@ -70,7 +79,7 @@ func (w *WidgetEntity) toAPIObject(schema map[string]*schema.Schema, data *schem
 
 	// The visualization ID is a string for the Terraform resource and an integer in the API.
 	if w.VisualizationID != "" {
-		visualizationID, err := strconv.Atoi(w.VisualizationID)
+		visualizationID, err := strconv.Atoi(extractVisualizationID(w.VisualizationID))
 		if err != nil {
 			return nil, err
 		}
@@ -244,6 +253,14 @@ func ResourceWidget() *schema.Resource {
 		WidgetEntity{},
 		func(m map[string]*schema.Schema) map[string]*schema.Schema {
 			m["text"].ConflictsWith = []string{"visualization_id"}
+
+			// Ignore the query ID part in composite visualization ID.
+			// It is present in this field if users refer to a visualization by the native
+			// Terraform resource ID (e.g. `databricks_sql_visualization.name.id`)
+			m["visualization_id"].DiffSuppressFunc = func(_, old, new string, d *schema.ResourceData) bool {
+				return extractVisualizationID(old) == extractVisualizationID(new)
+			}
+
 			return m
 		})
 

sqlanalytics/resource_widget_test.go

@@ -69,6 +69,66 @@ func TestWidgetCreateWithVisualization(t *testing.T) {
 	assert.Equal(t, "678", d.Get("visualization_id"))
 }
 
+func TestWidgetCreateWithVisualizationByResourceID(t *testing.T) {
+	i678 := 678
+
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:   "POST",
+				Resource: "/api/2.0/preview/sql/widgets",
+				ExpectedRequest: api.Widget{
+					DashboardID:     "some-uuid",
+					VisualizationID: &i678,
+				},
+				Response: api.Widget{
+					ID:          12345,
+					DashboardID: "some-uuid",
+				},
+			},
+			{
+				Method:   "GET",
+				Resource: "/api/2.0/preview/sql/dashboards/some-uuid",
+				Response: api.Dashboard{
+					ID: "some-uuid",
+					Widgets: []json.RawMessage{
+						json.RawMessage(`
+							{
+								"id": 12344
+							}
+						`),
+						json.RawMessage(`
+							{
+								"id": 12345,
+								"visualization": {
+									"id": 678
+								}
+							}
+						`),
+						json.RawMessage(`
+							{
+								"id": 12346
+							}
+						`),
+					},
+				},
+			},
+		},
+		Resource: ResourceWidget(),
+		Create:   true,
+		HCL: `
+			dashboard_id     = "some-uuid"
+			visualization_id = "doesnt-matter/678"
+		`,
+	}.Apply(t)
+
+	assert.NoError(t, err, err)
+	assert.Equal(t, "some-uuid/12345", d.Id(), "Resource ID should not be empty")
+	assert.Equal(t, "some-uuid", d.Get("dashboard_id"))
+	assert.Equal(t, "12345", d.Get("widget_id"))
+	assert.Equal(t, "678", d.Get("visualization_id"))
+}
+
 func TestWidgetCreateWithText(t *testing.T) {
 	sText := "widget text"