[Internal] refactored databricks_mws_permission_assignment to Go SDK (#3831)

nkvuong · web-flow · commit a59abf175757 · 2024-08-14T10:20:25.000Z
## Changes
- refactored `databricks_mws_permission_assignment` to Go SDK
- added `TestUcAccAssignGroupToWorkspace` to flaky tests

## Tests
&lt;!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
--&gt;

- [x] `make test` run locally
- [x] covered with integration tests in `internal/acceptance`
- [ ] relevant acceptance tests are passing
- [x] using Go SDK
diff --git a/common/util.go b/common/util.go
@@ -49,6 +49,11 @@ func MustInt64(s string) int64 {
 	return n
 }
 
+// GetInt64 returns the data for the given key and cast it to int64
+func GetInt64(d *schema.ResourceData, key string) int64 {
+	return int64(d.Get(key).(int))
+}
+
 // Reads the file content from a given path
 func ReadFileContent(source string) ([]byte, error) {
 	log.Printf("[INFO] Reading %s", source)
diff --git a/mws/resource_mws_permission_assignment.go b/mws/resource_mws_permission_assignment.go
@@ -2,131 +2,90 @@ package mws
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
 	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/terraform-provider-databricks/common"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
-func NewPermissionAssignmentAPI(ctx context.Context, m any) PermissionAssignmentAPI {
-	return PermissionAssignmentAPI{m.(*common.DatabricksClient), ctx}
-}
-
-type PermissionAssignmentAPI struct {
-	client  *common.DatabricksClient
-	context context.Context
-}
-
-type Permissions struct {
-	Permissions []string `json:"permissions"`
-}
-
-func (a PermissionAssignmentAPI) CreateOrUpdate(workspaceId, principalId int64, r Permissions) error {
-	if a.client.Config.AccountID == "" {
-		return errors.New("must have `account_id` on provider")
-	}
-	path := fmt.Sprintf(
-		"/accounts/%s/workspaces/%d/permissionassignments/principals/%d",
-		a.client.Config.AccountID, workspaceId, principalId)
-	return a.client.Put(a.context, path, r)
-}
-
-func (a PermissionAssignmentAPI) Remove(workspaceId, principalId string) error {
-	if a.client.Config.AccountID == "" {
-		return errors.New("must have `account_id` on provider")
-	}
-	path := fmt.Sprintf(
-		"/accounts/%s/workspaces/%s/permissionassignments/principals/%s",
-		a.client.Config.AccountID, workspaceId, principalId)
-	return a.client.Delete(a.context, path, nil)
-}
-
-type Principal struct {
-	DisplayName          string `json:"display_name"`
-	PrincipalID          int64  `json:"principal_id"`
-	ServicePrincipalName string `json:"service_principal_name,omitempty"`
-	UserName             string `json:"user_name,omitempty"`
-	GroupName            string `json:"group_name,omitempty"`
-}
-
-type PermissionAssignment struct {
-	Permissions []string `json:"permissions"`
-	Principal   Principal
-}
-
-type PermissionAssignmentList struct {
-	PermissionAssignments []PermissionAssignment `json:"permission_assignments"`
-}
-
-func (l PermissionAssignmentList) ForPrincipal(principalId int64) (res Permissions, err error) {
-	for _, v := range l.PermissionAssignments {
-		if v.Principal.PrincipalID != principalId {
+func getPermissionsByPrincipal(list iam.PermissionAssignments, principalId int64) (res iam.UpdateWorkspaceAssignments, err error) {
+	for _, v := range list.PermissionAssignments {
+		if v.Principal.PrincipalId != principalId {
 			continue
 		}
-		return Permissions{v.Permissions}, nil
+		return iam.UpdateWorkspaceAssignments{Permissions: v.Permissions}, nil
 	}
 	return res, apierr.NotFound(fmt.Sprintf("%d not found", principalId))
 }
 
-func (a PermissionAssignmentAPI) List(workspaceId int64) (list PermissionAssignmentList, err error) {
-	if a.client.Config.AccountID == "" {
-		return list, errors.New("must have `account_id` on provider")
-	}
-	path := fmt.Sprintf("/accounts/%s/workspaces/%d/permissionassignments",
-		a.client.Config.AccountID, workspaceId)
-	err = a.client.Get(a.context, path, nil, &list)
-	return
-}
-
 func ResourceMwsPermissionAssignment() common.Resource {
-	type entity struct {
-		WorkspaceId int64    `json:"workspace_id"`
-		PrincipalId int64    `json:"principal_id"`
-		Permissions []string `json:"permissions" tf:"slice_as_set"`
-	}
-	s := common.StructToSchema(entity{},
-		common.NoCustomize)
+	s := common.StructToSchema(iam.UpdateWorkspaceAssignments{},
+		func(m map[string]*schema.Schema) map[string]*schema.Schema {
+			common.CustomizeSchemaPath(m).AddNewField("workspace_id", &schema.Schema{
+				Type:     schema.TypeInt,
+				Required: true,
+			}).AddNewField("principal_id", &schema.Schema{
+				Type:     schema.TypeInt,
+				Required: true,
+			})
+			common.CustomizeSchemaPath(m, "permissions").SetRequired().SetSliceSet()
+			return m
+		})
 	pair := common.NewPairID("workspace_id", "principal_id").Schema(
 		func(m map[string]*schema.Schema) map[string]*schema.Schema {
 			return s
 		})
 	return common.Resource{
 		Schema: s,
 		Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
-			var assignment entity
+			acc, err := c.AccountClient()
+			if err != nil {
+				return err
+			}
+			var assignment iam.UpdateWorkspaceAssignments
 			common.DataToStructPointer(d, s, &assignment)
-			api := NewPermissionAssignmentAPI(ctx, c)
-			err := api.CreateOrUpdate(assignment.WorkspaceId, assignment.PrincipalId,
-				Permissions{assignment.Permissions})
+			assignment.PrincipalId = common.GetInt64(d, "principal_id")
+			assignment.WorkspaceId = common.GetInt64(d, "workspace_id")
+			_, err = acc.WorkspaceAssignment.Update(ctx, assignment)
 			if err != nil {
 				return err
 			}
 			pair.Pack(d)
 			return nil
 		},
 		Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			acc, err := c.AccountClient()
+			if err != nil {
+				return err
+			}
 			workspaceId, principalId, err := pair.Unpack(d)
 			if err != nil {
 				return fmt.Errorf("parse id: %w", err)
 			}
-			list, err := NewPermissionAssignmentAPI(ctx, c).List(common.MustInt64(workspaceId))
+			list, err := acc.WorkspaceAssignment.ListByWorkspaceId(ctx, common.MustInt64(workspaceId))
 			if err != nil {
 				return err
 			}
-			permissions, err := list.ForPrincipal(common.MustInt64(principalId))
+			permissions, err := getPermissionsByPrincipal(*list, common.MustInt64(principalId))
 			if err != nil {
 				return err
 			}
+			d.Set("workspace_id", common.MustInt64(workspaceId))
+			d.Set("principal_id", common.MustInt64(principalId))
 			return common.StructToData(permissions, s, d)
 		},
 		Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
+			acc, err := c.AccountClient()
+			if err != nil {
+				return err
+			}
 			workspaceId, principalId, err := pair.Unpack(d)
 			if err != nil {
 				return fmt.Errorf("parse id: %w", err)
 			}
-			return NewPermissionAssignmentAPI(ctx, c).Remove(workspaceId, principalId)
+			return acc.WorkspaceAssignment.DeleteByWorkspaceIdAndPrincipalId(ctx, common.MustInt64(workspaceId), common.MustInt64(principalId))
 		},
 	}
 }
diff --git a/mws/resource_mws_permission_assignment_test.go b/mws/resource_mws_permission_assignment_test.go
@@ -3,33 +3,36 @@ package mws
 import (
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/terraform-provider-databricks/qa"
+	"github.com/stretchr/testify/mock"
 )
 
 func TestPermissionAssignmentCreate(t *testing.T) {
 	qa.ResourceFixture{
-		Fixtures: []qa.HTTPFixture{
-			{
-				Method:   "PUT",
-				Resource: "/api/2.0/accounts/abc/workspaces/123/permissionassignments/principals/345",
-				ExpectedRequest: Permissions{
-					Permissions: []string{"USER"},
+		MockAccountClientFunc: func(m *mocks.MockAccountClient) {
+			e := m.GetMockWorkspaceAssignmentAPI().EXPECT()
+			e.Update(mock.Anything, iam.UpdateWorkspaceAssignments{
+				Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionUser},
+				PrincipalId: 345,
+				WorkspaceId: 123,
+			}).Return(&iam.PermissionAssignment{
+				Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionUser},
+				Principal: &iam.PrincipalOutput{
+					PrincipalId: 345,
 				},
-			},
-			{
-				Method:   "GET",
-				Resource: "/api/2.0/accounts/abc/workspaces/123/permissionassignments",
-				Response: PermissionAssignmentList{
-					PermissionAssignments: []PermissionAssignment{
-						{
-							Permissions: []string{"USER"},
-							Principal: Principal{
-								PrincipalID: 345,
-							},
+			}, nil)
+			e.ListByWorkspaceId(mock.Anything, int64(123)).Return(&iam.PermissionAssignments{
+				PermissionAssignments: []iam.PermissionAssignment{
+					{
+						Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionUser},
+						Principal: &iam.PrincipalOutput{
+							PrincipalId: 345,
 						},
 					},
 				},
-			},
+			}, nil)
 		},
 		Resource:  ResourceMwsPermissionAssignment(),
 		Create:    true,
@@ -42,23 +45,53 @@ func TestPermissionAssignmentCreate(t *testing.T) {
 	}.ApplyNoError(t)
 }
 
+func TestPermissionAssignmentRead(t *testing.T) {
+	qa.ResourceFixture{
+		MockAccountClientFunc: func(m *mocks.MockAccountClient) {
+			e := m.GetMockWorkspaceAssignmentAPI().EXPECT()
+			e.ListByWorkspaceId(mock.Anything, int64(123)).Return(&iam.PermissionAssignments{
+				PermissionAssignments: []iam.PermissionAssignment{
+					{
+						Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionUser},
+						Principal: &iam.PrincipalOutput{
+							PrincipalId: 345,
+						},
+					},
+					{
+						Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionUser},
+						Principal: &iam.PrincipalOutput{
+							PrincipalId: 456,
+						},
+					},
+				},
+			}, nil)
+		},
+		Resource:  ResourceMwsPermissionAssignment(),
+		Read:      true,
+		New:       true,
+		AccountID: "abc",
+		ID:        "123|456",
+	}.ApplyAndExpectData(t, map[string]any{
+		"workspace_id": 123,
+		"principal_id": 456,
+		"permissions":  []string{"USER"},
+	})
+}
+
 func TestPermissionAssignmentReadNotFound(t *testing.T) {
 	qa.ResourceFixture{
-		Fixtures: []qa.HTTPFixture{
-			{
-				Method:   "GET",
-				Resource: "/api/2.0/accounts/abc/workspaces/123/permissionassignments",
-				Response: PermissionAssignmentList{
-					PermissionAssignments: []PermissionAssignment{
-						{
-							Permissions: []string{"USER"},
-							Principal: Principal{
-								PrincipalID: 345,
-							},
+		MockAccountClientFunc: func(m *mocks.MockAccountClient) {
+			e := m.GetMockWorkspaceAssignmentAPI().EXPECT()
+			e.ListByWorkspaceId(mock.Anything, int64(123)).Return(&iam.PermissionAssignments{
+				PermissionAssignments: []iam.PermissionAssignment{
+					{
+						Permissions: []iam.WorkspacePermission{iam.WorkspacePermissionUser},
+						Principal: &iam.PrincipalOutput{
+							PrincipalId: 345,
 						},
 					},
 				},
-			},
+			}, nil)
 		},
 		Resource:  ResourceMwsPermissionAssignment(),
 		Read:      true,
@@ -70,11 +103,9 @@ func TestPermissionAssignmentReadNotFound(t *testing.T) {
 
 func TestPermissionAssignmentDelete(t *testing.T) {
 	qa.ResourceFixture{
-		Fixtures: []qa.HTTPFixture{
-			{
-				Method:   "DELETE",
-				Resource: "/api/2.0/accounts/abc/workspaces/123/permissionassignments/principals/456",
-			},
+		MockAccountClientFunc: func(m *mocks.MockAccountClient) {
+			e := m.GetMockWorkspaceAssignmentAPI().EXPECT()
+			e.DeleteByWorkspaceIdAndPrincipalId(mock.Anything, int64(123), int64(456)).Return(nil)
 		},
 		Resource:  ResourceMwsPermissionAssignment(),
 		Delete:    true,
@@ -86,7 +117,7 @@ func TestPermissionAssignmentDelete(t *testing.T) {
 func TestPermissionAssignmentFuzz_NoAccountID(t *testing.T) {
 	qa.ResourceCornerCases(t, ResourceMwsPermissionAssignment(),
 		qa.CornerCaseID("123|456"),
-		qa.CornerCaseExpectError("must have `account_id` on provider"))
+		qa.CornerCaseExpectError("invalid Databricks Account configuration"))
 }
 
 func TestPermissionAssignmentFuzz_InvalidID(t *testing.T) {
diff --git a/qa/testing.go b/qa/testing.go
@@ -458,7 +458,7 @@ func ResourceCornerCases(t *testing.T, resource common.Resource, cc ...CornerCas
 	}
 	HTTPFixturesApply(t, HTTPFailures, func(ctx context.Context, client *common.DatabricksClient) {
 		validData := r.TestResourceData()
-		client.Config.AccountID = config["account_id"]
+		client.Config.WithTesting().AccountID = config["account_id"]
 		for n, v := range m {
 			if v == nil {
 				continue
diff --git a/test-config.yaml b/test-config.yaml
@@ -12,3 +12,6 @@ ignored_tests:
     - package: github.com/databricks/terraform-provider-databricks/internal/acceptance
       test_name: TestMwsAccServicePrincipalResourceOnAws
       comment: Failures due to read-after-write inconsistency. Tracked in https://databricks.atlassian.net/browse/ES-1100061
+    - package: github.com/databricks/terraform-provider-databricks/internal/acceptance
+      test_name: TestUcAccAssignGroupToWorkspace
+      comment: Failures due to read-after-write inconsistency. Tracked in https://databricks.atlassian.net/browse/ES-1100061

Original file line number	Diff line number	Diff line change
`@@ -458,7 +458,7 @@ func ResourceCornerCases(t *testing.T, resource common.Resource, cc ...CornerCas`
`458`	`458`	`}`
`459`	`459`	`HTTPFixturesApply(t, HTTPFailures, func(ctx context.Context, client *common.DatabricksClient) {`
`460`	`460`	`validData := r.TestResourceData()`
`461`		`- client.Config.AccountID = config["account_id"]`
	`461`	`+ client.Config.WithTesting().AccountID = config["account_id"]`
`462`	`462`	`for n, v := range m {`
`463`	`463`	`if v == nil {`
`464`	`464`	`continue`