Add account-level API support for Unity Catalog objects

catalog/resource_metastore.go

@@ -5,69 +5,35 @@ import (
 	"log"
 	"strings"
 
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/terraform-provider-databricks/common"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
 )
 
-type MetastoresAPI struct {
-	client  *common.DatabricksClient
-	context context.Context
-}
-
-func NewMetastoresAPI(ctx context.Context, m any) MetastoresAPI {
-	return MetastoresAPI{m.(*common.DatabricksClient), context.WithValue(ctx, common.Api, common.API_2_1)}
-}
-
 type MetastoreInfo struct {
-	Name                                        string  `json:"name"`
-	StorageRoot                                 string  `json:"storage_root" tf:"force_new"`
-	DefaultDacID                                string  `json:"default_data_access_config_id,omitempty" tf:"suppress_diff"`
-	Owner                                       string  `json:"owner,omitempty" tf:"computed"`
-	MetastoreID                                 string  `json:"metastore_id,omitempty" tf:"computed"`
-	WorkspaceIDs                                []int64 `json:"workspace_ids,omitempty" tf:"computed"`
-	Region                                      string  `json:"region,omitempty" tf:"computed"`
-	Cloud                                       string  `json:"cloud,omitempty" tf:"computed"`
-	GlobalMetastoreId                           string  `json:"global_metastore_id,omitempty" tf:"computed"`
-	CreatedAt                                   int64   `json:"created_at,omitempty" tf:"computed"`
-	CreatedBy                                   string  `json:"created_by,omitempty" tf:"computed"`
-	UpdatedAt                                   int64   `json:"updated_at,omitempty" tf:"computed"`
-	UpdatedBy                                   string  `json:"updated_by,omitempty" tf:"computed"`
-	DeltaSharingScope                           string  `json:"delta_sharing_scope,omitempty" tf:"suppress_diff"`
-	DeltaSharingRecipientTokenLifetimeInSeconds int64   `json:"delta_sharing_recipient_token_lifetime_in_seconds,omitempty"`
-	DeltaSharingOrganizationName                string  `json:"delta_sharing_organization_name,omitempty"`
-}
-
-type CreateMetastore struct {
-	Name        string `json:"name"`
-	StorageRoot string `json:"storage_root"`
-}
-
-func (a MetastoresAPI) createMetastore(cm CreateMetastore) (mi MetastoreInfo, err error) {
-	err = a.client.Post(a.context, "/unity-catalog/metastores", cm, &mi)
-	return
-}
-
-func (a MetastoresAPI) getMetastore(id string) (mi MetastoreInfo, err error) {
-	err = a.client.Get(a.context, "/unity-catalog/metastores/"+id, nil, &mi)
-	return
-}
-
-func (a MetastoresAPI) updateMetastore(metastoreID string, update map[string]any) error {
-	return a.client.Patch(a.context, "/unity-catalog/metastores/"+metastoreID, update)
-}
-
-func (a MetastoresAPI) deleteMetastore(id string, force bool) error {
-	return a.client.Delete(a.context, "/unity-catalog/metastores/"+id, map[string]any{
-		"force": force,
-	})
+	Name                                        string `json:"name"`
+	StorageRoot                                 string `json:"storage_root" tf:"force_new"`
+	DefaultDacID                                string `json:"default_data_access_config_id,omitempty" tf:"suppress_diff"`
+	StorageRootCredentialId                     string `json:"storage_root_credential_id,omitempty" tf:"suppress_diff"`
+	Owner                                       string `json:"owner,omitempty" tf:"computed"`
+	MetastoreID                                 string `json:"metastore_id,omitempty" tf:"computed"`
+	Region                                      string `json:"region,omitempty" tf:"computed"`
+	Cloud                                       string `json:"cloud,omitempty" tf:"computed"`
+	GlobalMetastoreId                           string `json:"global_metastore_id,omitempty" tf:"computed"`
+	CreatedAt                                   int64  `json:"created_at,omitempty" tf:"computed"`
+	CreatedBy                                   string `json:"created_by,omitempty" tf:"computed"`
+	UpdatedAt                                   int64  `json:"updated_at,omitempty" tf:"computed"`
+	UpdatedBy                                   string `json:"updated_by,omitempty" tf:"computed"`
+	DeltaSharingScope                           string `json:"delta_sharing_scope,omitempty" tf:"suppress_diff"`
+	DeltaSharingRecipientTokenLifetimeInSeconds int64  `json:"delta_sharing_recipient_token_lifetime_in_seconds,omitempty"`
+	DeltaSharingOrganizationName                string `json:"delta_sharing_organization_name,omitempty"`
 }
 
 func ResourceMetastore() *schema.Resource {
 	s := common.StructToSchema(MetastoreInfo{},
 		func(m map[string]*schema.Schema) map[string]*schema.Schema {
-			delete(m, "metastore_id")
-			delete(m, "workspace_ids") // todo: bring it back when it works
 			m["force_destroy"] = &schema.Schema{
 				Type:     schema.TypeBool,
 				Optional: true,
@@ -84,32 +50,89 @@ func ResourceMetastore() *schema.Resource {
 			}
 			return m
 		})
-	update := updateFunctionFactory("/unity-catalog/metastores", []string{"owner", "name", "delta_sharing_scope",
-		"delta_sharing_recipient_token_lifetime_in_seconds", "delta_sharing_organization_name"})
 
 	return common.Resource{
 		Schema: s,
 		Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			var create CreateMetastore
+			var create catalog.CreateMetastore
+			var update catalog.UpdateMetastore
 			common.DataToStructPointer(d, s, &create)
-			mi, err := NewMetastoresAPI(ctx, c).createMetastore(create)
-			if err != nil {
-				return err
-			}
-			d.SetId(mi.MetastoreID)
-			return update(ctx, d, c)
+			common.DataToStructPointer(d, s, &update)
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				mi, err := acc.Metastores.Create(ctx,
+					catalog.AccountsCreateMetastore{
+						MetastoreInfo: &create,
+					})
+				if err != nil {
+					return err
+				}
+				d.SetId(mi.MetastoreInfo.MetastoreId)
+				_, err = acc.Metastores.Update(ctx, catalog.AccountsUpdateMetastore{
+					MetastoreId:   mi.MetastoreInfo.MetastoreId,
+					MetastoreInfo: &update,
+				})
+				if err != nil {
+					return err
+				}
+				return nil
+			}, func(w *databricks.WorkspaceClient) error {
+				mi, err := w.Metastores.Create(ctx, create)
+				if err != nil {
+					return err
+				}
+				d.SetId(mi.MetastoreId)
+				update.Id = mi.MetastoreId
+				_, err = w.Metastores.Update(ctx, update)
+				if err != nil {
+					return err
+				}
+				return nil
+			})
 		},
 		Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			mi, err := NewMetastoresAPI(ctx, c).getMetastore(d.Id())
