github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/graveler/staging/iterator.go (about)

     1  package staging
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/treeverse/lakefs/pkg/graveler"
     7  	"github.com/treeverse/lakefs/pkg/kv"
     8  )
     9  
    10  type Iterator struct {
    11  	ctx   context.Context
    12  	itr   *kv.PartitionIterator
    13  	entry *graveler.ValueRecord
    14  	err   error
    15  }
    16  
    17  // NewStagingIterator initiates the staging iterator with a batchSize
    18  func NewStagingIterator(ctx context.Context, kvStore kv.Store, st graveler.StagingToken, batchSize int) *Iterator {
    19  	itr := kv.NewPartitionIterator(ctx, kvStore, (&graveler.StagedEntryData{}).ProtoReflect().Type(), graveler.StagingTokenPartition(st), batchSize)
    20  	return &Iterator{
    21  		ctx: ctx,
    22  		itr: itr,
    23  	}
    24  }
    25  
    26  func (s *Iterator) Next() bool {
    27  	if s.Err() != nil {
    28  		return false
    29  	}
    30  	if !s.itr.Next() {
    31  		s.entry = nil
    32  		return false
    33  	}
    34  	entry := s.itr.Entry()
    35  	if entry == nil {
    36  		s.err = graveler.ErrInvalid
    37  		return false
    38  	}
    39  	key := entry.Value.(*graveler.StagedEntryData).Key
    40  	value := graveler.StagedEntryFromProto(entry.Value.(*graveler.StagedEntryData))
    41  	s.entry = &graveler.ValueRecord{
    42  		Key:   key,
    43  		Value: value,
    44  	}
    45  	return true
    46  }
    47  
    48  func (s *Iterator) SeekGE(key graveler.Key) {
    49  	s.itr.SeekGE(key)
    50  }
    51  
    52  func (s *Iterator) Value() *graveler.ValueRecord {
    53  	if s.Err() != nil {
    54  		return nil
    55  	}
    56  	// Tombstone handling
    57  	if s.entry != nil && s.entry.Value != nil && s.entry.Identity == nil {
    58  		s.entry.Value = nil
    59  	}
    60  	return s.entry
    61  }
    62  
    63  func (s *Iterator) Err() error {
    64  	if s.err == nil {
    65  		return s.itr.Err()
    66  	}
    67  	return s.err
    68  }
    69  
    70  func (s *Iterator) Close() {
    71  	s.itr.Close()
    72  }