diff --git a/common/client.go b/common/client.go index 53400261be..a9f26f345c 100644 --- a/common/client.go +++ b/common/client.go @@ -84,7 +84,7 @@ func (c *DatabricksClient) GetWorkspaceClientForUnifiedProvider( ctx context.Context, workspaceID string, ) (*databricks.WorkspaceClient, diag.Diagnostics) { // The provider can be configured at account level or workspace level. - if c.Config.IsAccountClient() { + if c.Config.ConfigType() == config.AccountConfig { return c.getWorkspaceClientForAccountConfiguredProvider(ctx, workspaceID) } return c.getWorkspaceClientForWorkspaceConfiguredProvider(ctx, workspaceID) @@ -240,6 +240,8 @@ func (c *DatabricksClient) SetWorkspaceClient(w *databricks.WorkspaceClient) { } func (c *DatabricksClient) WorkspaceClientForWorkspace(ctx context.Context, workspaceId int64) (*databricks.WorkspaceClient, error) { + // TODO: For unified hosts, we can just make the client using the WorkspaceId without making an API call to accounts, and + // the host should stay unchanged. We'll still want to make a copy of the config, of course. c.mu.Lock() defer c.mu.Unlock() if c.cachedWorkspaceClients == nil { @@ -354,7 +356,7 @@ func (c *DatabricksClient) AccountClientWithAccountIdFromPair(d *schema.Resource } func (c *DatabricksClient) AccountOrWorkspaceRequest(accCallback func(*databricks.AccountClient) error, wsCallback func(*databricks.WorkspaceClient) error) error { - if c.Config.IsAccountClient() { + if c.Config.ConfigType() == config.AccountConfig { a, err := c.AccountClient() if err != nil { return err @@ -427,7 +429,7 @@ func (c *DatabricksClient) addApiPrefix(r *http.Request) error { // scimVisitor is a separate method for the sake of unit tests func (c *DatabricksClient) scimVisitor(r *http.Request) error { - if c.Config.IsAccountClient() && c.Config.AccountID != "" { + if c.Config.ConfigType() == config.AccountConfig && c.Config.AccountID != "" { // until `/preview` is there for workspace scim, // `/api/2.0` is added by completeUrl visitor r.URL.Path = strings.ReplaceAll(r.URL.Path, "/api/2.0/preview", @@ -471,7 +473,7 @@ func (c *DatabricksClient) FormatURL(strs ...string) string { // ClientForHost creates a new DatabricksClient instance with the same auth parameters, // but for the given host. Authentication has to be reinitialized, as Google OIDC has // different authorizers, depending if it's workspace or Accounts API we're talking to. -func (c *DatabricksClient) ClientForHost(ctx context.Context, url string) (*DatabricksClient, error) { +func (c *DatabricksClient) ClientForHost(ctx context.Context, url string, workspaceId int64) (*DatabricksClient, error) { // create dummy http request req, _ := http.NewRequestWithContext(ctx, "GET", "/", nil) // Ensure that client is authenticated @@ -483,6 +485,8 @@ func (c *DatabricksClient) ClientForHost(ctx context.Context, url string) (*Data if err != nil { return nil, fmt.Errorf("cannot configure new client: %w", err) } + // Set the workspace ID for unified hosts, which need the X-Databricks-Org-Id header + cfg.WorkspaceId = fmt.Sprintf("%d", workspaceId) // TODO: should have made WorkspaceId an int64 in the go SDK config. It's not too late to change it. client, err := client.New(cfg) if err != nil { return nil, fmt.Errorf("cannot configure new client: %w", err) diff --git a/common/client_test.go b/common/client_test.go index e967dae9f4..e43875228e 100644 --- a/common/client_test.go +++ b/common/client_test.go @@ -237,11 +237,12 @@ func TestClientForHost(t *testing.T) { }) assert.NoError(t, err) assert.True(t, dc.IsAws()) - cc, err := dc.ClientForHost(context.Background(), "https://e2-workspace.cloud.databricks.com/") + cc, err := dc.ClientForHost(context.Background(), "https://e2-workspace.cloud.databricks.com/", 12345) assert.NoError(t, err) assert.Equal(t, dc.Config.Username, cc.Config.Username) assert.Equal(t, dc.Config.Password, cc.Config.Password) assert.NotEqual(t, dc.Config.Host, cc.Config.Host) + assert.Equal(t, "12345", cc.Config.WorkspaceId) } func TestClientForHostAuthError(t *testing.T) { @@ -254,7 +255,7 @@ func TestClientForHostAuthError(t *testing.T) { }, }, } - _, err := c.ClientForHost(context.Background(), "https://e2-workspace.cloud.databricks.com/") + _, err := c.ClientForHost(context.Background(), "https://e2-workspace.cloud.databricks.com/", 12345) assert.NoError(t, err) } diff --git a/exporter/context.go b/exporter/context.go index d64b7506a6..99346e31da 100644 --- a/exporter/context.go +++ b/exporter/context.go @@ -16,6 +16,7 @@ import ( "time" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/databricks/databricks-sdk-go/service/compute" "golang.org/x/exp/maps" @@ -397,7 +398,7 @@ func (ic *importContext) Run() error { return fmt.Errorf("the path %s is not a directory", ic.Directory) } - ic.accountLevel = ic.Client.Config.IsAccountClient() + ic.accountLevel = ic.Client.Config.ConfigType() == config.AccountConfig if ic.accountLevel { ic.meAdmin = true // TODO: check if we can get the current user from the account client diff --git a/go.mod b/go.mod index e50a343b5f..eef404e636 100644 --- a/go.mod +++ b/go.mod @@ -3,12 +3,11 @@ module github.com/databricks/terraform-provider-databricks go 1.24.0 require ( - github.com/databricks/databricks-sdk-go v0.86.0 + github.com/databricks/databricks-sdk-go v0.90.0 github.com/golang-jwt/jwt/v4 v4.5.2 github.com/hashicorp/go-cty v1.5.0 github.com/hashicorp/hcl v1.0.0 github.com/hashicorp/hcl/v2 v2.24.0 - github.com/hashicorp/terraform-json v0.27.2 github.com/hashicorp/terraform-plugin-framework v1.16.1 github.com/hashicorp/terraform-plugin-framework-validators v0.19.0 github.com/hashicorp/terraform-plugin-go v0.29.0 @@ -59,6 +58,7 @@ require ( github.com/hashicorp/hc-install v0.9.2 // indirect github.com/hashicorp/logutils v1.0.0 // indirect github.com/hashicorp/terraform-exec v0.23.1 // indirect + github.com/hashicorp/terraform-json v0.27.2 // indirect github.com/hashicorp/terraform-registry-address v0.4.0 // indirect github.com/hashicorp/terraform-svchost v0.1.1 // indirect github.com/hashicorp/yamux v0.1.2 // indirect diff --git a/go.sum b/go.sum index 375784f5d4..fef44bb59a 100644 --- a/go.sum +++ b/go.sum @@ -25,8 +25,8 @@ github.com/cloudflare/circl v1.6.1 h1:zqIqSPIndyBh1bjLVVDHMPpVKqp8Su/V+6MeDzzQBQ github.com/cloudflare/circl v1.6.1/go.mod h1:uddAzsPgqdMAYatqJ0lsjX1oECcQLIlRpzZh3pJrofs= github.com/cyphar/filepath-securejoin v0.4.1 h1:JyxxyPEaktOD+GAnqIqTf9A8tHyAG22rowi7HkoSU1s= github.com/cyphar/filepath-securejoin v0.4.1/go.mod h1:Sdj7gXlvMcPZsbhwhQ33GguGLDGQL7h7bg04C/+u9jI= -github.com/databricks/databricks-sdk-go v0.86.0 h1:Di1+NBQlfzBMUhY6w6gS2mtmNXIWycowoCsLCGFQPyU= -github.com/databricks/databricks-sdk-go v0.86.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE= +github.com/databricks/databricks-sdk-go v0.90.0 h1:yS9kBd/4LEp3WzzFZ/th49A+WXn0bb6rEWLAYL4T6/8= +github.com/databricks/databricks-sdk-go v0.90.0/go.mod h1:hWoHnHbNLjPKiTm5K/7bcIv3J3Pkgo5x9pPzh8K3RVE= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= diff --git a/mws/mws_workspaces_test.go b/mws/mws_workspaces_test.go index 1f12b1299f..fcfffaa3ef 100644 --- a/mws/mws_workspaces_test.go +++ b/mws/mws_workspaces_test.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "regexp" + "strconv" "testing" "time" @@ -146,8 +147,12 @@ func TestMwsAccWorkspacesTokenUpdate(t *testing.T) { 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") + workspaceIdStr, ok := state.Attributes["workspace_id"] + assert.True(t, ok, "workspace_id is absent from databricks_mws_workspaces instance state") + workspaceId, err := strconv.ParseInt(workspaceIdStr, 10, 64) + assert.NoError(t, err) - workspaceClient, err := client.ClientForHost(ctx, workspaceUrl) + workspaceClient, err := client.ClientForHost(ctx, workspaceUrl, workspaceId) assert.NoError(t, err) tokensAPI := tokens.NewTokensAPI(ctx, workspaceClient) @@ -219,8 +224,12 @@ func TestMwsAccWorkspacesTokenUpdate(t *testing.T) { 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") + workspaceIdStr, ok := state.Attributes["workspace_id"] + assert.True(t, ok, "workspace_id is absent from databricks_mws_workspaces instance state") + workspaceId, err := strconv.ParseInt(workspaceIdStr, 10, 64) + assert.NoError(t, err) - workspaceClient, err := client.ClientForHost(ctx, workspaceUrl) + workspaceClient, err := client.ClientForHost(ctx, workspaceUrl, workspaceId) assert.NoError(t, err) tokensAPI := tokens.NewTokensAPI(ctx, workspaceClient) diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index fe4754c517..ffdb9a4e46 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -208,7 +208,7 @@ func (a WorkspacesAPI) verifyWorkspaceReachable(ws Workspace) *resource.RetryErr defer cancel() // wait for DNS caches to refresh, as sometimes we cannot make // API calls to new workspaces immediately after it's created - wsClient, err := a.client.ClientForHost(a.context, ws.WorkspaceURL) + wsClient, err := a.client.ClientForHost(a.context, ws.WorkspaceURL, ws.WorkspaceID) if err != nil { return resource.NonRetryableError(err) } @@ -410,6 +410,7 @@ func (s SensitiveString) String() string { // ephemeral entity to use with StructToData() type WorkspaceToken struct { WorkspaceURL string `json:"workspace_url,omitempty"` + WorkspaceId int64 `json:"workspace_id,omitempty"` Token *Token `json:"token,omitempty"` } @@ -420,7 +421,7 @@ func CreateTokenIfNeeded(workspacesAPI WorkspacesAPI, if wsToken.Token == nil { return nil } - client, err := workspacesAPI.client.ClientForHost(workspacesAPI.context, wsToken.WorkspaceURL) + client, err := workspacesAPI.client.ClientForHost(workspacesAPI.context, wsToken.WorkspaceURL, wsToken.WorkspaceId) if err != nil { return err } @@ -458,7 +459,7 @@ func EnsureTokenExistsIfNeeded(a WorkspacesAPI, if wsToken.Token == nil { return nil } - client, err := a.client.ClientForHost(a.context, wsToken.WorkspaceURL) + client, err := a.client.ClientForHost(a.context, wsToken.WorkspaceURL, wsToken.WorkspaceId) if err != nil { return err } @@ -481,7 +482,9 @@ func EnsureTokenExistsIfNeeded(a WorkspacesAPI, } func removeTokenIfNeeded(a WorkspacesAPI, tokenID string, d *schema.ResourceData) error { - client, err := a.client.ClientForHost(a.context, d.Get("workspace_url").(string)) + workspaceID := int64(d.Get("workspace_id").(int)) + workspaceURL := d.Get("workspace_url").(string) + client, err := a.client.ClientForHost(a.context, workspaceURL, workspaceID) if err != nil { return err }