-			if err != nil {
-				return err
-			}
-			return common.StructToData(mi, s, d)
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				mi, err := acc.Metastores.GetByMetastoreId(ctx, d.Id())
+				if err != nil {
+					return err
+				}
+				return common.StructToData(mi, s, d)
+			}, func(w *databricks.WorkspaceClient) error {
+				mi, err := w.Metastores.GetById(ctx, d.Id())
+				if err != nil {
+					return err
+				}
+				return common.StructToData(mi, s, d)
+			})
+		},
+		Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			var update catalog.UpdateMetastore
+			common.DataToStructPointer(d, s, &update)
+
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				_, err := acc.Metastores.Update(ctx, catalog.AccountsUpdateMetastore{
+					MetastoreId:   d.Id(),
+					MetastoreInfo: &update,
+				})
+				if err != nil {
+					return err
+				}
+				return nil
+			}, func(w *databricks.WorkspaceClient) error {
+				update.Id = d.Id()
+				_, err := w.Metastores.Update(ctx, update)
+				if err != nil {
+					return err
+				}
+				return nil
+			})
 		},
-		Update: update,
 		Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
 			force := d.Get("force_destroy").(bool)
-			return NewMetastoresAPI(ctx, c).deleteMetastore(d.Id(), force)
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				return acc.Metastores.Delete(ctx, catalog.DeleteAccountMetastoreRequest{Force: force, MetastoreId: d.Id()})
+			}, func(w *databricks.WorkspaceClient) error {
+				return w.Metastores.Delete(ctx, catalog.DeleteMetastoreRequest{Force: force, Id: d.Id()})
+			})
 		},
 	}.ToResource()
 }

catalog/resource_metastore_assignment.go

