lookup_paginated.go

package lookup

import (
    "context"
    "fmt"
    "sync/atomic"

    openfgapb "go.buf.build/openfga/go/openfga/api/openfga/v1"
    "golang.org/x/sync/errgroup"
)

func Lookup(ctx context.Context, req *openfgapb.LookupRequest) (*openfgapb.LookupResponse, error) {

    storeID := req.GetStoreId()
    modelID := req.GetAuthorizationModelId()
    targetObjectType := req.GetObjectType()
    targetRelation := req.GetRelation()
    targetUser := req.GetUser()

    contToken := req.GetContinuationToken() // todo: use contToken in the database query

    // query all the relationship tuples for the given storeID and targetObjectType
    // note: this should be ordered by something stable for pagination, for example 'ulid'..
    iter, err := database.QueryRelationships(storeID, targetObjectType, contToken)
    if err != nil {
        // handle error
    }

    cctx, cancel := context.WithCancel(ctx)
    defer cancel()

    pagesize := req.GetPageSize()

    resultsChan := make(chan struct {
        objectID string
        ulid     string
    }, pagesize+1)

    go func() {
        defer close(resultsChan)

        g, gctx := errgroup.WithContext(cctx)
        g.SetLimit(100) // limit to 100 concurrent checks, for example

        var evaluated uint64
        for iter.Next() {
            if atomic.LoadUint64(&evaluated) >= pagesize+1 {
                return
            }

            var objectID, ulid string
            err := iter.Scan(&objectID, &ulid)
            if err != nil {
                return nil, err
            }

            fn := func() error {
                object := fmt.Sprintf("%s:%s", targetObjectType, objectID)
                allowed, err := Check(gctx, storeID, modelID, object, targetRelation, targetUser)
                if err != nil {
                    return err
                }

                if allowed {
                    resultsChan <- struct {
                        objectID string
                        ulid     string
                    }{
                        objectID,
                        ulid,
                    }

                    atomic.AddUint64(&evaluated, 1)
                }

                return nil
            }

            g.Go(fn)
        }

        if err := g.Wait(); err != nil {
            // handle error
        }
    }()

    var continuationToken string

    var objectIDs []string
    for result := range resultsChan {
        objectIDs = append(objectIDs, result.objectID)

        if len(objectIDs) == pagesize+1 {
            continuationToken = encodeContinuationToken(result.ulid)
        }
    }

    response := &openfgapb.LookupResponse{
        ContinuationToken: continuationToken,
    }

    if len(objectIDs) >= pagesize {
        response.ObjectIds = objectIDs[:pagesize]
    } else {
        response.ObjectIds = objectIDs
    }

    return response, nil
}

    contToken := req.GetContinuationToken() // todo: use contToken in the database query

but this was what I wanted to see.. :(

    iter, err := database.QueryRelationships(storeID, targetObjectType, contToken)

this can give duplicate objectIds if the database is Dynamo. Also the name is confusing.. it should be database.GetAllObjectIdsOfType, no?