Skip to content

Commit a162a2f

Browse files
andrewnesterandrewnester
authored
Added state upgrader to remove max_clusters_per_user and ebs volume attributes set to zero (#3551)
* Added state upgrader to remove max_clusters_per_user set to zero * fix for #3553 * fix typo
1 parent 6d08102 commit a162a2f

File tree

4 files changed

+129
-26
lines changed

4 files changed

+129
-26
lines changed

‎clusters/resource_cluster.go‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"strings"
88
"time"
99

10+
"github.com/hashicorp/go-cty/cty"
1011
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
1112
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
1213

@@ -19,7 +20,7 @@ const DefaultProvisionTimeout = 30 * time.Minute
1920
const DbfsDeprecationWarning = "For init scripts use 'volumes', 'workspace' or cloud storage location instead of 'dbfs'."
2021

2122
var clusterSchema = resourceClusterSchema()
22-
var clusterSchemaVersion = 2
23+
var clusterSchemaVersion = 3
2324

2425
func ResourceCluster() common.Resource {
2526
return common.Resource{
@@ -30,7 +31,46 @@ func ResourceCluster() common.Resource {
3031
Schema: clusterSchema,
3132
SchemaVersion: clusterSchemaVersion,
3233
Timeouts: resourceClusterTimeouts(),
34+
StateUpgraders: []schema.StateUpgrader{
35+
{
36+
Type: clusterSchemaV0(),
37+
Version: 2,
38+
Upgrade: removeZeroAwsEbsVolumeAttributes,
39+
},
40+
},
41+
}
42+
}
43+
44+
func clusterSchemaV0() cty.Type {
45+
return (&schema.Resource{
46+
Schema: clusterSchema}).CoreConfigSchema().ImpliedType()
47+
}
48+
49+
func removeZeroAwsEbsVolumeAttributes(ctx context.Context, rawState map[string]interface{}, meta interface{}) (map[string]interface{}, error) {
50+
newState := map[string]any{}
51+
for k, v := range rawState {
52+
switch k {
53+
case "aws_attributes":
54+
awsAttributes, ok := v.(map[string]any)
55+
if !ok {
56+
continue
57+
}
58+
59+
if awsAttributes["ebs_volume_count"] == 0 {
60+
log.Printf("[INFO] remove zero ebs_volume_count")
61+
delete(awsAttributes, "ebs_volume_count")
62+
}
63+
if awsAttributes["ebs_volume_size"] == 0 {
64+
log.Printf("[INFO] remove zero ebs_volume_size")
65+
delete(awsAttributes, "ebs_volume_size")
66+
}
67+
68+
newState[k] = awsAttributes
69+
default:
70+
newState[k] = v
71+
}
3372
}
73+
return newState, nil
3474
}
3575

3676
func resourceClusterTimeouts() *schema.ResourceTimeout {

‎policies/data_cluster_policy.go‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,16 @@ package policies
22

33
import (
44
"context"
5+
"log"
56

67
"github.com/databricks/databricks-sdk-go"
78
"github.com/databricks/terraform-provider-databricks/common"
9+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
810
)
911

1012
// DataSourceClusterPolicy returns information about cluster policy specified by name
1113
func DataSourceClusterPolicy() common.Resource {
12-
return common.WorkspaceData(func(ctx context.Context, data *struct {
14+
resource := common.WorkspaceData(func(ctx context.Context, data *struct {
1315
Id string `json:"id,omitempty" tf:"computed"`
1416
Name string `json:"name,omitempty" tf:"computed"`
1517
Definition string `json:"definition,omitempty" tf:"computed"`
@@ -32,4 +34,31 @@ func DataSourceClusterPolicy() common.Resource {
3234
data.MaxClustersPerUser = int(policy.MaxClustersPerUser)
3335
return nil
3436
})
37+
resource.SchemaVersion = 1
38+
resource.StateUpgraders = []schema.StateUpgrader{
39+
{
40+
Type: resource.ToResource().CoreConfigSchema().ImpliedType(),
41+
Version: 0,
42+
Upgrade: removeZeroMaxClustersPerUser,
43+
},
44+
}
45+
return resource
46+
}
47+
48+
func removeZeroMaxClustersPerUser(ctx context.Context, rawState map[string]interface{}, meta interface{}) (map[string]interface{}, error) {
49+
newState := map[string]any{}
50+
for k, v := range rawState {
51+
switch k {
52+
case "max_clusters_per_user":
53+
vv, ok := v.(int)
54+
if !ok || vv == 0 {
55+
log.Printf("[INFO] remove zero max_clusters_per_user")
56+
continue
57+
}
58+
newState[k] = v
59+
default:
60+
newState[k] = v
61+
}
62+
}
63+
return newState, nil
3564
}

‎policies/data_cluster_policy_test.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
package policies
22

33
import (
4+
"context"
45
"testing"
56

67
"github.com/databricks/databricks-sdk-go/apierr"
78
"github.com/databricks/databricks-sdk-go/service/compute"
89
"github.com/databricks/terraform-provider-databricks/qa"
10+
"github.com/stretchr/testify/assert"
911
)
1012

1113
func TestDataSourceClusterPolicy(t *testing.T) {
@@ -82,3 +84,21 @@ func TestDataSourceClusterPolicyNotFound(t *testing.T) {
8284
HCL: `name = "policy"`,
8385
}.ExpectError(t, "Policy named 'policy' does not exist")
8486
}
87+
88+
func TestDataSourceClusterPolicyStateUpgrader(t *testing.T) {
89+
state, err := removeZeroMaxClustersPerUser(context.Background(),
90+
map[string]any{
91+
"max_clusters_per_user": 0,
92+
}, nil)
93+
assert.NoError(t, err)
94+
_, ok := state["max_clusters_per_user"]
95+
assert.False(t, ok)
96+
97+
state, err = removeZeroMaxClustersPerUser(context.Background(),
98+
map[string]any{
99+
"max_clusters_per_user": 1,
100+
}, nil)
101+
assert.NoError(t, err)
102+
_, ok = state["max_clusters_per_user"]
103+
assert.True(t, ok)
104+
}

‎policies/resource_cluster_policy.go‎

Lines changed: 38 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"github.com/databricks/databricks-sdk-go"
77
"github.com/databricks/databricks-sdk-go/service/compute"
88
"github.com/databricks/terraform-provider-databricks/common"
9+
"github.com/hashicorp/go-cty/cty"
910

1011
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
1112
)
@@ -24,37 +25,45 @@ func isBuiltinPolicyFamily(ctx context.Context, w *databricks.WorkspaceClient, f
2425
return false, nil
2526
}
2627

27-
// ResourceClusterPolicy ...
28-
func ResourceClusterPolicy() common.Resource {
29-
s := common.StructToSchema(
30-
compute.CreatePolicy{},
31-
func(m map[string]*schema.Schema) map[string]*schema.Schema {
32-
m["policy_id"] = &schema.Schema{
33-
Type: schema.TypeString,
34-
Computed: true,
35-
}
36-
m["definition"].ConflictsWith = []string{"policy_family_definition_overrides", "policy_family_id"}
37-
m["definition"].Computed = true
38-
m["definition"].DiffSuppressFunc = common.SuppressDiffWhitespaceChange
28+
var rcpSchema = common.StructToSchema(
29+
compute.CreatePolicy{},
30+
func(m map[string]*schema.Schema) map[string]*schema.Schema {
31+
m["policy_id"] = &schema.Schema{
32+
Type: schema.TypeString,
33+
Computed: true,
34+
}
35+
m["definition"].ConflictsWith = []string{"policy_family_definition_overrides", "policy_family_id"}
36+
m["definition"].Computed = true
37+
m["definition"].DiffSuppressFunc = common.SuppressDiffWhitespaceChange
3938

40-
m["policy_family_definition_overrides"].ConflictsWith = []string{"definition"}
41-
m["policy_family_definition_overrides"].DiffSuppressFunc = common.SuppressDiffWhitespaceChange
42-
m["policy_family_id"].ConflictsWith = []string{"definition"}
43-
m["policy_family_definition_overrides"].RequiredWith = []string{"policy_family_id"}
39+
m["policy_family_definition_overrides"].ConflictsWith = []string{"definition"}
40+
m["policy_family_definition_overrides"].DiffSuppressFunc = common.SuppressDiffWhitespaceChange
41+
m["policy_family_id"].ConflictsWith = []string{"definition"}
42+
m["policy_family_definition_overrides"].RequiredWith = []string{"policy_family_id"}
4443

45-
return m
46-
})
44+
return m
45+
})
4746

47+
// ResourceClusterPolicy ...
48+
func ResourceClusterPolicy() common.Resource {
4849
return common.Resource{
49-
Schema: s,
50+
Schema: rcpSchema,
51+
SchemaVersion: 1,
52+
StateUpgraders: []schema.StateUpgrader{
53+
{
54+
Type: rcpSchemaV0(),
55+
Version: 0,
56+
Upgrade: removeZeroMaxClustersPerUser,
57+
},
58+
},
5059
Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
5160
w, err := c.WorkspaceClient()
5261
if err != nil {
5362
return err
5463
}
5564

5665
var request compute.CreatePolicy
57-
common.DataToStructPointer(d, s, &request)
66+
common.DataToStructPointer(d, rcpSchema, &request)
5867

5968
var clusterPolicy *compute.CreatePolicyResponse
6069
if request.PolicyFamilyId != "" {
@@ -69,7 +78,7 @@ func ResourceClusterPolicy() common.Resource {
6978
}
7079
clusterPolicy = &compute.CreatePolicyResponse{PolicyId: resp.PolicyId}
7180
var editRequest compute.EditPolicy
72-
common.DataToStructPointer(d, s, &editRequest)
81+
common.DataToStructPointer(d, rcpSchema, &editRequest)
7382
editRequest.PolicyId = resp.PolicyId
7483
err = w.ClusterPolicies.Edit(ctx, editRequest)
7584
} else {
@@ -93,7 +102,7 @@ func ResourceClusterPolicy() common.Resource {
93102
if err != nil {
94103
return err
95104
}
96-
return common.StructToData(resp, s, d)
105+
return common.StructToData(resp, rcpSchema, d)
97106
},
98107
Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
99108
w, err := c.WorkspaceClient()
@@ -102,7 +111,7 @@ func ResourceClusterPolicy() common.Resource {
102111
}
103112

104113
var request compute.EditPolicy
105-
common.DataToStructPointer(d, s, &request)
114+
common.DataToStructPointer(d, rcpSchema, &request)
106115
request.PolicyId = d.Id()
107116
if request.PolicyFamilyId != "" {
108117
request.Definition = ""
@@ -116,7 +125,7 @@ func ResourceClusterPolicy() common.Resource {
116125
return err
117126
}
118127
var request compute.EditPolicy
119-
common.DataToStructPointer(d, s, &request)
128+
common.DataToStructPointer(d, rcpSchema, &request)
120129
if request.PolicyFamilyId != "" {
121130
isBuiltin, err := isBuiltinPolicyFamily(ctx, w, request.PolicyFamilyId, request.Name)
122131
if err != nil {
@@ -133,3 +142,8 @@ func ResourceClusterPolicy() common.Resource {
133142
},
134143
}
135144
}
145+
146+
func rcpSchemaV0() cty.Type {
147+
return (&schema.Resource{
148+
Schema: rcpSchema}).CoreConfigSchema().ImpliedType()
149+
}

0 commit comments

Comments
 (0)