Merge pull request #395 from authzed/introduce-LR-page-limit

introduce lookup resources page limit

.golangci.yaml

@@ -10,7 +10,6 @@ linters:
   enable:
     - "bidichk"
     - "bodyclose"
-    - "deadcode"
     - "errcheck"
     - "errname"
     - "errorlint"
@@ -20,7 +19,6 @@ linters:
     - "gosec"
     - "gosimple"
     - "govet"
-    - "ifshort"
     - "importas"
     - "ineffassign"
     - "makezero"
@@ -30,12 +28,10 @@ linters:
     - "revive"
     - "rowserrcheck"
     - "staticcheck"
-    - "structcheck"
     - "stylecheck"
     - "tenv"
     - "typecheck"
     - "unconvert"
     - "unused"
-    - "varcheck"
     - "wastedassign"
     - "whitespace"

internal/commands/permission.go

@@ -101,6 +101,7 @@ func RegisterPermissionCmd(rootCmd *cobra.Command) *cobra.Command {
 	lookupResourcesCmd.Flags().Bool("json", false, "output as JSON")
 	lookupResourcesCmd.Flags().String("revision", "", "optional revision at which to check")
 	lookupResourcesCmd.Flags().String("caveat-context", "", "the caveat context to send along with the lookup, in JSON form")
+	lookupResourcesCmd.Flags().Uint32("page-limit", 0, "limit of relations returned per page")
 	registerConsistencyFlags(lookupResourcesCmd.Flags())
 
 	permissionCmd.AddCommand(lookupSubjectsCmd)
@@ -416,6 +417,8 @@ func expandCmdFunc(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+var newLookupResourcesPageCallbackForTests func(readByPage uint)
+
 func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 	objectNS := args[0]
 	relation := args[1]
@@ -424,6 +427,7 @@ func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	pageLimit := cobrautil.MustGetUint32(cmd, "page-limit")
 	caveatContext, err := GetCaveatContext(cmd)
 	if err != nil {
 		return err
@@ -438,46 +442,69 @@ func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
-	request := &v1.LookupResourcesRequest{
-		ResourceObjectType: objectNS,
-		Permission:         relation,
-		Subject: &v1.SubjectReference{
-			Object: &v1.ObjectReference{
-				ObjectType: subjectNS,
-				ObjectId:   subjectID,
-			},
-			OptionalRelation: subjectRel,
-		},
-		Context:     caveatContext,
-		Consistency: consistency,
-	}
-	log.Trace().Interface("request", request).Send()
-
-	respStream, err := client.LookupResources(cmd.Context(), request)
-	if err != nil {
-		return err
-	}
 
+	var cursor *v1.Cursor
+	var totalCount uint
 	for {
-		resp, err := respStream.Recv()
-		switch {
-		case errors.Is(err, io.EOF):
-			return nil
-		case err != nil:
+		request := &v1.LookupResourcesRequest{
+			ResourceObjectType: objectNS,
+			Permission:         relation,
+			Subject: &v1.SubjectReference{
+				Object: &v1.ObjectReference{
+					ObjectType: subjectNS,
+					ObjectId:   subjectID,
+				},
+				OptionalRelation: subjectRel,
+			},
+			Context:        caveatContext,
+			Consistency:    consistency,
+			OptionalLimit:  pageLimit,
+			OptionalCursor: cursor,
+		}
+		log.Trace().Interface("request", request).Uint32("page-limit", pageLimit).Send()
+
+		respStream, err := client.LookupResources(cmd.Context(), request)
+		if err != nil {
 			return err
-		default:
-			if cobrautil.MustGetBool(cmd, "json") {
-				prettyProto, err := PrettyProto(resp)
-				if err != nil {
-					return err
+		}
+
+		var count uint
+
+	stream:
+		for {
+			resp, err := respStream.Recv()
+			switch {
+			case errors.Is(err, io.EOF):
+				break stream
+			case err != nil:
+				return err
+			default:
+				count++
+				totalCount++
+				if cobrautil.MustGetBool(cmd, "json") {
+					prettyProto, err := PrettyProto(resp)
+					if err != nil {
+						return err
+					}
+
+					console.Println(string(prettyProto))
 				}
 
-				console.Println(string(prettyProto))
+				console.Println(prettyLookupPermissionship(resp.ResourceObjectId, resp.Permissionship, resp.PartialCaveatInfo))
+				cursor = resp.AfterResultCursor
 			}
+		}
 
-			console.Println(prettyLookupPermissionship(resp.ResourceObjectId, resp.Permissionship, resp.PartialCaveatInfo))
+		if newLookupResourcesPageCallbackForTests != nil {
+			newLookupResourcesPageCallbackForTests(count)
+		}
+		if count == 0 || pageLimit == 0 || count < uint(pageLimit) {
+			log.Trace().Interface("request", request).Uint32("page-limit", pageLimit).Uint("count", totalCount).Send()
+			break
 		}
 	}
+
+	return nil
 }
 
 func lookupSubjectsCmdFunc(cmd *cobra.Command, args []string) error {

internal/commands/permission_test.go

@@ -5,6 +5,11 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/authzed/spicedb/pkg/tuple"
+
+	"github.com/authzed/zed/internal/console"
+	zedtesting "github.com/authzed/zed/internal/testing"
+
 	"github.com/rs/zerolog"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/require"
@@ -98,3 +103,94 @@ func TestCheckErrorWithInvalidDebugInformation(t *testing.T) {
 	require.NotNil(t, err)
 	require.ErrorContains(t, err, "unknown field: invalid")
 }
+
+func TestLookupResourcesCommand(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	srv := zedtesting.NewTestServer(ctx, t)
+	go func() {
+		require.NoError(t, srv.Run(ctx))
+	}()
+	conn, err := srv.GRPCDialContext(ctx)
+	require.NoError(t, err)
+
+	originalClient := client.NewClient
+	defer func() {
+		client.NewClient = originalClient
+	}()
+
+	client.NewClient = zedtesting.ClientFromConn(conn)
+
+	c, err := zedtesting.ClientFromConn(conn)(nil)
+	require.NoError(t, err)
+
+	_, err = c.WriteSchema(ctx, &v1.WriteSchemaRequest{Schema: testSchema})
+	require.NoError(t, err)
+
+	var updates []*v1.RelationshipUpdate
+	for i := 0; i < 10; i++ {
+		updates = append(updates, &v1.RelationshipUpdate{
+			Operation:    v1.RelationshipUpdate_OPERATION_TOUCH,
+			Relationship: tuple.ParseRel(fmt.Sprintf("test/resource:%d#reader@test/user:1", i)),
+		})
+	}
+
+	_, err = c.WriteRelationships(ctx, &v1.WriteRelationshipsRequest{Updates: updates})
+	require.NoError(t, err)
+
+	// we override this to obtain the results being printed and validate them
+	previous := console.Println
+	defer func() {
+		console.Println = previous
+	}()
+	var count int
+	console.Println = func(values ...any) {
+		count += len(values)
+	}
+
+	// use test callback to make sure pagination is correct
+	var receivedPageSizes []uint
+	newLookupResourcesPageCallbackForTests = func(readPageSize uint) {
+		receivedPageSizes = append(receivedPageSizes, readPageSize)
+	}
+	defer func() {
+		newLookupResourcesPageCallbackForTests = nil
+	}()
+
+	// test no page size, server computes returns all resources in one go
+	cmd := testLookupResourcesCommand(t, 0)
+	err = lookupResourcesCmdFunc(cmd, []string{"test/resource", "read", "test/user:1"})
+	require.NoError(t, err)
+	require.Equal(t, 10, count)
+	require.EqualValues(t, []uint{10}, receivedPageSizes)
+
+	// use page size same as number of elements
+	count = 0
+	receivedPageSizes = nil
+	cmd = testLookupResourcesCommand(t, 10)
+	err = lookupResourcesCmdFunc(cmd, []string{"test/resource", "read", "test/user:1"})
+	require.NoError(t, err)
+	require.Equal(t, 10, count)
+	require.EqualValues(t, []uint{10, 0}, receivedPageSizes)
+
+	// use odd page size
+	count = 0
+	receivedPageSizes = nil
+	cmd = testLookupResourcesCommand(t, 3)
+	err = lookupResourcesCmdFunc(cmd, []string{"test/resource", "read", "test/user:1"})
+	require.NoError(t, err)
+	require.Equal(t, 10, count)
+	require.EqualValues(t, []uint{3, 3, 3, 1}, receivedPageSizes)
+}
+
+func testLookupResourcesCommand(t *testing.T, limit uint32) *cobra.Command {
+	return zedtesting.CreateTestCobraCommandWithFlagValue(t,
+		zedtesting.BoolFlag{FlagName: "consistency-full", FlagValue: true},
+		zedtesting.StringFlag{FlagName: "consistency-at-least"},
+		zedtesting.BoolFlag{FlagName: "consistency-min-latency", FlagValue: false},
+		zedtesting.StringFlag{FlagName: "consistency-at-exactly"},
+		zedtesting.StringFlag{FlagName: "revision"},
+		zedtesting.StringFlag{FlagName: "caveat-context"},
+		zedtesting.UintFlag32{FlagName: "page-limit", FlagValue: limit},
+		zedtesting.BoolFlag{FlagName: "json"})
+}

internal/commands/relationship_test.go

@@ -25,6 +25,8 @@ import (
 const testSchema = `definition test/resource {
 	relation reader: test/user
 	relation writer: test/user
+
+	permission read = reader + writer
 }
 
 definition test/user {}`

internal/console/console.go

@@ -23,7 +23,7 @@ var Errorf = func(format string, a ...any) {
 }
 
 // Println prints a line with optional values to the console.
-func Println(values ...any) {
+var Println = func(values ...any) {
 	for _, value := range values {
 		Printf("%v\n", value)
 	}

internal/testing/test_helpers.go

@@ -17,7 +17,7 @@ import (
 	"github.com/authzed/zed/internal/client"
 )
 
-func ClientFromConn(conn *grpc.ClientConn) func(cmd *cobra.Command) (client.Client, error) {
+func ClientFromConn(conn *grpc.ClientConn) func(_ *cobra.Command) (client.Client, error) {
 	return func(_ *cobra.Command) (client.Client, error) {
 		return &authzed.ClientWithExperimental{
 			Client: authzed.Client{