Add query_plan attribute to visualization (#1752)

Fixes #1751.

Author: pietern

sql/acceptance/sql_dashboard_test.go

@@ -58,6 +58,24 @@ func TestAccDashboard(t *testing.T) {
 
 					options = jsonencode({})
 				}
+
+				resource "databricks_sql_visualization" "q1v2" {
+					query_id = databricks_sql_query.q1.id
+					type = "table"
+					name = "My Table (1)"
+
+					options = jsonencode({})
+
+					# Note: this resource differs from the one above in that
+					# the query plan is set. This tests that it can either
+					# be unset or set and in both cases yield a consistent result.
+					query_plan = jsonencode({
+						# The value should be non-empty to check we don't have
+						# plan changes due to whitespace changes in JSON serialization.
+						groups = [
+						]
+					})
+				}
 			`,
 		},
 	})

sql/api/visualization.go

@@ -13,4 +13,5 @@ type Visualization struct {
 	Name        string          `json:"name"`
 	Description string          `json:"description,omitempty"`
 	Options     json.RawMessage `json:"options,omitempty"`
+	QueryPlan   json.RawMessage `json:"query_plan,omitempty"`
 }

sql/api/visualization_test.go

@@ -15,6 +15,7 @@ func TestVisualizationMarshalUnmarshal(t *testing.T) {
 		Name:        "name",
 		Description: "description",
 		Options:     json.RawMessage("{}"),
+		QueryPlan:   json.RawMessage(`{"foo":"bar"}`),
 	}
 
 	out, err := json.Marshal(v)

sql/resource_sql_visualization.go

@@ -23,6 +23,7 @@ type VisualizationEntity struct {
 	Name        string `json:"name"`
 	Description string `json:"description,omitempty"`
 	Options     string `json:"options"`
+	QueryPlan   string `json:"query_plan,omitempty"`
 }
 
 func (v *VisualizationEntity) toAPIObject(schema map[string]*schema.Schema, data *schema.ResourceData) (*api.Visualization, error) {
@@ -37,6 +38,9 @@ func (v *VisualizationEntity) toAPIObject(schema map[string]*schema.Schema, data
 	av.Name = v.Name
 	av.Description = v.Description
 	av.Options = json.RawMessage(v.Options)
+	if v.QueryPlan != "" {
+		av.QueryPlan = json.RawMessage(v.QueryPlan)
+	}
 	return &av, nil
 }
 
@@ -49,6 +53,11 @@ func (v *VisualizationEntity) fromAPIObject(av *api.Visualization, schema map[st
 	v.Description = av.Description
 	v.Options = string(av.Options)
 
+	// If the query plan attribute is not set; it may come back as `null` from the API.
+	if av.QueryPlan != nil && !bytes.Equal(av.QueryPlan, []byte("null")) {
+		v.QueryPlan = string(av.QueryPlan)
+	}
+
 	// Transform to ResourceData.
 	return common.StructToData(*v, schema, data)
 }
@@ -125,26 +134,28 @@ func jsonRemarshal(in []byte) ([]byte, error) {
 	return out, nil
 }
 
+func suppressWhitespaceChangesInJSON(_, old, new string, d *schema.ResourceData) bool {
+	oldp, err := jsonRemarshal([]byte(old))
+	if err != nil {
+		log.Printf("[WARN] Unable to remarshal value %#v", old)
+		return false
+	}
+	newp, err := jsonRemarshal([]byte(new))
+	if err != nil {
+		log.Printf("[WARN] Unable to remarshal value %#v", new)
+		return false
+	}
+	return bytes.Equal(oldp, newp)
+}
+
 func ResourceSqlVisualization() *schema.Resource {
 	p := common.NewPairSeparatedID("query_id", "visualization_id", "/")
 	s := common.StructToSchema(
 		VisualizationEntity{},
 		func(m map[string]*schema.Schema) map[string]*schema.Schema {
-			// We care only about logical changes to the JSON payload in `options`.
-			m["options"].DiffSuppressFunc = func(_, old, new string, d *schema.ResourceData) bool {
-				oldp, err := jsonRemarshal([]byte(old))
-				if err != nil {
-					log.Printf("[WARN] Unable to remarshal value %#v", old)
-					return false
-				}
-				newp, err := jsonRemarshal([]byte(new))
-				if err != nil {
-					log.Printf("[WARN] Unable to remarshal value %#v", new)
-					return false
-				}
-				return bytes.Equal(oldp, newp)
-			}
-
+			// We care only about logical changes to the JSON payload in `options` and `query_plan`.
+			m["options"].DiffSuppressFunc = suppressWhitespaceChangesInJSON
+			m["query_plan"].DiffSuppressFunc = suppressWhitespaceChangesInJSON
 			return m
 		})
 

sql/resource_sql_visualization_test.go

@@ -137,6 +137,106 @@ func TestVisualizationRead(t *testing.T) {
 	assert.Less(t, 0, len(d.Get("options").(string)))
 }
 
+func TestVisualizationReadWithQueryPlan(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:   "GET",
+				Resource: "/api/2.0/preview/sql/queries/foo",
+				Response: api.Query{
+					ID: "foo",
+					Visualizations: []json.RawMessage{
+						json.RawMessage(`
+							{
+								"id": 12344
+							}
+						`),
+						json.RawMessage(`
+							{
+								"id": 12345,
+								"type": "CHART",
+								"name": "My Chart",
+								"description": "Some Description",
+								"options": {},
+								"query_plan": {"foo":"qux"}
+							}
+						`),
+						json.RawMessage(`
+							{
+								"id": 12345
+							}
+						`),
+					},
+				},
+			},
+		},
+		Resource: ResourceSqlVisualization(),
+		Read:     true,
+		ID:       "foo/12345",
+		State: map[string]any{
+			"query_id":    "foo",
+			"type":        "chart",
+			"name":        "My Chart",
+			"description": "Some Description",
+			"options":     "{}",
+			"query_plan":  `{"foo":"bar"}`,
+		},
+	}.Apply(t)
+
+	assert.NoError(t, err)
+	assert.Equal(t, `{"foo":"qux"}`, d.Get("query_plan").(string))
+}
+
+func TestVisualizationReadWithNullQueryPlan(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:   "GET",
+				Resource: "/api/2.0/preview/sql/queries/foo",
+				Response: api.Query{
+					ID: "foo",
+					Visualizations: []json.RawMessage{
+						json.RawMessage(`
+							{
+								"id": 12344
+							}
+						`),
+						json.RawMessage(`
+							{
+								"id": 12345,
+								"type": "CHART",
+								"name": "My Chart",
+								"description": "Some Description",
+								"options": {},
+								"query_plan": null
+							}
+						`),
+						json.RawMessage(`
+							{
+								"id": 12345
+							}
+						`),
+					},
+				},
+			},
+		},
+		Resource: ResourceSqlVisualization(),
+		Read:     true,
+		New:      true,
+		ID:       "foo/12345",
+		State: map[string]any{
+			"query_id":    "foo",
+			"type":        "chart",
+			"name":        "My Chart",
+			"description": "Some Description",
+			"options":     "{}",
+		},
+	}.Apply(t)
+
+	assert.NoError(t, err)
+	assert.Equal(t, "", d.Get("query_plan").(string))
+}
+
 func TestVisualizationReadNotFound(t *testing.T) {
 	d, err := qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{