[INC-300] baton-github: optimize Grants() for org Resource

Author: Bencheng21

pkg/connector/org.go

@@ -182,78 +182,49 @@ func (o *orgResourceType) Grants(
 	resource *v2.Resource,
 	pToken *pagination.Token,
 ) ([]*v2.Grant, string, annotations.Annotations, error) {
-	bag, page, err := parsePageToken(pToken.Token, resource.Id)
-	if err != nil {
-		return nil, "", nil, err
-	}
-
-	opts := github.ListMembersOptions{
-		ListOptions: github.ListOptions{
-			Page:    page,
-			PerPage: maxPageSize,
-		},
-	}
-
 	orgName, err := o.orgCache.GetOrgName(ctx, resource.Id)
 	if err != nil {
 		return nil, "", nil, err
 	}
 
-	users, resp, err := o.client.Organizations.ListMembers(ctx, orgName, &opts)
-	if err != nil {
-		if isNotFoundError(resp) {
-			return nil, "", nil, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf("org: %s not found", orgName))
-		}
-		errMsg := "github-connectorv2: failed to list org members"
-		if isRatelimited(resp) {
-			return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
-		}
-		return nil, "", nil, fmt.Errorf("%s: %w", errMsg, err)
-	}
-
-	nextPage, reqAnnos, err := parseResp(resp)
-	if err != nil {
-		return nil, "", nil, fmt.Errorf("github-connectorv2: failed to parse response: %w", err)
-	}
-
-	pageToken, err := bag.NextToken(nextPage)
-	if err != nil {
-		return nil, "", nil, err
-	}
-
-	var rv []*v2.Grant
-	for _, user := range users {
-		membership, _, err := o.client.Organizations.GetOrgMembership(ctx, user.GetLogin(), orgName)
-		if err != nil {
-			return nil, "", nil, fmt.Errorf("github-connectorv2: failed to get org memberships for user: %w", err)
-		}
-		if membership.GetState() == "pending" {
-			continue
-		}
-
-		ur, err := userResource(ctx, user, user.GetEmail(), nil)
-		if err != nil {
-			return nil, "", nil, err
+	rv := []*v2.Grant{}
+	for _, role := range []string{orgRoleMember, orgRoleAdmin} {
+		opts := github.ListMembersOptions{
+			Role: role,
+			ListOptions: github.ListOptions{
+				Page:    0,
+				PerPage: maxPageSize,
+			},
 		}
 
-		roleName := strings.ToLower(membership.GetRole())
-		switch roleName {
-		case orgRoleAdmin:
-			rv = append(rv, o.orgRoleGrant(orgRoleAdmin, resource, ur.Id, user.GetID()))
-			rv = append(rv, o.orgRoleGrant(orgRoleMember, resource, ur.Id, user.GetID()))
+		for {
+			users, resp, err := o.client.Organizations.ListMembers(ctx, orgName, &opts)
+			if err != nil {
+				if isNotFoundError(resp) {
+					return nil, "", nil, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf("org: %s not found", orgName))
+				}
+				errMsg := "github-connectorv2: failed to list org members"
+				if isRatelimited(resp) {
+					return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+				}
+				return nil, "", nil, fmt.Errorf("%s: %w", errMsg, err)
+			}
 
-		case orgRoleMember:
-			rv = append(rv, o.orgRoleGrant(orgRoleMember, resource, ur.Id, user.GetID()))
+			for _, user := range users {
+				ur, err := userResource(ctx, user, user.GetEmail(), nil)
+				if err != nil {
+					return nil, "", nil, err
+				}
+				rv = append(rv, o.orgRoleGrant(role, resource, ur.Id, user.GetID()))
+			}
 
-		default:
-			ctxzap.Extract(ctx).Warn("Unknown GitHub Role Name",
-				zap.String("role_name", roleName),
-				zap.String("github_username", user.GetLogin()),
-			)
+			if resp.NextPage == 0 {
+				break
+			}
+			opts.ListOptions.Page = resp.NextPage
 		}
 	}
-
-	return rv, pageToken, reqAnnos, nil
+	return rv, "", nil, nil
 }
 
 func (o *orgResourceType) Grant(ctx context.Context, principal *v2.Resource, en *v2.Entitlement) (annotations.Annotations, error) {