Merge branch 'main' of https://github.com/ConductorOne/baton-slack in…

…to jakcinmarina/provisioning
ConductorOne · Mar 7, 2024 · 9cc3a85 · 9cc3a85
2 parents 8d154f9 + 19c69c8
commit 9cc3a85
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 8 deletions.
diff --git 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
@@ -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
@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/slack-go/slack"
 )
@@ -282,6 +284,14 @@ func (c *Client) SetWorkspaceRole(ctx context.Context, teamID, userID, roleID st
 	return 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())
 
@@ -303,5 +313,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
 }