[BB-1771]Fix Grants (#33)

Author: FeliLucero1

pkg/connector/groups.go

@@ -189,58 +189,65 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to parse group resource id: %w", err)
 	}
 
-	groupTrait, err := rs.GetGroupTrait(resource)
-	if err != nil {
-		return nil, "", nil, err
-	}
-
 	var workspaceId string
 	isWorkspaceGroup := parentId.ResourceType == workspaceResourceType.Id
 	if isWorkspaceGroup {
 		workspaceId = parentId.Resource
 	}
 
 	// membership grants
-	membersPayload, ok := rs.GetProfileStringValue(groupTrait.Profile, "members")
-	if ok {
-		members := strings.Split(membersPayload, ",")
-
-		for _, m := range members {
-			pp := strings.Split(m, "/")
-			if len(pp) != 2 {
-				return nil, "", nil, fmt.Errorf("databricks-connector: invalid member format of %s: %w", m, err)
-			}
+	// Always fetch the group with members attribute to ensure we get the members
+	// regardless of authentication type (OAuth vs personal access token)
+	group, rateLimitData, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
+	if err != nil {
+		return nil, "", nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
+	}
 
-			memberType, memberID := pp[0], pp[1]
-			var resourceId *v2.ResourceId
-			var anns []protoreflect.ProtoMessage
+	annos := annotations.Annotations{}
+	if rateLimitData != nil {
+		annos.WithRateLimiting(rateLimitData)
+	}
 
-			switch memberType {
-			case "Users":
-				resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: memberID}
-			case "Groups":
-				rid, expandAnnotation, err := groupGrantExpansion(ctx, memberID, parentId)
-				if err != nil {
-					return rv, "", nil, err
-				}
-				resourceId = rid
-				anns = append(anns, expandAnnotation)
-			case "ServicePrincipals":
-				resourceId = &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: memberID}
-			default:
-				return nil, "", nil, fmt.Errorf("databricks-connector: invalid member type: %s", memberType)
-			}
+	for _, member := range group.Members {
+		// member.Ref contains the type and ID separated by "/", e.g., "Users/123" or "Groups/456"
+		pp := strings.Split(member.Ref, "/")
+		if len(pp) != 2 {
+			return nil, "", nil, fmt.Errorf("databricks-connector: invalid member format of %s", member.Ref)
+		}
+
+		memberType, memberID := pp[0], pp[1]
+		var resourceId *v2.ResourceId
+		var anns []protoreflect.ProtoMessage
 
-			rv = append(rv, grant.NewGrant(resource, groupMemberEntitlement, resourceId, grant.WithAnnotation(anns...)))
+		switch memberType {
+		case "Users":
+			resourceId = &v2.ResourceId{ResourceType: userResourceType.Id, Resource: memberID}
+		case "Groups":
+			rid, expandAnnotation, err := groupGrantExpansion(ctx, memberID, parentId)
+			if err != nil {
+				return rv, "", nil, err
+			}
+			resourceId = rid
+			anns = append(anns, expandAnnotation)
+		case "ServicePrincipals":
+			resourceId = &v2.ResourceId{ResourceType: servicePrincipalResourceType.Id, Resource: memberID}
+		default:
+			return nil, "", nil, fmt.Errorf("databricks-connector: invalid member type: %s", memberType)
 		}
+
+		rv = append(rv, grant.NewGrant(resource, groupMemberEntitlement, resourceId, grant.WithAnnotation(anns...)))
 	}
 
 	// role permissions grants
-	ruleSets, _, err := g.client.ListRuleSets(ctx, workspaceId, GroupsType, groupId.Resource)
+	ruleSets, rateLimitDataRuleSets, err := g.client.ListRuleSets(ctx, workspaceId, GroupsType, groupId.Resource)
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("databricks-connector: failed to list role rule sets for group %s: %w", resource.Id.Resource, err)
 	}
 
+	if rateLimitDataRuleSets != nil {
+		annos.WithRateLimiting(rateLimitDataRuleSets)
+	}
+
 	for _, ruleSet := range ruleSets {
 		for _, p := range ruleSet.Principals {
 			resourceId, err := prepareResourceId(ctx, g.client, workspaceId, p)
@@ -261,7 +268,7 @@ func (g *groupBuilder) Grants(ctx context.Context, resource *v2.Resource, pToken
 		}
 	}
 
-	return rv, "", nil, nil
+	return rv, "", annos, nil
 }
 
 func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitlement *v2.Entitlement) (annotations.Annotations, error) {
@@ -306,7 +313,7 @@ func (g *groupBuilder) Grant(ctx context.Context, principal *v2.Resource, entitl
 	membershipEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupMemberEntitlement)
 	managerEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupManagerEntitlement)
 	if entitlement.Id == membershipEntitlementID {
-		group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource)
+		group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
 		}
@@ -449,7 +456,7 @@ func (g *groupBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations
 	membershipEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupMemberEntitlement)
 	managerEntitlementID := ent.NewEntitlementID(entitlement.Resource, groupManagerEntitlement)
 	if entitlement.Id == membershipEntitlementID {
-		group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource)
+		group, _, err := g.client.GetGroup(ctx, workspaceId, groupId.Resource, databricks.NewGroupAttrVars())
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to get group %s: %w", groupId.Resource, err)
 		}

pkg/connector/roles.go

@@ -347,7 +347,7 @@ func (r *roleBuilder) Grant(ctx context.Context, principal *v2.Resource, entitle
 		}
 
 	case groupResourceType.Id:
-		g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource)
+		g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource, databricks.NewGroupRolesAttrVars())
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to get group: %w", err)
 		}
@@ -428,7 +428,7 @@ func (r *roleBuilder) Revoke(ctx context.Context, grant *v2.Grant) (annotations.
 		}
 
 	case groupResourceType.Id:
-		g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource)
+		g, _, err := r.client.GetGroup(ctx, workspaceId, principal.Id.Resource, databricks.NewGroupRolesAttrVars())
 		if err != nil {
 			return nil, fmt.Errorf("databricks-connector: failed to get group: %w", err)
 		}

pkg/databricks/client.go

@@ -254,7 +254,7 @@ func (c *Client) ListGroups(
 	return res.Resources, res.Total, ratelimitData, nil
 }
 
-func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string) (
+func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string, vars ...Vars) (
 	*Group,
 	*v2.RateLimitDescription,
 	error,
@@ -267,7 +267,7 @@ func (c *Client) GetGroup(ctx context.Context, workspaceId, groupId string) (
 	}
 
 	var res *Group
-	ratelimitData, err := c.Get(ctx, u, &res)
+	ratelimitData, err := c.Get(ctx, u, &res, vars...)
 	if err != nil {
 		return nil, ratelimitData, err
 	}