Merge pull request #95 from ConductorOne/BB1134

[BB-1134] githbu: parse rate limit from response

Author: Bencheng21

pkg/connector/helpers.go

@@ -170,7 +170,12 @@ func extractRateLimitData(response *github.Response) (*v2.RateLimitDescription,
 		ra = &timestamppb.Timestamp{Seconds: ts}
 	}
 
+	status := v2.RateLimitDescription_STATUS_OK
+	if r <= 0 {
+		status = v2.RateLimitDescription_STATUS_OVERLIMIT
+	}
 	return &v2.RateLimitDescription{
+		Status:    status,
 		Limit:     l,
 		Remaining: r,
 		ResetAt:   ra,
@@ -220,3 +225,10 @@ func isNotFoundError(resp *github.Response) bool {
 	}
 	return resp.StatusCode == http.StatusNotFound
 }
+
+func isRatelimited(resp *github.Response) bool {
+	if resp == nil {
+		return false
+	}
+	return resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusTooManyRequests
+}

pkg/connector/invitation.go

@@ -99,7 +99,7 @@ func (i *invitationResourceType) List(ctx context.Context, parentID *v2.Resource
 		invitationResources = append(invitationResources, ir)
 	}
 	annotations.WithRateLimiting(restApiRateLimit)
-	return invitationResources, pageToken, nil, nil
+	return invitationResources, pageToken, annotations, nil
 }
 
 func (i *invitationResourceType) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
@@ -155,7 +155,7 @@ func (i *invitationResourceType) CreateAccount(
 	}
 	return &v2.CreateAccountResponse_SuccessResult{
 		Resource: r,
-	}, nil, nil, nil
+	}, nil, annotations, nil
 }
 
 func (i *invitationResourceType) Delete(ctx context.Context, resourceId *v2.ResourceId) (annotations.Annotations, error) {

pkg/connector/org.go

@@ -204,7 +204,11 @@ func (o *orgResourceType) Grants(
 		if isNotFoundError(resp) {
 			return nil, "", nil, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf("org: %s not found", orgName))
 		}
-		return nil, "", nil, fmt.Errorf("github-connectorv2: failed to list org members: %w", err)
+		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)

pkg/connector/org_role_test.go

@@ -58,7 +58,7 @@ func TestOrgRole(t *testing.T) {
 			grants = append(grants, nextGrants...)
 
 			require.Nil(t, err)
-			test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
+			test.AssertHasRatelimitAnnotations(t, grantsAnnotations)
 			if nextToken == "" {
 				break
 			}
@@ -98,7 +98,7 @@ func TestOrgRole(t *testing.T) {
 		require.Nil(t, err)
 		require.Empty(t, resources)
 		require.Empty(t, nextToken)
-		test.AssertNoRatelimitAnnotations(t, annotations)
+		test.AssertHasRatelimitAnnotations(t, annotations)
 
 		// Test Grants with permission error
 		role, _ := orgRoleResource(ctx, &OrganizationRole{
@@ -112,6 +112,6 @@ func TestOrgRole(t *testing.T) {
 		require.Empty(t, grants)
 		// The token should contain the initial state for users
 		require.NotEmpty(t, nextToken)
-		test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
+		test.AssertHasRatelimitAnnotations(t, grantsAnnotations)
 	})
 }

pkg/connector/org_test.go

@@ -39,7 +39,7 @@ func TestOrganization(t *testing.T) {
 
 		grants, nextToken, grantsAnnotations, err := client.Grants(ctx, organization, &pagination.Token{})
 		require.Nil(t, err)
-		test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
+		test.AssertHasRatelimitAnnotations(t, grantsAnnotations)
 		require.Equal(t, "", nextToken)
 		require.Len(t, grants, 2)
 

pkg/connector/repository_test.go

@@ -53,7 +53,7 @@ func TestRepository(t *testing.T) {
 			grants = append(grants, nextGrants...)
 
 			require.Nil(t, err)
-			test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
+			test.AssertHasRatelimitAnnotations(t, grantsAnnotations)
 			if nextToken == "" {
 				break
 			}

pkg/connector/team.go

@@ -186,6 +186,9 @@ func (o *teamResourceType) Grants(ctx context.Context, resource *v2.Resource, pT
 		if isNotFoundError(resp) {
 			return nil, "", nil, uhttp.WrapErrors(codes.NotFound, fmt.Sprintf("org: %d not found", org.GetID()))
 		}
+		if isRatelimited(resp) {
+			return nil, "", nil, uhttp.WrapErrors(codes.Unavailable, "too many requests", err)
+		}
 		return nil, "", nil, fmt.Errorf("github-connectorv2: failed to fetch team members: %w", err)
 	}
 

pkg/connector/team_test.go

@@ -41,7 +41,7 @@ func TestTeam(t *testing.T) {
 
 		grants, nextToken, grantsAnnotations, err := client.Grants(ctx, team, &pagination.Token{})
 		require.Nil(t, err)
-		test.AssertNoRatelimitAnnotations(t, grantsAnnotations)
+		test.AssertHasRatelimitAnnotations(t, grantsAnnotations)
 		require.Equal(t, "", nextToken)
 		require.Len(t, grants, 1)
 

pkg/connector/user_test.go

@@ -58,7 +58,7 @@ func TestUsersList(t *testing.T) {
 				&pagination.Token{},
 			)
 			require.Nil(t, err)
-			test.AssertNoRatelimitAnnotations(t, annotations)
+			test.AssertHasRatelimitAnnotations(t, annotations)
 			require.Equal(t, "", nextToken)
 			require.Len(t, users, 1)
 			require.Equal(t, *githubUser.Login, users[0].Id.Resource)

test/testhelpers.go

@@ -1,14 +1,13 @@
 package test
 
 import (
-	"slices"
 	"testing"
 
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 )
 
-func AssertNoRatelimitAnnotations(
+func AssertHasRatelimitAnnotations(
 	t *testing.T,
 	actualAnnotations annotations.Annotations,
 ) {
@@ -22,13 +21,7 @@ func AssertNoRatelimitAnnotations(
 		if err != nil {
 			continue
 		}
-		if slices.Contains(
-			[]v2.RateLimitDescription_Status{
-				v2.RateLimitDescription_STATUS_ERROR,
-				v2.RateLimitDescription_STATUS_OVERLIMIT,
-			},
-			ratelimitDescription.Status,
-		) {
+		if ratelimitDescription.Status != v2.RateLimitDescription_STATUS_OVERLIMIT {
 			t.Fatal("request was ratelimited, expected not to be ratelimited")
 		}
 	}