More pagination fixes (#28)

ConductorOne · Dec 29, 2023 · d489b2a · d489b2a
1 parent e2c3912
commit d489b2a
Show file tree

Hide file tree

Showing 7 changed files with 11 additions and 31 deletions.
diff --git a/pkg/connector/account.go b/pkg/connector/account.go
@@ -82,17 +82,7 @@ func (o *accountResourceType) List(ctx context.Context, _ *v2.ResourceId, pt *pa
 		rv = append(rv, userResource)
 	}
 
-	hasNextPage := resp.NextToken != nil
-	if !hasNextPage {
-		return rv, "", nil, nil
-	}
-
-	nextPage, err := bag.NextToken(awsSdk.ToString(resp.NextToken))
-	if err != nil {
-		return nil, "", nil, fmt.Errorf("aws-connector: failed to marshal pagination bag: %w", err)
-	}
-
-	return rv, nextPage, nil, nil
+	return paginate(rv, bag, resp.NextToken)
 }
 
 func (o *accountResourceType) Entitlements(ctx context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {

diff --git a/pkg/connector/iam_group.go b/pkg/connector/iam_group.go
@@ -72,17 +72,7 @@ func (o *iamGroupResourceType) List(ctx context.Context, _ *v2.ResourceId, pt *p
 		rv = append(rv, groupResource)
 	}
 
-	hasNextPage := resp.IsTruncated && resp.Marker != nil
-	if !hasNextPage {
-		return rv, "", nil, nil
-	}
-
-	nextPage, err := bag.NextToken(awsSdk.ToString(resp.Marker))
-	if err != nil {
-		return nil, "", nil, fmt.Errorf("aws-connector: failed to marshal pagination bag: %w", err)
-	}
-
-	return rv, nextPage, nil, nil
+	return paginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
 }
 
 func (o *iamGroupResourceType) Entitlements(ctx context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
@@ -132,7 +122,7 @@ func (o *iamGroupResourceType) Grants(ctx context.Context, resource *v2.Resource
 		rv = append(rv, grant)
 	}
 
-	return PaginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
+	return paginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
 }
 
 func iamGroupBuilder(iamClient *iam.Client) *iamGroupResourceType {

diff --git a/pkg/connector/iam_user.go b/pkg/connector/iam_user.go
@@ -62,7 +62,7 @@ func (o *iamUserResourceType) List(ctx context.Context, _ *v2.ResourceId, pt *pa
 		rv = append(rv, userResource)
 	}
 
-	return PaginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
+	return paginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
 }
 
 func (o *iamUserResourceType) Entitlements(_ context.Context, _ *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {

diff --git a/pkg/connector/pagination.go b/pkg/connector/pagination.go
@@ -5,7 +5,7 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/pagination"
 )
 
-func Paginate[RV any](rv RV, bag *pagination.Bag, pageToken *string) (RV, string, annotations.Annotations, error) {
+func paginate[RV any](rv RV, bag *pagination.Bag, pageToken *string) (RV, string, annotations.Annotations, error) {
 	if pageToken == nil || *pageToken == "" {
 		return rv, "", nil, nil
 	}
@@ -18,10 +18,10 @@ func Paginate[RV any](rv RV, bag *pagination.Bag, pageToken *string) (RV, string
 	return rv, token, nil, nil
 }
 
-func PaginateTruncation[RV any](rv RV, bag *pagination.Bag, pageToken *string, isTruncated bool) (RV, string, annotations.Annotations, error) {
+func paginateTruncation[RV any](rv RV, bag *pagination.Bag, pageToken *string, isTruncated bool) (RV, string, annotations.Annotations, error) {
 	if !isTruncated {
 		return rv, "", nil, nil
 	}
 
-	return Paginate(rv, bag, pageToken)
+	return paginate(rv, bag, pageToken)
 }
diff --git a/pkg/connector/role.go b/pkg/connector/role.go
@@ -69,7 +69,7 @@ func (o *roleResourceType) List(ctx context.Context, _ *v2.ResourceId, pt *pagin
 		rv = append(rv, roleResource)
 	}
 
-	return PaginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
+	return paginateTruncation(rv, bag, resp.Marker, resp.IsTruncated)
 }
 
 func (o *roleResourceType) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {

diff --git a/pkg/connector/sso_group.go b/pkg/connector/sso_group.go
@@ -76,7 +76,7 @@ func (o *ssoGroupResourceType) List(ctx context.Context, _ *v2.ResourceId, pt *p
 		rv = append(rv, groupResource)
 	}
 
-	return Paginate(rv, bag, resp.NextToken)
+	return paginate(rv, bag, resp.NextToken)
 }
 
 func (o *ssoGroupResourceType) Entitlements(_ context.Context, resource *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {
@@ -156,7 +156,7 @@ func (o *ssoGroupResourceType) Grants(ctx context.Context, resource *v2.Resource
 		rv = append(rv, grant)
 	}
 
-	rv, nextPage, anno, err := Paginate(rv, bag, resp.NextToken)
+	rv, nextPage, anno, err := paginate(rv, bag, resp.NextToken)
 	if err != nil {
 		return nil, "", nil, fmt.Errorf("aws-connector: failed to marshal pagination bag [%s]: %w", awsSdk.ToString(input.GroupId), err)
 	}

diff --git a/pkg/connector/sso_user.go b/pkg/connector/sso_user.go
@@ -83,7 +83,7 @@ func (o *ssoUserResourceType) List(ctx context.Context, _ *v2.ResourceId, pt *pa
 		rv = append(rv, userResource)
 	}
 
-	return Paginate(rv, bag, resp.NextToken)
+	return paginate(rv, bag, resp.NextToken)
 }
 
 func (o *ssoUserResourceType) Entitlements(_ context.Context, _ *v2.Resource, _ *pagination.Token) ([]*v2.Entitlement, string, annotations.Annotations, error) {