Support account-level provider with workspace-level resources

clusters/data_spark_version.go

@@ -131,15 +131,20 @@ func DataSourceSparkVersion() *schema.Resource {
 
 		s["photon"].Deprecated = "Specify runtime_engine=\"PHOTON\" in the cluster configuration"
 		s["graviton"].Deprecated = "Not required anymore - it's automatically enabled on the Graviton-based node types"
+		common.AddWorkspaceIdField(s)
 		return s
 	})
 
 	return &schema.Resource{
 		Schema: s,
 		ReadContext: func(ctx context.Context, d *schema.ResourceData, m any) diag.Diagnostics {
+			c, err := m.(*common.DatabricksClient).InConfiguredWorkspace(ctx, d)
+			if err != nil {
+				return diag.FromErr(err)
+			}
 			var this SparkVersionRequest
 			common.DataToStructPointer(d, s, &this)
-			version, err := NewClustersAPI(ctx, m).LatestSparkVersion(this)
+			version, err := NewClustersAPI(ctx, c).LatestSparkVersion(this)
 			if err != nil {
 				return diag.FromErr(err)
 			}

clusters/resource_cluster.go

@@ -166,6 +166,7 @@ func (ClusterResourceProvider) CustomizeSchema(s map[string]*schema.Schema) map[
 			Schema: common.StructToSchema(MountInfo{}, nil),
 		},
 	})
+	common.AddWorkspaceIdField(s)
 
 	return s
 }

common/client.go

