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

     1  package ref
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sort"
     7  
     8  	"github.com/treeverse/lakefs/pkg/graveler"
     9  	"github.com/treeverse/lakefs/pkg/kv"
    10  )
    11  
    12  // CompareFunc type used for sorting in InMemIterator, it is a strictly bigger comparison function required for the
    13  // sort.Slice algorithm, implementors need to decide how to handle equal values
    14  type CompareFunc func(i, j int) bool
    15  
    16  var ErrIteratorClosed = errors.New("iterator already closed")
    17  
    18  // BranchSimpleIterator Iterates over repository's branches in a sorted way, since the branches are already sorted in DB according to BranchID
    19  type BranchSimpleIterator struct {
    20  	ctx           context.Context
    21  	store         kv.Store
    22  	itr           *kv.PrimaryIterator
    23  	repoPartition string
    24  	value         *graveler.BranchRecord
    25  	err           error
    26  }
    27  
    28  func NewBranchSimpleIterator(ctx context.Context, store kv.Store, repo *graveler.RepositoryRecord) (*BranchSimpleIterator, error) {
    29  	repoPartition := graveler.RepoPartition(repo)
    30  	it, err := kv.NewPrimaryIterator(ctx, store, (&graveler.BranchData{}).ProtoReflect().Type(),
    31  		repoPartition, []byte(graveler.BranchPath("")), kv.IteratorOptionsFrom([]byte("")))
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	return &BranchSimpleIterator{
    37  		ctx:           ctx,
    38  		store:         store,
    39  		itr:           it,
    40  		repoPartition: repoPartition,
    41  		value:         nil,
    42  		err:           nil,
    43  	}, nil
    44  }
    45  
    46  func (bi *BranchSimpleIterator) Next() bool {
    47  	if bi.Err() != nil {
    48  		return false
    49  	}
    50  	if !bi.itr.Next() {
    51  		bi.value = nil
    52  		return false
    53  	}
    54  	entry := bi.itr.Entry()
    55  	if entry == nil {
    56  		bi.err = graveler.ErrInvalid
    57  		return false
    58  	}
    59  	value, ok := entry.Value.(*graveler.BranchData)
    60  	if !ok {
    61  		bi.err = graveler.ErrReadingFromStore
    62  		return false
    63  	}
    64  
    65  	bi.value = &graveler.BranchRecord{
    66  		BranchID: graveler.BranchID(value.Id),
    67  		Branch:   branchFromProto(value),
    68  	}
    69  	return true
    70  }
    71  
    72  func (bi *BranchSimpleIterator) SeekGE(id graveler.BranchID) {
    73  	if bi.Err() == nil {
    74  		bi.itr.Close() // Close previous before creating new iterator
    75  		bi.itr, bi.err = kv.NewPrimaryIterator(bi.ctx, bi.store, (&graveler.BranchData{}).ProtoReflect().Type(),
    76  			bi.repoPartition, []byte(graveler.BranchPath("")), kv.IteratorOptionsFrom([]byte(graveler.BranchPath(id))))
    77  	}
    78  }
    79  
    80  func (bi *BranchSimpleIterator) Value() *graveler.BranchRecord {
    81  	if bi.Err() != nil {
    82  		return nil
    83  	}
    84  	return bi.value
    85  }
    86  
    87  func (bi *BranchSimpleIterator) Err() error {
    88  	if bi.err != nil {
    89  		return bi.err
    90  	}
    91  	return bi.itr.Err()
    92  }
    93  
    94  func (bi *BranchSimpleIterator) Close() {
    95  	bi.err = ErrIteratorClosed
    96  	if bi.itr != nil {
    97  		bi.itr.Close()
    98  	}
    99  }
   100  
   101  // BranchByCommitIterator iterates over repository's branches ordered by Commit ID. Currently, implemented as in-mem iterator
   102  type BranchByCommitIterator struct {
   103  	ctx    context.Context
   104  	values []*graveler.BranchRecord
   105  	value  *graveler.BranchRecord
   106  	idx    int
   107  	err    error
   108  }
   109  
   110  func (b *BranchByCommitIterator) SortByCommitID(i, j int) bool {
   111  	return b.values[i].CommitID.String() <= b.values[j].CommitID.String()
   112  }
   113  
   114  func NewBranchByCommitIterator(ctx context.Context, store kv.Store, repo *graveler.RepositoryRecord) (*BranchByCommitIterator, error) {
   115  	bi := &BranchByCommitIterator{
   116  		ctx:    ctx,
   117  		values: make([]*graveler.BranchRecord, 0),
   118  	}
   119  	itr, err := NewBranchSimpleIterator(ctx, store, repo)
   120  	if err != nil {
   121  		return nil, err
   122  	}
   123  	defer itr.Close()
   124  	for itr.Next() {
   125  		value := itr.Value()
   126  		if value == nil { // nil only if internal itr has errors
   127  			return nil, itr.Err()
   128  		}
   129  		bi.values = append(bi.values, itr.Value())
   130  	}
   131  	if itr.Err() != nil {
   132  		return nil, err
   133  	}
   134  
   135  	sort.Slice(bi.values, bi.SortByCommitID)
   136  	return bi, nil
   137  }
   138  
   139  func (b *BranchByCommitIterator) Next() bool {
   140  	if b.idx >= len(b.values) {
   141  		return false
   142  	}
   143  	b.value = b.values[b.idx]
   144  	b.idx++
   145  	return true
   146  }
   147  
   148  func (b *BranchByCommitIterator) SeekGE(_ graveler.BranchID) {
   149  	panic("Not Implemented")
   150  }
   151  
   152  func (b *BranchByCommitIterator) Value() *graveler.BranchRecord {
   153  	if b.Err() != nil {
   154  		return nil
   155  	}
   156  	return b.value
   157  }
   158  
   159  func (b *BranchByCommitIterator) Err() error {
   160  	return b.err
   161  }
   162  
   163  func (b *BranchByCommitIterator) Close() {
   164  	b.err = ErrIteratorClosed
   165  }