From bcac615a29419458a92557412dc196c5b440a494 Mon Sep 17 00:00:00 2001
From: Toni Spets <toni.spets@beeper.com>
Date: Wed, 31 Jan 2024 13:29:55 +0200
Subject: [PATCH] Add slice iterator for dbiter

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.
---
 dbutil/iter.go | 70 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 68 insertions(+), 2 deletions(-)

diff --git a/dbutil/iter.go b/dbutil/iter.go
index 4953bcb..c22a4cf 100644
--- a/dbutil/iter.go
+++ b/dbutil/iter.go
@@ -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 {
@@ -23,6 +27,8 @@ 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.
@@ -30,6 +36,11 @@ func NewRowIter[T any](rows Rows, convertFn func(Scannable) (T, error)) RowIter[
 	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
@@ -46,21 +57,32 @@ 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()
+
+	err := i.Rows.Err()
+	if err != nil {
+		i.err = err
+	} else {
+		i.err = ErrAlreadyIterated
+	}
+	return err
 }
 
 func (i *rowIterImpl[T]) AsList() (list []T, err error) {
@@ -70,3 +92,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
+}