@@ -48,6 +48,16 @@ type DatabricksClient struct {
 	cachedWorkspaceClient *databricks.WorkspaceClient
 	cachedAccountClient   *databricks.AccountClient
 	mu                    sync.Mutex
+
+	cachedWorkspaceClients   map[int64]*databricks.WorkspaceClient
+	cachedWorkspaceClientsMu sync.Mutex
+}
+
+func addCachedMe(w *databricks.WorkspaceClient) {
+	internalImpl := w.CurrentUser.Impl()
+	w.CurrentUser.WithImpl(&cachedMe{
+		internalImpl: internalImpl,
+	})
 }
 
 func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error) {
@@ -60,14 +70,59 @@ func (c *DatabricksClient) WorkspaceClient() (*databricks.WorkspaceClient, error
 	if err != nil {
 		return nil, err
 	}
-	internalImpl := w.CurrentUser.Impl()
-	w.CurrentUser.WithImpl(&cachedMe{
-		internalImpl: internalImpl,
-	})
+	addCachedMe(w)
 	c.cachedWorkspaceClient = w
 	return w, nil
 }
 
+func (c *DatabricksClient) InConfiguredWorkspace(ctx context.Context, d *schema.ResourceData) (*DatabricksClient, error) {
+	workspaceId, ok := d.GetOk("workspace_id")
+	if !ok {
+		return c, nil
+	}
+	w, err := c.getConfiguredWorkspaceClient(ctx, int64(workspaceId.(int)))
+	if err != nil {
+		return nil, err
+	}
+	apiClient, err := client.New(w.Config)
+	if err != nil {
+		return nil, err
+	}
+	return &DatabricksClient{
+		DatabricksClient:      apiClient,
+		cachedWorkspaceClient: w,
+	}, nil
+}
+
+func (c *DatabricksClient) getConfiguredWorkspaceClient(ctx context.Context, workspaceId int64) (*databricks.WorkspaceClient, error) {
+	if w, ok := c.cachedWorkspaceClients[workspaceId]; ok {
+		return w, nil
+	}
+	c.cachedWorkspaceClientsMu.Lock()
+	defer c.cachedWorkspaceClientsMu.Unlock()
+	if w, ok := c.cachedWorkspaceClients[workspaceId]; ok {
+		return w, nil
+	}
+	a, err := c.AccountClient()
+	if err != nil {
+		return nil, err
+	}
+	if c.cachedWorkspaceClients == nil {
+		c.cachedWorkspaceClients = make(map[int64]*databricks.WorkspaceClient)
+	}
+	ws, err := a.Workspaces.GetByWorkspaceId(ctx, workspaceId)
+	if err != nil {
+		return nil, err
+	}
+	w, err := a.GetWorkspaceClient(*ws)
+	if err != nil {
+		return nil, err
+	}
+	addCachedMe(w)
+	c.cachedWorkspaceClients[workspaceId] = w
+	return w, nil
+}
+
 // Set the cached workspace client.
 func (c *DatabricksClient) SetWorkspaceClient(w *databricks.WorkspaceClient) {
 	c.mu.Lock()
@@ -97,6 +152,9 @@ func (c *DatabricksClient) setAccountId(accountId string) error {
 }
 
 func (c *DatabricksClient) AccountClient() (*databricks.AccountClient, error) {
+	if c.cachedAccountClient != nil {
+		return c.cachedAccountClient, nil
+	}
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.cachedAccountClient != nil {

common/resource.go

@@ -87,7 +87,10 @@ func (r Resource) ToResource() *schema.Resource {
 	if r.Update != nil {
 		update = func(ctx context.Context, d *schema.ResourceData,
 			m any) diag.Diagnostics {
-			c := m.(*DatabricksClient)
+			c, err := m.(*DatabricksClient).InConfiguredWorkspace(ctx, d)
+			if err != nil {
+				return diag.FromErr(err)
+			}
 			if err := recoverable(r.Update)(ctx, d, c); err != nil {
 				err = nicerError(ctx, err, "update")
 				return diag.FromErr(err)
@@ -124,7 +127,11 @@ func (r Resource) ToResource() *schema.Resource {
 		m any) diag.Diagnostics {
 		return func(ctx context.Context, d *schema.ResourceData,
 			m any) diag.Diagnostics {
-			err := recoverable(r.Read)(ctx, d, m.(*DatabricksClient))
+			c, err := m.(*DatabricksClient).InConfiguredWorkspace(ctx, d)
+			if err != nil {
+				return diag.FromErr(err)
+			}
+			err = recoverable(r.Read)(ctx, d, c)
 			// TODO: https://github.com/databricks/terraform-provider-databricks/issues/2021
 			if ignoreMissing && apierr.IsMissing(err) {
 				log.Printf("[INFO] %s[id=%s] is removed on backend",
@@ -146,8 +153,11 @@ func (r Resource) ToResource() *schema.Resource {
 		CustomizeDiff:  r.saferCustomizeDiff(),
 		CreateContext: func(ctx context.Context, d *schema.ResourceData,
 			m any) diag.Diagnostics {
-			c := m.(*DatabricksClient)
-			err := recoverable(r.Create)(ctx, d, c)
+			c, err := m.(*DatabricksClient).InConfiguredWorkspace(ctx, d)
+			if err != nil {
+				return diag.FromErr(err)
+			}
+			err = recoverable(r.Create)(ctx, d, c)
 			if err != nil {
 				err = nicerError(ctx, err, "create")
 				return diag.FromErr(err)
@@ -162,7 +172,11 @@ func (r Resource) ToResource() *schema.Resource {
 		UpdateContext: update,
 		DeleteContext: func(ctx context.Context, d *schema.ResourceData,
 			m any) diag.Diagnostics {
-			err := recoverable(r.Delete)(ctx, d, m.(*DatabricksClient))
+			c, err := m.(*DatabricksClient).InConfiguredWorkspace(ctx, d)
+			if err != nil {
+				return diag.FromErr(err)
+			}
+			err = recoverable(r.Delete)(ctx, d, c)
 			if apierr.IsMissing(err) {
 				// TODO: https://github.com/databricks/terraform-provider-databricks/issues/2021
 				log.Printf("[INFO] %s[id=%s] is removed on backend",
@@ -447,3 +461,13 @@ func AddAccountIdField(s map[string]*schema.Schema) map[string]*schema.Schema {
 	}
 	return s
 }
+
+func AddWorkspaceIdField(s map[string]*schema.Schema) map[string]*schema.Schema {
+	s["workspace_id"] = &schema.Schema{
+		Type:        schema.TypeInt,
+		Optional:    true,
+		ForceNew:    true,
+		Description: "The workspace id of the workspace, for example 1234567890. This can be retrieved from `databricks_mws_workspaces.<YOUR_WORKSPACE>.workspace_id`. This attribute is required when using an account-level provider.",
+	}
+	return s
+}

go.mod

@@ -82,3 +82,5 @@ require (
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace github.com/databricks/databricks-sdk-go => /Users/miles/databricks-sdk-go

internal/acceptance/cluster_test.go

@@ -69,3 +69,27 @@ func TestAccClusterResource_CreateSingleNodeCluster(t *testing.T) {
 		}`,
 	})
 }
+
+func TestMwsAccClusterResource_CreateClusterWithAccountLevelProvider(t *testing.T) {
+	accountLevel(t, step{
+		Template: `
+		data "databricks_spark_version" "latest" {
+			workspace_id = 5549111504720597
+		}
+		resource "databricks_cluster" "this" {
+			workspace_id = 5549111504720597
+			cluster_name = "singlenode-{var.RANDOM}"
+			spark_version = data.databricks_spark_version.latest.id
+			instance_pool_id = "1127-101603-faked620-pool-c10aq68g"
+			num_workers = 0
+			autotermination_minutes = 10
+			spark_conf = {
+				"spark.databricks.cluster.profile" = "singleNode"
+				"spark.master" = "local[*]"
+			}
+			custom_tags = {
+				"ResourceClass" = "SingleNode"
+			}
+		}`,
+	})
+}