Fixed mount read failing if cluster no longer exists (#2634)

* Fix mount read failing if cluster no longer exits

* fmt

* add acceptance test

* Add pinned cluster test

* lint

* a comment

* -

* put closure inside func

* fix bug in func check

* address comment

* Add docs

* Address comments

Author: shreyas-goenka

common/client.go

@@ -42,7 +42,7 @@ func (a *cachedMe) Me(ctx context.Context) (*iam.User, error) {
 type DatabricksClient struct {
 	*client.DatabricksClient
 
-	// callback used to create API1.2 call wrapper, which simplifies unit tessting
+	// callback used to create API1.2 call wrapper, which simplifies unit testing
 	commandFactory        func(context.Context, *DatabricksClient) CommandExecutor
 	cachedWorkspaceClient *databricks.WorkspaceClient
 	cachedAccountClient   *databricks.AccountClient

docs/resources/mount.md

@@ -7,6 +7,8 @@ This resource will [mount your cloud storage](https://docs.databricks.com/data/d
 
 **Note** When `cluster_id` is not specified, it will create the smallest possible cluster in the default availability zone with name equal to or starting with `terraform-mount` for the shortest possible amount of time. To avoid mount failure due to potentially quota or capacity issues with the default cluster, we recommend specifying a cluster to use for mounting.
 
+**Note** CRUD operations on a databricks mount require a running cluster. Due to limitations of terraform and the databricks mounts APIs, if the cluster the mount was most recently created / updated using no longer exists AND the mount is destroyed as a part of a terraform apply, we mark it as deleted without cleaning it up from the workspace.
+
 This resource provides two ways of mounting a storage account:
 
 1. Use a storage-specific configuration block - this could be used for the most cases, as it will fill most of the necessary details. Currently we support following configuration blocks:

internal/acceptance/init_test.go

@@ -83,9 +83,17 @@ func unityAccountLevel(t *testing.T, steps ...step) {
 	run(t, steps)
 }
 
+// A step in a terraform acceptance test
 type step struct {
+	// Terraform HCL for resources to materialize in this test step.
 	Template string
-	Check    func(*terraform.State) error
+
+	// This function is called after the template is applied. Useful for making assertions
+	// or doing cleanup.
+	Check func(*terraform.State) error
+
+	// Setup function called before the template is materialized.
+	PreConfig func()
 
 	Destroy                   bool
 	ExpectNonEmptyPlan        bool
@@ -148,7 +156,8 @@ func environmentTemplate(t *testing.T, template string, otherVars ...map[string]
 	return commands.TrimLeadingWhitespace(template)
 }
 
-// Test wrapper over terraform testing framework
+// Test wrapper over terraform testing framework. Multiple steps share the same
+// terraform state context.
 func run(t *testing.T, steps []step) {
 	cloudEnv := os.Getenv("CLOUD_ENV")
 	if cloudEnv == "" {
@@ -170,13 +179,7 @@ func run(t *testing.T, steps []step) {
 	}
 	ts := []resource.TestStep{}
 	ctx := context.Background()
-	type testResource struct {
-		ID       string
-		Name     string
-		Resource *schema.Resource
-	}
 
-	resourcesEverCreated := map[testResource]bool{}
 	stepConfig := ""
 	for i, s := range steps {
 		if s.Template != "" {
@@ -185,6 +188,7 @@ func run(t *testing.T, steps []step) {
 		stepNum := i
 		thisStep := s
 		stepCheck := thisStep.Check
+		stepPreConfig := s.PreConfig
 		ts = append(ts, resource.TestStep{
 			PreConfig: func() {
 				if stepConfig == "" {
@@ -193,6 +197,10 @@ func run(t *testing.T, steps []step) {
 				logger.Infof(ctx, "Test %s (%s) step %d config is:\n%s",
 					t.Name(), cloudEnv, stepNum,
 					commands.TrimLeadingWhitespace(stepConfig))
+
+				if stepPreConfig != nil {
+					stepPreConfig()
+				}
 			},
 			Config:                    stepConfig,
 			Destroy:                   s.Destroy,
@@ -205,20 +213,17 @@ func run(t *testing.T, steps []step) {
 			Check: func(state *terraform.State) error {
 				// get configured client from provider
 				client := provider.Meta().(*common.DatabricksClient)
+
+				// Default check for all runs. Asserts that the read operation succeeds.
 				for n, is := range state.RootModule().Resources {
 					p := strings.Split(n, ".")
+
+					// Skip data resources.
 					if p[0] == "data" {
 						continue
 					}
 					r := provider.ResourcesMap[p[0]]
-					resourcesEverCreated[testResource{
-						ID:       is.Primary.ID,
-						Name:     p[1],
-						Resource: r,
-					}] = true
-					dia := r.ReadContext(ctx, r.Data(&terraform.InstanceState{
-						ID: is.Primary.ID,
-					}), client)
+					dia := r.ReadContext(ctx, r.Data(is.Primary), client)
 					if dia != nil {
 						return fmt.Errorf("%v", dia)
 					}

internal/acceptance/mounts_test.go

@@ -0,0 +1,103 @@
+package acceptance
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"golang.org/x/exp/maps"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+var mountHcl = `
+data "databricks_spark_version" "latest" {}
+
+# Test cluster to create the mount using.
+resource "databricks_cluster" "this" {
+	cluster_name = "acc-test-mounts-{var.STICKY_RANDOM}"
+	spark_version = data.databricks_spark_version.latest.id
+	instance_pool_id = "{env.TEST_INSTANCE_POOL_ID}"
+	num_workers = 1
+
+	aws_attributes {
+		instance_profile_arn = "{env.TEST_INSTANCE_PROFILE_ARN}"
+	}
+}
+
+resource "databricks_mount" "my_mount" {
+	name = "test-mount-{var.STICKY_RANDOM}"
+	cluster_id = databricks_cluster.this.id
+	
+	s3 {
+		bucket_name      = "{env.TEST_S3_BUCKET_NAME}"
+	}
+}`
+
+func TestAccCreateDatabricksMount(t *testing.T) {
+	workspaceLevel(t,
+		step{
+			Template: mountHcl,
+		})
+}
+
+func TestAccCreateDatabricksMountIsFineOnClusterRecreate(t *testing.T) {
+	clusterId1 := ""
+	clusterId2 := ""
+
+	workspaceLevel(t,
+		// Step 1 creates the cluster and mount.
+		step{
+			Template: mountHcl,
+			Check: func(s *terraform.State) error {
+				resources := s.RootModule().Resources
+				cluster := resources["databricks_cluster.this"]
+				if cluster == nil {
+					return fmt.Errorf("expected to find databricks_cluster.this in resources keys: %v", maps.Keys(resources))
+				}
+				clusterId1 = cluster.Primary.ID
+
+				// Assert cluster id is not empty. This is later used to ensure
+				// the cluster has been recreated.
+				assert.NotEmpty(t, clusterId1)
+
+				// Assert the mount points to the created cluster.
+				mount := resources["databricks_mount.my_mount"]
+				assert.Equal(t, clusterId1, mount.Primary.Attributes["cluster_id"])
+				return nil
+			},
+		},
+		// Step 2: Manually delete the cluster, and then reapply the config. The mount
+		// will be recreated in this case.
+		step{
+			PreConfig: func() {
+				w, err := databricks.NewWorkspaceClient(&databricks.Config{})
+				require.NoError(t, err)
+				err = w.Clusters.PermanentDeleteByClusterId(context.Background(), clusterId1)
+				assert.NoError(t, err, "failed to delete the cluster, id: "+clusterId1)
+			},
+			Template: mountHcl,
+			Check: func(s *terraform.State) error {
+				resources := s.RootModule().Resources
+				cluster := resources["databricks_cluster.this"]
+				if cluster == nil {
+					return fmt.Errorf("expected to find databricks_cluster.this in resources keys: %v", maps.Keys(resources))
+				}
+				clusterId2 = cluster.Primary.ID
+
+				// Assert cluster was indeed recreated
+				assert.NotEmpty(t, clusterId1)
+				assert.NotEmpty(t, clusterId2)
+				assert.NotEqual(t, clusterId1, clusterId2)
+
+				// Assert the mount points to the newly recreated cluster.
+				mount := resources["databricks_mount.my_mount"]
+				assert.Equal(t, clusterId2, mount.Primary.Attributes["cluster_id"])
+				return nil
+			},
+		},
+	)
+}

storage/resource_mount.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/terraform-provider-databricks/clusters"
 	"github.com/databricks/terraform-provider-databricks/common"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -37,6 +39,42 @@ func (cb mountCallback) preProcess(r *schema.Resource) func(
 	}
 }
 
+// If a cluster_id is specified for the mount, this function validates that the
+// cluster actually exists. If the cluster does not exist, we assume (incorrectly)
+// that the mount also does not exist and remove it from the in memory representation
+// of the terraform state.
+//
+// In cases of update requests this would lead to an attempted recreation of the mount.
+//
+// In cases of delete, this could orphan the mount, ie that the mount still exists
+// but is not tracked by the terraform state.
+//
+// This is needed to resolve issue: https://github.com/databricks/terraform-provider-databricks/issues/1864
+// because without a running cluster we do not have a way to do CRUD on mounts.
+// The argument being that possibly orphaning mounts is a better experience than
+// requiring manual edits to tfstate.
+//
+// Why do we need a valid cluster_id to be set here?
+// Answer: Downstream read and delete code relies on a valid cluster_id being set
+// in the terraform state. If the cluster_id is invalid then there is no way to update
+// it using native terraform apply workflows. The only workaround is to manually edit the
+// tfstate file replacing the existing invalid cluster_id with a new valid one.
+func removeFromStateIfClusterDoesNotExist(cb func(context.Context, *schema.ResourceData, any) diag.Diagnostics) func(context.Context, *schema.ResourceData, any) diag.Diagnostics {
+	return func(ctx context.Context, d *schema.ResourceData, m any) diag.Diagnostics {
+		clusterId := d.Get("cluster_id").(string)
+		if clusterId != "" {
+			clustersAPI := clusters.NewClustersAPI(ctx, m)
+			_, err := clustersAPI.Get(clusterId)
+			if apierr.IsMissing(err) {
+				// Return with empty ID to indicate the resource no longer exists.
+				d.SetId("")
+				return nil
+			}
+		}
+		return cb(ctx, d, m)
+	}
+}
+
 func ResourceDatabricksMountSchema() map[string]*schema.Schema {
 	tpl := GenericMount{}
 	scm := common.StructToSchema(tpl, func(s map[string]*schema.Schema) map[string]*schema.Schema {
@@ -62,8 +100,8 @@ func ResourceMount() *schema.Resource {
 	tpl := GenericMount{}
 	r := commonMountResource(tpl, ResourceDatabricksMountSchema())
 	r.CreateContext = mountCallback(mountCreate).preProcess(r)
-	r.ReadContext = mountCallback(mountRead).preProcess(r)
-	r.DeleteContext = mountCallback(mountDelete).preProcess(r)
+	r.ReadContext = removeFromStateIfClusterDoesNotExist(mountCallback(mountRead).preProcess(r))
+	r.DeleteContext = removeFromStateIfClusterDoesNotExist(mountCallback(mountDelete).preProcess(r))
 	r.Importer = nil
 	r.Timeouts = &schema.ResourceTimeout{
 		Default: schema.DefaultTimeout(20 * time.Minute),