Juju 7058/fix pagination for entitlement (#1435)

* fix paging mechanism for openfga tuples
canonical · Nov 18, 2024 · 33003cf · 33003cf
1 parent 5fc14d9
commit 33003cf
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 52 deletions.
diff --git a/internal/common/pagination/pagination.go b/internal/common/pagination/pagination.go
@@ -84,7 +84,7 @@ type OpenFGAPagination struct {
 
 // NewOpenFGAFilter creates a filter for token pagination.
 func NewOpenFGAFilter(limit int, token string) OpenFGAPagination {
-	if limit < 0 {
+	if limit <= 0 {
 		limit = defaultOpenFGAPageSize
 	}
 	if limit > maxOpenFGAPageSize {

diff --git a/internal/jimm/relation.go b/internal/jimm/relation.go
@@ -110,29 +110,60 @@ func (j *JIMM) ListObjectRelations(ctx context.Context, user *openfga.User, obje
 	if !user.JimmAdmin {
 		return nil, e, errors.E(op, errors.CodeUnauthorized, "unauthorized")
 	}
-	continuationToken, kind, err := pagination.DecodeEntitlementToken(entitlementToken)
+	responseTuples, nextToken, err := j.getObjectRelationsPage(ctx, object, pageSize, entitlementToken)
 	if err != nil {
-		return nil, e, err
+		return nil, e, errors.E(op, err)
 	}
+	// verify next page contains some entries. Otherwise return empty nextToken.
+	if len(responseTuples) == int(pageSize) && nextToken.String() != "" {
+		responseTuples, _, err := j.getObjectRelationsPage(ctx, object, 1, nextToken)
+		if err != nil {
+			return nil, e, errors.E(op, "error getting next page to verify it cointains something", err)
+		}
+		if len(responseTuples) == 0 {
+			nextToken = pagination.EntitlementToken{}
+		}
+	}
+	return responseTuples, nextToken, nil
+}
+
+func (j *JIMM) getObjectRelationsPage(ctx context.Context, object string, pageSize int32, entitlementToken pagination.EntitlementToken) ([]openfga.Tuple, pagination.EntitlementToken, error) {
+	var err error
+	var e pagination.EntitlementToken
 	tuple := &openfga.Tuple{}
 	tuple.Object, err = j.parseAndValidateTag(ctx, object)
 	if err != nil {
 		return nil, e, err
 	}
-	tuple.Target, err = j.parseAndValidateTag(ctx, kind.String())
-	if err != nil {
-		return nil, e, err
-	}
-
-	responseTuples, nextContinuationToken, err := j.OpenFGAClient.ReadRelatedObjects(ctx, *tuple, pageSize, continuationToken)
-	if err != nil {
-		return nil, e, errors.E(op, err)
-	}
-	nextEntitlementToken, err := pagination.NextEntitlementToken(kind, nextContinuationToken)
-	if err != nil {
-		return nil, e, err
+	var responseTuples []openfga.Tuple
+	nextToken := entitlementToken
+	// loop around entity kinds, each with a different continuation token.
+	for {
+		nextContinuationToken, kind, err := pagination.DecodeEntitlementToken(nextToken)
+		if err != nil {
+			return nil, e, err
+		}
+		tuple.Target, err = j.parseAndValidateTag(ctx, kind.String())
+		if err != nil {
+			return nil, e, err
+		}
+		t, nextContinuationToken, err := j.OpenFGAClient.ReadRelatedObjects(ctx, *tuple, pageSize, nextContinuationToken)
+		if err != nil {
+			return nil, e, err
+		}
+		responseTuples = append(responseTuples, t...)
+		// nolint:gosec
+		pageSize -= int32(len(t))
+		nextToken, err = pagination.NextEntitlementToken(kind, nextContinuationToken)
+		if err != nil {
+			return nil, e, err
+		}
+		// break on a full page or no other entries.
+		if pageSize <= 0 || nextToken.String() == "" {
+			break
+		}
 	}
-	return responseTuples, nextEntitlementToken, nil
+	return responseTuples, nextToken, nil
 }
 
 // parseTuples translate the api request struct containing tuples to a slice of openfga tuple keys.

diff --git a/internal/jimm/relation_test.go b/internal/jimm/relation_test.go
@@ -192,7 +192,7 @@ func TestListObjectRelations(t *testing.T) {
 	u := openfga.NewUser(&dbmodel.Identity{Name: "[email protected]"}, ofgaClient)
 	u.JimmAdmin = true
 
-	user, _, controller, model, _, _, _ := createTestControllerEnvironment(ctx, c, j.Database)
+	user, group, controller, model, _, cloud, _ := createTestControllerEnvironment(ctx, c, j.Database)
 	c.Assert(err, qt.IsNil)
 
 	err = j.AddRelation(ctx, u, []apiparams.RelationshipTuple{
@@ -211,28 +211,61 @@ func TestListObjectRelations(t *testing.T) {
 			Relation:     names.AuditLogViewerRelation.String(),
 			TargetObject: controller.ResourceTag().String(),
 		},
+		{
+			Object:       user.Tag().String(),
+			Relation:     names.AdministratorRelation.String(),
+			TargetObject: controller.ResourceTag().String(),
+		},
+		{
+			Object:       user.Tag().String(),
+			Relation:     names.AdministratorRelation.String(),
+			TargetObject: cloud.ResourceTag().String(),
+		},
+		{
+			Object:       user.Tag().String(),
+			Relation:     names.CanAddModelRelation.String(),
+			TargetObject: cloud.ResourceTag().String(),
+		},
+		{
+			Object:       user.Tag().String(),
+			Relation:     names.MemberRelation.String(),
+			TargetObject: group.ResourceTag().String(),
+		},
 	})
+
 	c.Assert(err, qt.IsNil)
 	type ExpectedTuple struct {
 		expectedRelation string
 		expectedTargetId string
 	}
 
 	testCases := []struct {
-		description    string
-		object         string
-		initialToken   pagination.EntitlementToken
-		expectedError  string
-		expectedLength int
-		expectedTuples []ExpectedTuple
+		description          string
+		object               string
+		initialToken         pagination.EntitlementToken
+		pageSize             int32
+		expectNumPages       int
+		expectedError        string
+		expectedTuplesLength int
+		expectedTuples       []ExpectedTuple
 	}{
 		{
-			description:    "test listing all relations",
-			object:         user.Tag().String(),
-			expectedLength: 3,
+			description:          "test listing all relations in single page",
+			object:               user.Tag().String(),
+			pageSize:             10,
+			expectNumPages:       1,
+			expectedTuplesLength: 7,
+		},
+		{
+			description:          "test listing all relations in multiple pages",
+			object:               user.Tag().String(),
+			pageSize:             2,
+			expectNumPages:       4,
+			expectedTuplesLength: 7,
 		},
 		{
 			description:   "invalid initial token",
+			object:        user.Tag().String(),
 			initialToken:  pagination.NewEntitlementToken("bar"),
 			expectedError: "failed to decode pagination token.*",
 		},
@@ -247,19 +280,22 @@ func TestListObjectRelations(t *testing.T) {
 		c.Run(t.description, func(c *qt.C) {
 			token := t.initialToken
 			tuples := []openfga.Tuple{}
+			numPages := 0
 			for {
-				res, nextToken, err := j.ListObjectRelations(ctx, u, t.object, 10, token)
+				res, nextToken, err := j.ListObjectRelations(ctx, u, t.object, t.pageSize, token)
 				if t.expectedError != "" {
 					c.Assert(err, qt.ErrorMatches, t.expectedError)
 					break
 				}
 				tuples = append(tuples, res...)
+				numPages += 1
 				if nextToken.String() == "" {
 					break
 				}
 				token = nextToken
 			}
-			c.Assert(tuples, qt.HasLen, t.expectedLength)
+			c.Assert(numPages, qt.Equals, t.expectNumPages)
+			c.Assert(tuples, qt.HasLen, t.expectedTuplesLength)
 			for i, expectedTuple := range t.expectedTuples {
 				c.Assert(tuples[i].Relation.String(), qt.Equals, expectedTuple.expectedRelation)
 				c.Assert(tuples[i].Target.ID, qt.Equals, expectedTuple.expectedTargetId)

diff --git a/internal/jimmhttp/rebac_admin/groups_integration_test.go b/internal/jimmhttp/rebac_admin/groups_integration_test.go
@@ -177,18 +177,10 @@ func (s rebacAdminSuite) TestGetGroupEntitlementsIntegration(c *gc.C) {
 	emptyPageToken := ""
 	req := resources.GetGroupsItemEntitlementsParams{NextPageToken: &emptyPageToken}
 	var entitlements []resources.EntityEntitlement
-	for {
-		res, err := s.groupSvc.GetGroupEntitlements(ctx, group.UUID, &req)
-		c.Assert(err, gc.IsNil)
-		c.Assert(res, gc.Not(gc.IsNil))
-		entitlements = append(entitlements, res.Data...)
-		if res.Next.PageToken == nil {
-			break
-		}
-		c.Assert(*res.Meta.PageToken, gc.Equals, *req.NextPageToken)
-		c.Assert(*res.Next.PageToken, gc.Not(gc.Equals), "")
-		req.NextPageToken = res.Next.PageToken
-	}
+	res, err := s.groupSvc.GetGroupEntitlements(ctx, group.UUID, &req)
+	c.Assert(err, gc.IsNil)
+	c.Assert(res, gc.Not(gc.IsNil))
+	entitlements = append(entitlements, res.Data...)
 	c.Assert(entitlements, gc.HasLen, 6)
 	modelEntitlementCount := 0
 	controllerEntitlementCount := 0

diff --git a/internal/jimmhttp/rebac_admin/identities_integration_test.go b/internal/jimmhttp/rebac_admin/identities_integration_test.go
@@ -221,18 +221,10 @@ func (s *identitiesSuite) TestIdentityEntitlements(c *gc.C) {
 	emptyPageToken := ""
 	req := resources.GetIdentitiesItemEntitlementsParams{NextPageToken: &emptyPageToken}
 	var entitlements []resources.EntityEntitlement
-	for {
-		res, err := identitySvc.GetIdentityEntitlements(ctx, user.Id(), &req)
-		c.Assert(err, gc.IsNil)
-		c.Assert(res, gc.Not(gc.IsNil))
-		entitlements = append(entitlements, res.Data...)
-		if res.Next.PageToken == nil || *res.Next.PageToken == "" {
-			break
-		}
-		c.Assert(*res.Meta.PageToken, gc.Equals, *req.NextPageToken)
-		c.Assert(*res.Next.PageToken, gc.Not(gc.Equals), "")
-		req.NextPageToken = res.Next.PageToken
-	}
+	res, err := identitySvc.GetIdentityEntitlements(ctx, user.Id(), &req)
+	c.Assert(err, gc.IsNil)
+	c.Assert(res, gc.Not(gc.IsNil))
+	entitlements = append(entitlements, res.Data...)
 	c.Assert(entitlements, gc.HasLen, 7)
 	modelEntitlementCount := 0
 	controllerEntitlementCount := 0