make use of lookup-subjects limit for pagination

Author: vroldanbet

.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"

go.mod

@@ -106,6 +106,7 @@ func RegisterPermissionCmd(rootCmd *cobra.Command) *cobra.Command {
 	lookupSubjectsCmd.Flags().Bool("json", false, "output as JSON")
 	lookupSubjectsCmd.Flags().String("revision", "", "optional revision at which to check")
 	lookupSubjectsCmd.Flags().String("caveat-context", "", "the caveat context to send along with the lookup, in JSON form")
+	lookupSubjectsCmd.Flags().Uint32("page-limit", 100, "limit of subject returned per page")
 	registerConsistencyFlags(lookupSubjectsCmd.Flags())
 
 	return permissionCmd
@@ -188,7 +189,7 @@ func checkCmdFunc(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	client, err := client.NewClient(cmd)
+	c, err := client.NewClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -218,7 +219,7 @@ func checkCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 
 	var trailerMD metadata.MD
-	resp, err := client.CheckPermission(ctx, request, grpc.Trailer(&trailerMD))
+	resp, err := c.CheckPermission(ctx, request, grpc.Trailer(&trailerMD))
 	if err != nil {
 		var debugInfo *v1.DebugInformation
 		if resp != nil {
@@ -370,7 +371,7 @@ func expandCmdFunc(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	client, err := client.NewClient(cmd)
+	c, err := client.NewClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -385,7 +386,7 @@ func expandCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 	log.Trace().Interface("request", request).Send()
 
-	resp, err := client.ExpandPermissionTree(cmd.Context(), request)
+	resp, err := c.ExpandPermissionTree(cmd.Context(), request)
 	if err != nil {
 		return err
 	}
@@ -425,7 +426,7 @@ func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	client, err := client.NewClient(cmd)
+	c, err := client.NewClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -444,7 +445,7 @@ func lookupResourcesCmdFunc(cmd *cobra.Command, args []string) error {
 	}
 	log.Trace().Interface("request", request).Send()
 
-	respStream, err := client.LookupResources(cmd.Context(), request)
+	respStream, err := c.LookupResources(cmd.Context(), request)
 	if err != nil {
 		return err
 	}
@@ -492,7 +493,7 @@ func lookupSubjectsCmdFunc(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	client, err := client.NewClient(cmd)
+	c, err := client.NewClient(cmd)
 	if err != nil {
 		return err
 	}
@@ -507,21 +508,37 @@ func lookupSubjectsCmdFunc(cmd *cobra.Command, args []string) error {
 		Context:                 caveatContext,
 		Consistency:             consistency,
 	}
-	log.Trace().Interface("request", request).Send()
-
-	respStream, err := client.LookupSubjects(cmd.Context(), request)
-	if err != nil {
-		return err
-	}
 
+	limit := cobrautil.MustGetUint32(cmd, "page-limit")
+	request.OptionalConcreteLimit = limit
+	log.Trace().Interface("request", request).Send()
+	lastCursor := request.OptionalCursor
 	for {
-		resp, err := respStream.Recv()
-		switch {
-		case errors.Is(err, io.EOF):
-			return nil
-		case err != nil:
+		request.OptionalCursor = lastCursor
+		respStream, err := c.LookupSubjects(cmd.Context(), request)
+		if err != nil {
 			return err
-		default:
+		}
+
+		var cursorToken string
+		if lastCursor != nil {
+			cursorToken = lastCursor.Token
+		}
+
+		log.Trace().Interface("request", request).Str("cursor", cursorToken).Msg("reading subjects page")
+		var relCount uint32
+		for {
+			resp, err := respStream.Recv()
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			if err != nil {
+				return err
+			}
+
+			lastCursor = resp.AfterResultCursor
+			relCount++
+
 			if cobrautil.MustGetBool(cmd, "json") {
 				prettyProto, err := PrettyProto(resp)
 				if err != nil {
@@ -533,8 +550,16 @@ func lookupSubjectsCmdFunc(cmd *cobra.Command, args []string) error {
 			console.Printf("%s:%s%s\n",
 				subjectType,
 				prettyLookupPermissionship(resp.Subject.SubjectObjectId, resp.Subject.Permissionship, resp.Subject.PartialCaveatInfo),
-				excludedSubjectsString(resp.ExcludedSubjects),
-			)
+				excludedSubjectsString(resp.ExcludedSubjects))
+		}
+
+		if relCount < limit || limit == 0 {
+			return nil
+		}
+
+		if relCount > limit {
+			log.Warn().Uint32("limit-specified", limit).Uint32("relationships-received", relCount).Msg("page limit ignored, pagination may not be supported by the server, consider updating SpiceDB")
+			return nil
 		}
 	}
 }

go.sum

@@ -0,0 +1,72 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
+	"github.com/authzed/spicedb/pkg/tuple"
+	"github.com/authzed/zed/internal/client"
+	"github.com/authzed/zed/internal/console"
+	zedtesting "github.com/authzed/zed/internal/testing"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestLookupSubjectsCmd(t *testing.T) {
+	cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t,
+		zedtesting.BoolFlag{FlagName: "json"},
+		zedtesting.StringFlag{FlagName: "caveat-context"},
+		zedtesting.BoolFlag{FlagName: "consistency-full", FlagValue: true},
+		zedtesting.StringFlag{FlagName: "consistency-at-least"},
+		zedtesting.StringFlag{FlagName: "revision"},
+		zedtesting.StringFlag{FlagName: "consistency-at-exactly"},
+		zedtesting.UintFlag32{FlagName: "page-limit", FlagValue: 1})
+
+	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)(cmd)
+	require.NoError(t, err)
+
+	_, err = c.WriteSchema(ctx, &v1.WriteSchemaRequest{Schema: testSchema})
+	require.NoError(t, err)
+
+	var updates []*v1.RelationshipUpdate
+	for i := 0; i < 1000; i++ {
+		updates = append(updates, &v1.RelationshipUpdate{
+			Operation:    v1.RelationshipUpdate_OPERATION_TOUCH,
+			Relationship: tuple.ParseRel(fmt.Sprintf("test/resource:1#reader@test/user:%d", i)),
+		})
+	}
+	wr := &v1.WriteRelationshipsRequest{
+		Updates: updates,
+	}
+	_, err = c.WriteRelationships(ctx, wr)
+	require.NoError(t, err)
+
+	originalFunc := console.Printf
+	var results []string
+	console.Printf = func(format string, a ...interface{}) {
+		results = append(results, fmt.Sprintf(format, a...))
+	}
+	defer func() {
+		console.Printf = originalFunc
+	}()
+
+	err = lookupSubjectsCmdFunc(cmd, []string{"test/resource:1", "read", "test/user"})
+	require.NoError(t, err)
+	require.Len(t, results, 1000)
+}

internal/commands/permission.go

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

internal/commands/permission_test.go

@@ -82,6 +82,41 @@ func (wc wasmClient) BulkCheckPermission(ctx context.Context, in *v1.BulkCheckPe
 	return client.BulkCheckPermission(ctx, in, opts...)
 }
 
+func (wc wasmClient) ExperimentalReflectSchema(ctx context.Context, in *v1.ExperimentalReflectSchemaRequest, opts ...grpc.CallOption) (*v1.ExperimentalReflectSchemaResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalReflectSchema(ctx, in, opts...)
+}
+
+func (wc wasmClient) ExperimentalComputablePermissions(ctx context.Context, in *v1.ExperimentalComputablePermissionsRequest, opts ...grpc.CallOption) (*v1.ExperimentalComputablePermissionsResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalComputablePermissions(ctx, in, opts...)
+}
+
+func (wc wasmClient) ExperimentalDependentRelations(ctx context.Context, in *v1.ExperimentalDependentRelationsRequest, opts ...grpc.CallOption) (*v1.ExperimentalDependentRelationsResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalDependentRelations(ctx, in, opts...)
+}
+
+func (wc wasmClient) ExperimentalDiffSchema(ctx context.Context, in *v1.ExperimentalDiffSchemaRequest, opts ...grpc.CallOption) (*v1.ExperimentalDiffSchemaResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalDiffSchema(ctx, in, opts...)
+}
+
+func (wc wasmClient) ExperimentalRegisterRelationshipCounter(ctx context.Context, in *v1.ExperimentalRegisterRelationshipCounterRequest, opts ...grpc.CallOption) (*v1.ExperimentalRegisterRelationshipCounterResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalRegisterRelationshipCounter(ctx, in, opts...)
+}
+
+func (wc wasmClient) ExperimentalCountRelationships(ctx context.Context, in *v1.ExperimentalCountRelationshipsRequest, opts ...grpc.CallOption) (*v1.ExperimentalCountRelationshipsResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalCountRelationships(ctx, in, opts...)
+}
+
+func (wc wasmClient) ExperimentalUnregisterRelationshipCounter(ctx context.Context, in *v1.ExperimentalUnregisterRelationshipCounterRequest, opts ...grpc.CallOption) (*v1.ExperimentalUnregisterRelationshipCounterResponse, error) {
+	client := v1.NewExperimentalServiceClient(wc.conn)
+	return client.ExperimentalUnregisterRelationshipCounter(ctx, in, opts...)
+}
+
 func (wc wasmClient) Watch(ctx context.Context, in *v1.WatchRequest, opts ...grpc.CallOption) (v1.WatchService_WatchClient, error) {
 	client := v1.NewWatchServiceClient(wc.conn)
 	return client.Watch(ctx, in, opts...)