Allow rotating token block in databricks_mws_workspaces resource by only changing comment field (#2114)

shreyas-goenka · web-flow · commit 424c60941c22 · 2023-04-03T14:54:57.000Z
Tested manually for the following cases

Without this PR the provider recreates the entire workspace on a token update
With changes in this PR only the token is refreshed
When both token and storage_configuration_id are changed then the entire workspace is recreated
Additional unit tests also added that allow checks that patch workspace calls are not made when only token is changed

Also added an integration test to check tokens are successfully updated
diff --git a/internal/acceptance/init_test.go b/internal/acceptance/init_test.go
@@ -272,6 +272,24 @@ func resourceCheck(name string,
 	}
 }
 
+// resourceCheckWithState calls back a function with client and resource instance state
+func resourceCheckWithState(name string,
+	cb func(ctx context.Context, client *common.DatabricksClient, state *terraform.InstanceState) error) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		rs, ok := s.RootModule().Resources[name]
+		if !ok {
+			return fmt.Errorf("not found: %s", name)
+		}
+		client, err := client.New(&config.Config{})
+		if err != nil {
+			panic(err)
+		}
+		return cb(context.Background(), &common.DatabricksClient{
+			DatabricksClient: client,
+		}, rs.Primary)
+	}
+}
+
 const fullCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
 const hexCharset = "0123456789abcdef"
 
diff --git a/internal/acceptance/mws_workspaces_test.go b/internal/acceptance/mws_workspaces_test.go
@@ -1,7 +1,13 @@
 package acceptance
 
 import (
+	"context"
 	"testing"
+
+	"github.com/databricks/terraform-provider-databricks/common"
+	"github.com/databricks/terraform-provider-databricks/tokens"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestMwsAccWorkspaces(t *testing.T) {
@@ -54,6 +60,129 @@ func TestMwsAccWorkspaces(t *testing.T) {
 	})
 }
 