@@ -2,53 +2,18 @@ package catalog
 
 import (
 	"context"
-	"fmt"
+	"strconv"
 
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/terraform-provider-databricks/common"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
-type MetastoreAssignmentAPI struct {
-	client  *common.DatabricksClient
-	context context.Context
-}
-
-func NewMetastoreAssignmentAPI(ctx context.Context, m any) MetastoreAssignmentAPI {
-	return MetastoreAssignmentAPI{m.(*common.DatabricksClient), context.WithValue(ctx, common.Api, common.API_2_1)}
-}
-
-type MetastoreAssignment struct {
-	WorkspaceID        int64  `json:"workspace_id" tf:"force_new"`
-	MetastoreID        string `json:"metastore_id"`
-	DefaultCatalogName string `json:"default_catalog_name,omitempty" tf:"default:hive_metastore"`
-}
-
-func (a MetastoreAssignmentAPI) createMetastoreAssignment(ma MetastoreAssignment) error {
-	path := fmt.Sprintf("/unity-catalog/workspaces/%d/metastore", ma.WorkspaceID)
-	return a.client.Put(a.context, path, ma)
-}
-
-func (a MetastoreAssignmentAPI) updateMetastoreAssignment(ma MetastoreAssignment) error {
-	path := fmt.Sprintf("/unity-catalog/workspaces/%d/metastore", ma.WorkspaceID)
-	return a.client.Patch(a.context, path, ma)
-}
-
-func (a MetastoreAssignmentAPI) getAssignedMetastoreID() (string, error) {
-	var mi MetastoreInfo
-	err := a.client.Get(a.context, "/unity-catalog/metastore_summary", nil, &mi)
-	return mi.MetastoreID, err
-}
-
-func (a MetastoreAssignmentAPI) deleteMetastoreAssignment(workspaceID, metastoreID string) error {
-	path := fmt.Sprintf("/unity-catalog/workspaces/%s/metastore", workspaceID)
-	return a.client.Delete(a.context, path, map[string]string{
-		"metastore_id": metastoreID,
-	})
-}
-
 func ResourceMetastoreAssignment() *schema.Resource {
-	s := common.StructToSchema(MetastoreAssignment{},
+	s := common.StructToSchema(catalog.MetastoreAssignment{},
 		func(m map[string]*schema.Schema) map[string]*schema.Schema {
+			m["default_catalog_name"].Default = "hive_metastore"
 			return m
 		})
 	pi := common.NewPairID("workspace_id", "metastore_id").Schema(
@@ -58,30 +23,95 @@ func ResourceMetastoreAssignment() *schema.Resource {
 	return common.Resource{
 		Schema: s,
 		Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			var ma MetastoreAssignment
-			common.DataToStructPointer(d, s, &ma)
-			if err := NewMetastoreAssignmentAPI(ctx, c).createMetastoreAssignment(ma); err != nil {
-				return err
-			}
-			pi.Pack(d)
-			return nil
+			workspaceId := int64(d.Get("workspace_id").(int))
+			metastoreId := d.Get("metastore_id").(string)
+			var create catalog.CreateMetastoreAssignment
+			common.DataToStructPointer(d, s, &create)
+			create.WorkspaceId = workspaceId
+
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				err := acc.MetastoreAssignments.Create(ctx,
+					catalog.AccountsCreateMetastoreAssignment{
+						WorkspaceId:         workspaceId,
+						MetastoreId:         metastoreId,
+						MetastoreAssignment: &create,
+					})
+				if err != nil {
+					return err
+				}
+				pi.Pack(d)
+				return nil
+			}, func(w *databricks.WorkspaceClient) error {
+				err := w.Metastores.Assign(ctx, create)
+				if err != nil {
+					return err
+				}
+				pi.Pack(d)
+				return nil
+			})
 		},
 		Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			metastoreID, err := NewMetastoreAssignmentAPI(ctx, c).getAssignedMetastoreID()
-			d.Set("metastore_id", metastoreID)
-			return err
+			first, _, err := pi.Unpack(d)
+			if err != nil {
+				return err
+			}
+			workspaceId, err := strconv.ParseInt(first, 10, 64)
+			if err != nil {
+				return err
+			}
+
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				ma, err := acc.MetastoreAssignments.GetByWorkspaceId(ctx, workspaceId)
+				if err != nil {
+					return err
+				}
+				return common.StructToData(ma, s, d)
+			}, func(w *databricks.WorkspaceClient) error {
+				//this only works when managing the metastore assigned to the current workspace.
+				//plus we don't know the workspace we're logged into.
+				ma, err := w.Metastores.Current(ctx)
+				if err != nil {
+					return err
+				}
+				d.Set("metastore_id", ma.MetastoreId)
+				return nil
+			})
 		},
 		Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			var ma MetastoreAssignment
-			common.DataToStructPointer(d, s, &ma)
-			return NewMetastoreAssignmentAPI(ctx, c).updateMetastoreAssignment(ma)
+			workspaceId := int64(d.Get("workspace_id").(int))
+			metastoreId := d.Get("metastore_id").(string)
+			var update catalog.UpdateMetastoreAssignment
+			common.DataToStructPointer(d, s, &update)
+			update.WorkspaceId = workspaceId
+
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				return acc.MetastoreAssignments.Update(ctx,
+					catalog.AccountsUpdateMetastoreAssignment{
+						WorkspaceId:         workspaceId,
+						MetastoreId:         metastoreId,
+						MetastoreAssignment: &update,
+					})
+			}, func(w *databricks.WorkspaceClient) error {
+				return w.Metastores.UpdateAssignment(ctx, update)
+			})
 		},
 		Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			workspaceID, metastoreID, err := pi.Unpack(d)
+			first, metastoreId, err := pi.Unpack(d)
+			if err != nil {
+				return err
+			}
+			workspaceId, err := strconv.ParseInt(first, 10, 64)
 			if err != nil {
 				return err
 			}
-			return NewMetastoreAssignmentAPI(ctx, c).deleteMetastoreAssignment(workspaceID, metastoreID)
+			return c.WorkspaceOrAccountRequest(func(acc *databricks.AccountClient) error {
+				return acc.MetastoreAssignments.DeleteByWorkspaceIdAndMetastoreId(ctx, workspaceId, metastoreId)
+			}, func(w *databricks.WorkspaceClient) error {
+				return w.Metastores.Unassign(ctx, catalog.UnassignRequest{
+					MetastoreId: metastoreId,
+					WorkspaceId: workspaceId,
+				})
+			})
 		},
 	}.ToResource()
 }