Add support for workspace ID in prepareResourceId

ConductorOne · Dec 20, 2024 · 393a768 · 393a768
1 parent b27117b
commit 393a768
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 11 deletions.
diff --git a/pkg/connector/account.go b/pkg/connector/account.go
@@ -125,7 +125,7 @@ func (a *accountBuilder) Grants(ctx context.Context, resource *v2.Resource, pTok
 		// rule set contains role and its principals, each one with resource type and resource id seperated by "/"
 		if strings.Contains(ruleSet.Role, MarketplaceAdminRole) {
 			for _, p := range ruleSet.Principals {
-				resourceId, err := prepareResourceID(ctx, a.client, p)
+				resourceId, err := prepareResourceId(ctx, a.client, "", p)
 				if err != nil {
 					return nil, "", nil, fmt.Errorf("databricks-connector: failed to prepare resource id for principal %s: %w", p, err)
 				}

diff --git a/pkg/connector/groups.go b/pkg/connector/groups.go
@@ -233,7 +233,7 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken
 
 	for _, ruleSet := range ruleSets {
 		for _, p := range ruleSet.Principals {
-			resourceId, err := prepareResourceID(ctx, g.client, p)
+			resourceId, err := prepareResourceId(ctx, g.client, workspaceId, p)
 			if err != nil {
 				return nil, "", nil, fmt.Errorf("databricks-connector: failed to prepare resource id for principal %s: %w", p, err)
 			}

diff --git a/pkg/connector/helpers.go b/pkg/connector/helpers.go
@@ -62,9 +62,9 @@ func annotationsForUserResourceType() annotations.Annotations {
 	return annos
 }
 
-// prepareResourceID prepares a resource ID for a user, group, or service principal.
+// prepareResourceId prepares a resource ID for a user, group, or service principal.
 // It's used when we need to parse results from listing rule sets.
-func prepareResourceID(ctx context.Context, c *databricks.Client, principal string) (*v2.ResourceId, error) {
+func prepareResourceId(ctx context.Context, c *databricks.Client, workspaceId string, principal string) (*v2.ResourceId, error) {
 	pp := strings.Split(principal, "/")
 	if len(pp) != 2 {
 		return nil, fmt.Errorf("invalid principal format: %s", principal)
@@ -77,21 +77,21 @@ func prepareResourceID(ctx context.Context, c *databricks.Client, principal stri
 
 	switch principalType {
 	case UsersType:
-		userID, _, err := c.FindUserID(ctx, principal)
+		userID, _, err := c.FindUserID(ctx, workspaceId, principal)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find user %s: %w", principal, err)
 		}
 
 		resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: userID}
 	case GroupsType:
-		groupID, _, err := c.FindGroupID(ctx, principal)
+		groupID, _, err := c.FindGroupID(ctx, workspaceId, principal)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find group %s: %w", principal, err)
 		}
 
 		resourceId = &v2.ResourceId{ResourceType: groupResourceType.Id, Resource: groupID}
 	case ServicePrincipalsType:
-		servicePrincipalID, _, err := c.FindServicePrincipalID(ctx, principal)
+		servicePrincipalID, _, err := c.FindServicePrincipalID(ctx, workspaceId, principal)
 		if err != nil {
 			return nil, fmt.Errorf("failed to find service principal %s: %w", principal, err)
 		}

diff --git a/pkg/connector/service-principals.go b/pkg/connector/service-principals.go
@@ -182,7 +182,7 @@ func (s *servicePrincipalBuilder) Grants(ctx context.Context, resource *v2.Resou
 	var rv []*v2.Grant
 	for _, ruleSet := range ruleSets {
 		for _, p := range ruleSet.Principals {
-			resourceId, err := prepareResourceID(ctx, s.client, p)
+			resourceId, err := prepareResourceId(ctx, s.client, workspaceId, p)
 			if err != nil {
 				return nil, "", nil, fmt.Errorf("databricks-connector: failed to prepare resource id for principal %s: %w", p, err)
 			}

diff --git a/pkg/databricks/client.go b/pkg/databricks/client.go
@@ -171,6 +171,7 @@ func (c *Client) UpdateUser(ctx context.Context, workspaceId string, user *User)
 
 func (c *Client) FindUserID(
 	ctx context.Context,
+	workspaceId string,
 	username string,
 ) (
 	string,
@@ -179,7 +180,7 @@ func (c *Client) FindUserID(
 ) {
 	users, _, ratelimitData, err := c.ListUsers(
 		ctx,
-		"",
+		workspaceId,
 		&PaginationVars{Count: 1},
 		NewFilterVars(fmt.Sprintf("userName eq '%s'", username)),
 	)
@@ -282,6 +283,7 @@ func (c *Client) UpdateGroup(ctx context.Context, workspaceId string, group *Gro
 
 func (c *Client) FindGroupID(
 	ctx context.Context,
+	workspaceId string,
 	displayName string,
 ) (
 	string,
@@ -290,7 +292,7 @@ func (c *Client) FindGroupID(
 ) {
 	groups, _, ratelimitData, err := c.ListGroups(
 		ctx,
-		"",
+		workspaceId,
 		&PaginationVars{Count: 1},
 		NewFilterVars(fmt.Sprintf("displayName eq '%s'", displayName)),
 	)
@@ -399,6 +401,7 @@ func (c *Client) UpdateServicePrincipal(
 
 func (c *Client) FindServicePrincipalID(
 	ctx context.Context,
+	workspaceId string,
 	appID string,
 ) (
 	string,
@@ -407,7 +410,7 @@ func (c *Client) FindServicePrincipalID(
 ) {
 	servicePrincipals, _, ratelimitData, err := c.ListServicePrincipals(
 		ctx,
-		"",
+		workspaceId,
 		&PaginationVars{Count: 1},
 		NewFilterVars(fmt.Sprintf("applicationId eq '%s'", appID)),
 	)