Added limits and continuation tokens

project-kessel · May 21, 2024 · 555f289 · 555f289
1 parent 1f3a4b5
commit 555f289
Show file tree

Hide file tree

Showing 6 changed files with 125 additions and 32 deletions.
diff --git a/api/rebac/v1/lookup.pb.go b/api/rebac/v1/lookup.pb.go
diff --git a/api/rebac/v1/lookup.proto b/api/rebac/v1/lookup.proto
@@ -21,8 +21,12 @@ message LookupSubjectsRequest {
     string subject_type = 1;
     string relation = 2;
     ObjectReference object = 3;
+    optional uint32 limit = 4;
+    optional string continuation_token = 5;
 }
 
 message LookupSubjectsResponse {
     SubjectReference subject = 1;
+    bool has_more = 2;
+    string continuation_token = 3;
 }
diff --git a/internal/biz/relationships.go b/internal/biz/relationships.go
@@ -10,9 +10,15 @@ import (
 // relationship domain objects re-used from the api layer for now, but otherwise would be defined here
 type TouchSemantics bool
 
+type ContinuationToken string
+type SubjectResult struct {
+	Subject      *v1.SubjectReference
+	Continuation ContinuationToken
+}
+
 type ZanzibarRepository interface {
 	Check(ctx context.Context, request *v1.CheckRequest) (*v1.CheckResponse, error)
-	LookupSubjects(context.Context, string, string, *v1.ObjectReference) (chan *v1.SubjectReference, chan error, error)
+	LookupSubjects(ctx context.Context, subject_type, relation string, object *v1.ObjectReference, limit uint32, continuation ContinuationToken) (chan *SubjectResult, chan error, error)
 	CreateRelationships(context.Context, []*v1.Relationship, TouchSemantics) error
 	ReadRelationships(context.Context, *v1.RelationshipFilter) ([]*v1.Relationship, error)
 	DeleteRelationships(context.Context, *v1.RelationshipFilter) error

diff --git a/internal/data/spicedb.go b/internal/data/spicedb.go
@@ -76,7 +76,14 @@ func NewSpiceDbRepository(c *conf.Data, logger log.Logger) (*SpiceDbRepository,
 	return &SpiceDbRepository{client}, cleanup, nil
 }
 
-func (s *SpiceDbRepository) LookupSubjects(ctx context.Context, subject_type string, relation string, object *apiV1.ObjectReference) (chan *apiV1.SubjectReference, chan error, error) {
+func (s *SpiceDbRepository) LookupSubjects(ctx context.Context, subject_type, relation string, object *apiV1.ObjectReference, limit uint32, continuation biz.ContinuationToken) (chan *biz.SubjectResult, chan error, error) {
+	var cursor *v1.Cursor = nil
+	if continuation != "" {
+		cursor = &v1.Cursor{
+			Token: string(continuation),
+		}
+	}
+
 	client, err := s.client.LookupSubjects(ctx, &v1.LookupSubjectsRequest{
 		Resource: &v1.ObjectReference{
 			ObjectType: object.Type,
@@ -85,14 +92,15 @@ func (s *SpiceDbRepository) LookupSubjects(ctx context.Context, subject_type str
 		Permission:              relation,
 		SubjectObjectType:       subject_type,
 		OptionalSubjectRelation: "",
-		OptionalConcreteLimit:   0,
+		OptionalConcreteLimit:   limit,
+		OptionalCursor:          cursor,
 	})
 
 	if err != nil {
 		return nil, nil, err
 	}
 
-	subjects := make(chan *apiV1.SubjectReference)
+	subjects := make(chan *biz.SubjectResult)
 	errs := make(chan error, 1)
 
 	go func() {
@@ -107,13 +115,21 @@ func (s *SpiceDbRepository) LookupSubjects(ctx context.Context, subject_type str
 				return
 			}
 
+			continuation := biz.ContinuationToken("")
+			if msg.AfterResultCursor != nil {
+				continuation = biz.ContinuationToken(msg.AfterResultCursor.Token)
+			}
+
 			subj := msg.GetSubject()
-			subjects <- &apiV1.SubjectReference{
-				Object: &apiV1.ObjectReference{
-					Type: subject_type,
-					Id:   subj.SubjectObjectId,
+			subjects <- &biz.SubjectResult{
+				Subject: &apiV1.SubjectReference{
+					Object: &apiV1.ObjectReference{
+						Type: subject_type,
+						Id:   subj.SubjectObjectId,
+					},
+					Relation: "",
 				},
-				Relation: "",
+				Continuation: continuation,
 			}
 		}
 	}()

diff --git a/internal/service/lookup.go b/internal/service/lookup.go
@@ -20,18 +20,28 @@ func NewLookupSubjectsService(repo biz.ZanzibarRepository) *LookupService {
 
 func (s *LookupService) Subjects(req *pb.LookupSubjectsRequest, conn pb.Lookup_SubjectsServer) error {
 	ctx := context.TODO() //Doesn't get context from grpc?
+	limit := uint32(1000)
+	if req.Limit != nil {
+		limit = *req.Limit
+	}
+
+	continuation := biz.ContinuationToken("")
+	if req.ContinuationToken != nil {
+		continuation = biz.ContinuationToken(*req.ContinuationToken)
+	}
 	subs, errs, err := s.repo.LookupSubjects(ctx, req.SubjectType, req.Relation, &pb.ObjectReference{
 		Type: req.Object.Type, //Need null check
 		Id:   req.Object.Id,
-	})
+	}, limit, continuation)
 
 	if err != nil {
 		return err
 	}
 
 	for sub := range subs {
 		err = conn.Send(&pb.LookupSubjectsResponse{
-			Subject: sub,
+			Subject:           sub.Subject,
+			ContinuationToken: string(sub.Continuation),
 		})
 		if err != nil {
 			return err

diff --git a/openapi.yaml b/openapi.yaml
@@ -160,6 +160,15 @@ paths:
                   in: query
                   schema:
                     type: string
+                - name: limit
+                  in: query
+                  schema:
+                    type: integer
+                    format: uint32
+                - name: continuationToken
+                  in: query
+                  schema:
+                    type: string
             responses:
                 "200":
                     description: OK
@@ -210,6 +219,10 @@ components:
             properties:
                 subject:
                     $ref: '#/components/schemas/api.rebac.v1.SubjectReference'
+                hasMore:
+                    type: boolean
+                continuationToken:
+                    type: string
         api.rebac.v1.ObjectReference:
             type: object
             properties: