Skip to content

Commit

Permalink
Add slice iterator for dbiter
Browse files Browse the repository at this point in the history
In addition return dbutil.ErrAlreadyIterated if the iterator is tried
to run twice after successful completion or the previous error that had
occured during iteration.
  • Loading branch information
hifi committed Jan 31, 2024
1 parent 19d1266 commit 58adfc4
Showing 1 changed file with 65 additions and 2 deletions.
67 changes: 65 additions & 2 deletions dbutil/iter.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@

package dbutil

import "errors"

var ErrAlreadyIterated = errors.New("this iterator has been already iterated")

// RowIter is a wrapper for [Rows] that allows conveniently iterating over rows
// with a predefined scanner function.
type RowIter[T any] interface {
Expand All @@ -23,13 +27,20 @@ type RowIter[T any] interface {
type rowIterImpl[T any] struct {
Rows
ConvertRow func(Scannable) (T, error)

err error
}

// NewRowIter creates a new RowIter from the given Rows and scanner function.
func NewRowIter[T any](rows Rows, convertFn func(Scannable) (T, error)) RowIter[T] {
return &rowIterImpl[T]{Rows: rows, ConvertRow: convertFn}
}

// NewRowIterWithError creates a new RowIter from the given Rows and scanner function with default error. If not nil, it will be returned without calling iterator function.
func NewRowIterWithError[T any](rows Rows, convertFn func(Scannable) (T, error), err error) RowIter[T] {
return &rowIterImpl[T]{Rows: rows, ConvertRow: convertFn, err: err}
}

func ScanSingleColumn[T any](rows Scannable) (val T, err error) {
err = rows.Scan(&val)
return
Expand All @@ -46,21 +57,29 @@ func ScanDataStruct[T NewableDataStruct[T]](rows Scannable) (T, error) {
}

func (i *rowIterImpl[T]) Iter(fn func(T) (bool, error)) error {
if i == nil || i.Rows == nil {
if i == nil {
return nil
} else if i.Rows == nil || i.err != nil {
return i.err
}
defer i.Rows.Close()

for i.Rows.Next() {
if item, err := i.ConvertRow(i.Rows); err != nil {
i.err = err
return err
} else if cont, err := fn(item); err != nil {
i.err = err
return err
} else if !cont {
break
}
}
return i.Rows.Err()

if err := i.Rows.Err(); err != nil {
i.err = err
}
return i.err
}

func (i *rowIterImpl[T]) AsList() (list []T, err error) {
Expand All @@ -70,3 +89,47 @@ func (i *rowIterImpl[T]) AsList() (list []T, err error) {
})
return
}

type sliceIterImpl[T any] struct {
items []T
err error
}

func NewSliceIter[T any](items []T) RowIter[T] {
return &sliceIterImpl[T]{items: items}
}

func NewSliceIterWithError[T any](items []T, err error) RowIter[T] {
return &sliceIterImpl[T]{items: items, err: err}
}

func (i *sliceIterImpl[T]) Iter(fn func(T) (bool, error)) error {
if i == nil {
return nil
} else if i.err != nil {
return i.err
}

for _, item := range i.items {
if cont, err := fn(item); err != nil {
i.err = err
return err
} else if !cont {
break
}
}

i.err = ErrAlreadyIterated
return nil
}

func (i *sliceIterImpl[T]) AsList() ([]T, error) {
if i == nil {
return nil, nil
} else if i.err != nil {
return nil, i.err
}

i.err = ErrAlreadyIterated
return i.items, nil
}

0 comments on commit 58adfc4

Please sign in to comment.