diff --git a/pkg/connector/enterprise_roles.go b/pkg/connector/enterprise_roles.go index 9bef1fc2..98878b81 100644 --- a/pkg/connector/enterprise_roles.go +++ b/pkg/connector/enterprise_roles.go @@ -10,6 +10,7 @@ import ( ent "github.com/conductorone/baton-sdk/pkg/types/entitlement" "github.com/conductorone/baton-sdk/pkg/types/grant" resources "github.com/conductorone/baton-sdk/pkg/types/resource" + enterprise "github.com/conductorone/baton-slack/pkg/slack" ) diff --git a/pkg/connector/helpers.go b/pkg/connector/helpers.go index d8eede43..a0af8654 100644 --- a/pkg/connector/helpers.go +++ b/pkg/connector/helpers.go @@ -9,6 +9,8 @@ import ( "github.com/conductorone/baton-sdk/pkg/pagination" "github.com/slack-go/slack" "google.golang.org/protobuf/types/known/timestamppb" + + enterprise "github.com/conductorone/baton-slack/pkg/slack" ) func parsePageToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, error) { @@ -31,14 +33,24 @@ func parsePageToken(i string, resourceID *v2.ResourceId) (*pagination.Bag, error func annotationsForError(err error) (annotations.Annotations, error) { annos := annotations.Annotations{} var rateLimitErr *slack.RateLimitedError - if !errors.As(err, &rateLimitErr) { - return annos, err + if errors.As(err, &rateLimitErr) { + annos.WithRateLimiting(&v2.RateLimitDescription{ + Limit: 0, + Remaining: 0, + ResetAt: timestamppb.New(time.Now().Add(rateLimitErr.RetryAfter)), + }) + return annos, nil + } + + var enterpriseRateLimitErr *enterprise.RateLimitError + if errors.As(err, &enterpriseRateLimitErr) { + annos.WithRateLimiting(&v2.RateLimitDescription{ + Limit: 0, + Remaining: 0, + ResetAt: timestamppb.New(time.Now().Add(enterpriseRateLimitErr.RetryAfter)), + }) + return annos, nil } - annos.WithRateLimiting(&v2.RateLimitDescription{ - Limit: 0, - Remaining: 0, - ResetAt: timestamppb.New(time.Now().Add(rateLimitErr.RetryAfter)), - }) - return annos, nil + return annos, err } diff --git a/pkg/slack/client.go b/pkg/slack/client.go index d9c1c502..03ec4775 100644 --- a/pkg/slack/client.go +++ b/pkg/slack/client.go @@ -6,7 +6,9 @@ import ( "fmt" "net/http" "net/url" + "strconv" "strings" + "time" "github.com/slack-go/slack" ) @@ -218,9 +220,8 @@ func (c *Client) GetUserGroups(ctx context.Context, teamID string) ([]slack.User } var res struct { - Ok bool `json:"ok"` + BaseResponse UserGroups []slack.UserGroup `json:"usergroups"` - Error string `json:"error"` } err = c.doRequest(ctx, userGroupsUrl, &res, http.MethodPost, nil, values) @@ -235,6 +236,14 @@ func (c *Client) GetUserGroups(ctx context.Context, teamID string) ([]slack.User return res.UserGroups, nil } +type RateLimitError struct { + RetryAfter time.Duration +} + +func (r *RateLimitError) Error() string { + return fmt.Sprintf("rate limited, retry after: %s", r.RetryAfter.String()) +} + func (c *Client) doRequest(ctx context.Context, url string, res interface{}, method string, payload []byte, values url.Values) error { reqBody := strings.NewReader(values.Encode()) @@ -256,5 +265,14 @@ func (c *Client) doRequest(ctx context.Context, url string, res interface{}, met return err } + if resp.StatusCode == http.StatusTooManyRequests { + retryAfter := resp.Header.Get("Retry-After") + retryAfterSec, err := strconv.Atoi(retryAfter) + if err != nil { + return fmt.Errorf("error parsing retry after header: %w", err) + } + return &RateLimitError{RetryAfter: time.Second * time.Duration(retryAfterSec)} + } + return nil }