+func TestMwsAccWorkspacesTokenUpdate(t *testing.T) {
+	accountLevel(t, step{
+		Template: `
+		resource "databricks_mws_credentials" "this" {
+			account_id       = "{env.DATABRICKS_ACCOUNT_ID}"
+			credentials_name = "credentials-ws-{var.RANDOM}"
+			role_arn         = "{env.TEST_CROSSACCOUNT_ARN}"
+		}
+		resource "databricks_mws_customer_managed_keys" "this" {
+			account_id   = "{env.DATABRICKS_ACCOUNT_ID}"
+			aws_key_info {
+				key_arn   = "{env.TEST_MANAGED_KMS_KEY_ARN}"
+				key_alias = "{env.TEST_MANAGED_KMS_KEY_ALIAS}"
+			}
+			use_cases = ["MANAGED_SERVICES"]
+		}
+		resource "databricks_mws_storage_configurations" "this" {
+			account_id                 = "{env.DATABRICKS_ACCOUNT_ID}"
+			storage_configuration_name = "storage-ws-{var.RANDOM}"
+			bucket_name                = "{env.TEST_ROOT_BUCKET}"
+		}
+		resource "databricks_mws_workspaces" "this" {
+			account_id      = "{env.DATABRICKS_ACCOUNT_ID}"
+			workspace_name  = "terra-{var.RANDOM}"
+			aws_region      = "{env.AWS_REGION}"
+	
+			credentials_id = databricks_mws_credentials.this.credentials_id
+			storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id
+			managed_services_customer_managed_key_id = databricks_mws_customer_managed_keys.this.customer_managed_key_id
+
+			token {
+				comment = "test foo"
+			}
+		}`,
+		Check: resourceCheckWithState("databricks_mws_workspaces.this",
+			func(ctx context.Context, client *common.DatabricksClient, state *terraform.InstanceState) error {
+				workspaceUrl, ok := state.Attributes["workspace_url"]
+				assert.True(t, ok, "workspace_url is absent from databricks_mws_workspaces instance state")
+
+				workspaceClient, err := client.ClientForHost(ctx, workspaceUrl)
+				assert.NoError(t, err)
+
+				tokensAPI := tokens.NewTokensAPI(ctx, workspaceClient)
+				tokens, err := tokensAPI.List()
+				assert.NoError(t, err)
+
+				foundFoo := false
+				foundBar := false
+				for _, token := range tokens {
+					if token.Comment == "test foo" {
+						foundFoo = true
+					}
+					if token.Comment == "test bar" {
+						foundBar = true
+					}
+				}
+				assert.True(t, foundFoo)
+				assert.False(t, foundBar)
+				return nil
+			}),
+	},
+		step{
+			Template: `
+		resource "databricks_mws_credentials" "this" {
+			account_id       = "{env.DATABRICKS_ACCOUNT_ID}"
+			credentials_name = "credentials-ws-{var.RANDOM}"
+			role_arn         = "{env.TEST_CROSSACCOUNT_ARN}"
+		}
+		resource "databricks_mws_customer_managed_keys" "this" {
+			account_id   = "{env.DATABRICKS_ACCOUNT_ID}"
+			aws_key_info {
+				key_arn   = "{env.TEST_MANAGED_KMS_KEY_ARN}"
+				key_alias = "{env.TEST_MANAGED_KMS_KEY_ALIAS}"
+			}
+			use_cases = ["MANAGED_SERVICES"]
+		}
+		resource "databricks_mws_storage_configurations" "this" {
+			account_id                 = "{env.DATABRICKS_ACCOUNT_ID}"
+			storage_configuration_name = "storage-ws-{var.RANDOM}"
+			bucket_name                = "{env.TEST_ROOT_BUCKET}"
+		}
+		resource "databricks_mws_workspaces" "this" {
+			account_id      = "{env.DATABRICKS_ACCOUNT_ID}"
+			workspace_name  = "terra-{var.RANDOM}"
+			aws_region      = "{env.AWS_REGION}"
+	
+			credentials_id = databricks_mws_credentials.this.credentials_id
+			storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id
+			managed_services_customer_managed_key_id = databricks_mws_customer_managed_keys.this.customer_managed_key_id
+
+			token {
+				comment = "test bar"
+			}
+		}`,
+			Check: resourceCheckWithState("databricks_mws_workspaces.this",
+				func(ctx context.Context, client *common.DatabricksClient, state *terraform.InstanceState) error {
+					workspaceUrl, ok := state.Attributes["workspace_url"]
+					assert.True(t, ok, "workspace_url is absent from databricks_mws_workspaces instance state")
+
+					workspaceClient, err := client.ClientForHost(ctx, workspaceUrl)
+					assert.NoError(t, err)
+
+					tokensAPI := tokens.NewTokensAPI(ctx, workspaceClient)
+					tokens, err := tokensAPI.List()
+					assert.NoError(t, err)
+
+					foundFoo := false
+					foundBar := false
+					for _, token := range tokens {
+						if token.Comment == "test foo" {
+							foundFoo = true
+						}
+						if token.Comment == "test bar" {
+							foundBar = true
+						}
+					}
+					assert.False(t, foundFoo)
+					assert.True(t, foundBar)
+					return nil
+				}),
+		})
+}
+
 func TestMwsAccGcpWorkspaces(t *testing.T) {
 	accountLevel(t, step{
 		Template: `
diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go
@@ -560,9 +560,11 @@ func ResourceMwsWorkspaces() *schema.Resource {
 				workspace.CustomerManagedKeyID = ""
 			}
 			workspacesAPI := NewWorkspacesAPI(ctx, c)
-			err := workspacesAPI.UpdateRunning(workspace, d.Timeout(schema.TimeoutUpdate))
-			if err != nil {
-				return err
+			if d.HasChangeExcept("token") {
+				err := workspacesAPI.UpdateRunning(workspace, d.Timeout(schema.TimeoutUpdate))
+				if err != nil {
+					return err
+				}
 			}
 			return UpdateTokenIfNeeded(workspacesAPI, workspaceSchema, d)
 		},
diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go
@@ -990,6 +990,48 @@ func TestWorkspace_WaitForResolve(t *testing.T) {
 }
 
 func updateWorkspaceTokenFixture(t *testing.T, fixtures []qa.HTTPFixture, state map[string]string, hcl string) {
+	accountsAPI := []qa.HTTPFixture{
+		{
+			Method:       "GET",
+			ReuseRequest: true,
+			Resource:     "/api/2.0/accounts/c/workspaces/0",
+		},
+	}
+	tokensAPI := []qa.HTTPFixture{
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/token/list",
+			Response: `{}`, // we just need a JSON for this
+		},
+	}
+	tokensAPI = append(tokensAPI, fixtures...)
+	// outer HTTP server is used for inner request for "just created" workspace
+	qa.HTTPFixturesApply(t, tokensAPI, func(ctx context.Context, wsClient *common.DatabricksClient) {
+		// a bit hacky, but the whole thing is more readable
+		accountsAPI[0].Response = Workspace{
+			WorkspaceStatus: "RUNNING",
+			WorkspaceURL:    wsClient.Config.Host,
+		}
+		state["workspace_url"] = wsClient.Config.Host
+		state["workspace_name"] = "b"
+		state["account_id"] = "c"
+		state["network_id"] = "d"
+		state["is_no_public_ip_enabled"] = "false"
+		qa.ResourceFixture{
+			Fixtures:      accountsAPI,
+			Resource:      ResourceMwsWorkspaces(),
+			InstanceState: state,
+			Update:        true,
+			ID:            "a",
+			HCL: hcl + `
+			workspace_name = "b"
+			account_id = "c",
+			network_id = "d"`,
+		}.Apply(t)
+	})
+}
+
+func updateWorkspaceTokenFixtureWithPatch(t *testing.T, fixtures []qa.HTTPFixture, state map[string]string, hcl string) {
 	accountsAPI := []qa.HTTPFixture{
 		{
 			Method:   "PATCH",
@@ -1019,6 +1061,7 @@ func updateWorkspaceTokenFixture(t *testing.T, fixtures []qa.HTTPFixture, state
 		state["workspace_url"] = wsClient.Config.Host
 		state["workspace_name"] = "b"
 		state["account_id"] = "c"
+		state["storage_customer_managed_key_id"] = "1234"
 		state["is_no_public_ip_enabled"] = "false"
 		qa.ResourceFixture{
 			Fixtures:      accountsAPI,
@@ -1028,7 +1071,8 @@ func updateWorkspaceTokenFixture(t *testing.T, fixtures []qa.HTTPFixture, state
 			ID:            "a",
 			HCL: hcl + `
 			workspace_name = "b"
-			account_id = "c"`,
+			account_id = "c",
+			storage_customer_managed_key_id = "1234"`,
 		}.Apply(t)
 	})
 }
