diff --git a/pkg/connector/client/helpers.go b/pkg/connector/client/helpers.go index debb5891..1dc28a0d 100644 --- a/pkg/connector/client/helpers.go +++ b/pkg/connector/client/helpers.go @@ -186,6 +186,27 @@ func MapSlackErrorToGRPCCode(slackError string) codes.Code { return codes.Unknown } +// IsRateLimited checks whether the annotations contain a RateLimitDescription +// with STATUS_OVERLIMIT. +func IsRateLimited(annos *annotations.Annotations) bool { + if annos == nil { + return false + } + rl := &v2.RateLimitDescription{} + ok, err := annos.Pick(rl) + if err != nil || !ok { + return false + } + return rl.Status == v2.RateLimitDescription_STATUS_OVERLIMIT +} + +// RateLimitOverride returns a RateLimitDescription with a 60s wait. +// Use this for endpoints where Slack's Retry-After header is too short +// to allow meaningful progress (e.g. users.info in a per-member loop). +func RateLimitOverride() *v2.RateLimitDescription { + return rateLimitDescription(60 * time.Second) +} + func rateLimitDescription(retryAfter time.Duration) *v2.RateLimitDescription { return &v2.RateLimitDescription{ Status: v2.RateLimitDescription_STATUS_OVERLIMIT, diff --git a/pkg/connector/client/helpers_test.go b/pkg/connector/client/helpers_test.go new file mode 100644 index 00000000..c71924b6 --- /dev/null +++ b/pkg/connector/client/helpers_test.go @@ -0,0 +1,108 @@ +package client + +import ( + "testing" + "time" + + v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" + "github.com/conductorone/baton-sdk/pkg/annotations" + "github.com/slack-go/slack" +) + +func TestIsRateLimited(t *testing.T) { + t.Run("nil annotations", func(t *testing.T) { + if IsRateLimited(nil) { + t.Error("expected false for nil annotations") + } + }) + + t.Run("empty annotations", func(t *testing.T) { + var annos annotations.Annotations + if IsRateLimited(&annos) { + t.Error("expected false for empty annotations") + } + }) + + t.Run("overlimit annotation", func(t *testing.T) { + var annos annotations.Annotations + annos.WithRateLimiting(rateLimitDescription(4 * time.Second)) + if !IsRateLimited(&annos) { + t.Error("expected true for overlimit annotation") + } + }) + + t.Run("ok status annotation", func(t *testing.T) { + var annos annotations.Annotations + annos.WithRateLimiting(&v2.RateLimitDescription{ + Status: v2.RateLimitDescription_STATUS_OK, + Remaining: 50, + }) + if IsRateLimited(&annos) { + t.Error("expected false for OK status annotation") + } + }) +} + +func TestWrapErrorSetsRateLimitAnnotation(t *testing.T) { + t.Run("rate limit error populates annotations", func(t *testing.T) { + var annos annotations.Annotations + err := &slack.RateLimitedError{RetryAfter: 4 * time.Second} + wrappedErr := WrapError(err, "test", &annos) + if wrappedErr == nil { + t.Fatal("expected non-nil error") + } + if !IsRateLimited(&annos) { + t.Error("expected rate limit annotation after WrapError with RateLimitedError") + } + }) + + t.Run("non-rate-limit error does not populate rate limit annotation", func(t *testing.T) { + var annos annotations.Annotations + err := slack.SlackErrorResponse{Err: "user_not_found"} + wrappedErr := WrapError(err, "test", &annos) + if wrappedErr == nil { + t.Fatal("expected non-nil error") + } + if IsRateLimited(&annos) { + t.Error("expected no rate limit annotation for user_not_found error") + } + }) +} + +func TestRateLimitOverride(t *testing.T) { + rl := RateLimitOverride() + if rl.Status != v2.RateLimitDescription_STATUS_OVERLIMIT { + t.Errorf("expected STATUS_OVERLIMIT, got %v", rl.Status) + } + if rl.Remaining != 0 { + t.Errorf("expected Remaining=0, got %d", rl.Remaining) + } + resetIn := time.Until(rl.ResetAt.AsTime()) + if resetIn < 55*time.Second || resetIn > 65*time.Second { + t.Errorf("expected ResetAt ~60s from now, got %v", resetIn) + } +} + +func TestRateLimitOverrideReplacesExisting(t *testing.T) { + // Simulate: WrapError sets a 4s annotation, then we override to 60s + var annos annotations.Annotations + annos.WithRateLimiting(rateLimitDescription(4 * time.Second)) + + if !IsRateLimited(&annos) { + t.Fatal("expected rate limited after initial annotation") + } + + // Override + annos.WithRateLimiting(RateLimitOverride()) + + rl := &v2.RateLimitDescription{} + ok, err := annos.Pick(rl) + if err != nil || !ok { + t.Fatal("expected to find rate limit annotation after override") + } + + resetIn := time.Until(rl.ResetAt.AsTime()) + if resetIn < 55*time.Second { + t.Errorf("expected override to set ResetAt ~60s from now, got %v", resetIn) + } +} diff --git a/pkg/connector/user.go b/pkg/connector/user.go index b292fad1..95861269 100644 --- a/pkg/connector/user.go +++ b/pkg/connector/user.go @@ -27,7 +27,7 @@ func (o *userResourceType) scimUserResource(ctx context.Context, scimUser client // NOTE: this is mainly to maintain compatibility with existing profile in non scim flow. slackUser, err := o.client.GetUserInfoContext(ctx, scimUser.ID) if err != nil { - return nil, client.WrapError(err, fmt.Sprintf("fetching user info for SCIM user %s", scimUser.ID), nil) + return nil, err } profile := make(map[string]interface{}) @@ -244,7 +244,11 @@ func (o *userResourceType) listScimAPI(ctx context.Context, parentResourceID *v2 for _, user := range response.Resources { userResource, err := o.scimUserResource(ctx, user, parentResourceID) if err != nil { - return nil, &resource.SyncOpResults{Annotations: annos}, err + wrappedErr := client.WrapError(err, fmt.Sprintf("fetching user info for SCIM user %s", user.ID), &annos) + if client.IsRateLimited(&annos) { + annos.WithRateLimiting(client.RateLimitOverride()) + } + return nil, &resource.SyncOpResults{Annotations: annos}, wrappedErr } rv = append(rv, userResource) } diff --git a/pkg/connector/user_group.go b/pkg/connector/user_group.go index 5c84d6cc..eb1b8630 100644 --- a/pkg/connector/user_group.go +++ b/pkg/connector/user_group.go @@ -166,7 +166,11 @@ func (o *userGroupResourceType) Grants( for _, member := range page { user, err := o.client.GetUserInfoContext(ctx, member) if err != nil { - return nil, &resource.SyncOpResults{Annotations: outputAnnotations}, client.WrapError(err, fmt.Sprintf("fetching user info for member %s", member), &outputAnnotations) + wrappedErr := client.WrapError(err, fmt.Sprintf("fetching user info for member %s", member), &outputAnnotations) + if client.IsRateLimited(&outputAnnotations) { + outputAnnotations.WithRateLimiting(client.RateLimitOverride()) + } + return nil, &resource.SyncOpResults{Annotations: outputAnnotations}, wrappedErr } ur, err := userResource(ctx, user, res.Id) if err != nil {