@@ -1054,6 +1098,53 @@ func TestUpdateWorkspace_AddToken(t *testing.T) {
 	}, `token {}`)
 }
 
+func TestUpdateWorkspace_AddTokenAndChangeNetworkId(t *testing.T) {
+	updateWorkspaceTokenFixtureWithPatch(t, []qa.HTTPFixture{
+		{
+			Method:   "POST",
+			Resource: "/api/2.0/token/create",
+			ExpectedRequest: Token{
+				LifetimeSeconds: 2.592e+06,
+				Comment:         "Terraform PAT",
+			},
+			Response: tokens.TokenResponse{
+				TokenValue: "sensitive",
+				TokenInfo: &tokens.TokenInfo{
+					TokenID: "abcdef",
+				},
+			},
+		},
+	}, map[string]string{
+		"network_id": "alpha",
+		// no token in state
+	}, `
+	network_id = "beta"
+	token {}
+	`)
+}
+
+func TestUpdateWorkspace_DeleteTokenAndChangeNetworkId(t *testing.T) {
+	updateWorkspaceTokenFixtureWithPatch(t, []qa.HTTPFixture{
+		{
+			Method:   "POST",
+			Resource: "/api/2.0/token/delete",
+			ExpectedRequest: map[string]any{
+				"token_id": "abcdef",
+			},
+		},
+	}, map[string]string{
+		"token.#":                  "1",
+		"token.0.comment":          "Terraform PAT",
+		"token.0.lifetime_seconds": "2592000",
+		"token.0.token_id":         "abcdef",
+		"token.0.token_value":      "sensitive",
+		"network_id":               "alpha",
+	}, `
+	network_id = "beta"
+	`)
+
+}
+
 func TestUpdateWorkspace_DeleteToken(t *testing.T) {
 	updateWorkspaceTokenFixture(t, []qa.HTTPFixture{
 		{
@@ -1072,6 +1163,44 @@ func TestUpdateWorkspace_DeleteToken(t *testing.T) {
 	}, ``)
 }
 
+func TestUpdateWorkspace_ReplaceTokenAndChangeNetworkId(t *testing.T) {
+	updateWorkspaceTokenFixtureWithPatch(t, []qa.HTTPFixture{
+		{
+			Method:   "POST",
+			Resource: "/api/2.0/token/delete",
+			ExpectedRequest: map[string]any{
+				"token_id": "abcdef",
+			},
+		},
+		{
+			Method:   "POST",
+			Resource: "/api/2.0/token/create",
+			ExpectedRequest: Token{
+				LifetimeSeconds: 2.592e+06,
+				Comment:         "I am Batman!",
+			},
+			Response: tokens.TokenResponse{
+				TokenValue: "new-value",
+				TokenInfo: &tokens.TokenInfo{
+					TokenID: "new-id",
+				},
+			},
+		},
+	}, map[string]string{
+		"token.#":                  "1",
+		"token.0.comment":          "Terraform PAT",
+		"token.0.lifetime_seconds": "2592000",
+		"token.0.token_id":         "abcdef",
+		"token.0.token_value":      "sensitive",
+		"network_id":               "alpha",
+	},
+		`
+	network_id = "beta"
+	token { 
+		comment = "I am Batman!"
+	}`)
+}
+
 func TestUpdateWorkspace_ReplaceToken(t *testing.T) {
 	updateWorkspaceTokenFixture(t, []qa.HTTPFixture{
 		{
