Add NewLimitIterator to cap iterator output to N items

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### New Features and Improvements
 
+* Added `NewLimitIterator` to `listing` package for lazy iteration with a cap on output items ([#1555](https://github.com/databricks/databricks-sdk-go/pull/1555)).
+
 ### Bug Fixes
 
 ### Documentation

listing/listing.go

@@ -242,3 +242,42 @@ func (s *SliceIterator[T]) Next(_ context.Context) (T, error) {
 	*s = (*s)[1:]
 	return v, nil
 }
+
+type limitIterator[T any] struct {
+	it        Iterator[T]
+	remaining int
+}
+
+// NewLimitIterator wraps an iterator and yields at most n items. If n <= 0,
+// the inner iterator is returned unchanged (no limit is applied). Note that
+// this differs from ToSliceN, where n == 0 means unlimited but negative n
+// yields zero items.
+func NewLimitIterator[T any](iter Iterator[T], n int) Iterator[T] {
+	if n <= 0 {
+		return iter
+	}
+	return &limitIterator[T]{
+		it:        iter,
+		remaining: n,
+	}
+}
+
+func (i *limitIterator[T]) HasNext(ctx context.Context) bool {
+	if i.remaining <= 0 {
+		return false
+	}
+	return i.it.HasNext(ctx)
+}
+
+func (i *limitIterator[T]) Next(ctx context.Context) (T, error) {
+	var t T
+	if i.remaining <= 0 {
+		return t, ErrNoMoreItems
+	}
+	t, err := i.it.Next(ctx)
+	if err != nil {
+		return t, err
+	}
+	i.remaining--
+	return t, nil
+}

listing/listing_test.go

@@ -249,3 +249,133 @@ func TestSliceIterator(t *testing.T) {
 	})
 
 }
+
+func TestLimitIterator(t *testing.T) {
+	t.Run("caps results when limit less than total", func(t *testing.T) {
+		rrs := []requestResponse{
+			{req: "page1", page: map[string][]int{"page": {1, 2}}},
+			{req: "page2", page: map[string][]int{"page": {3, 4}}},
+			{req: "page3", page: map[string][]int{"page": {5, 6}}},
+		}
+		iterator := makeIterator(rrs)
+		limited := listing.NewLimitIterator[int](iterator, 3)
+
+		items, err := listing.ToSlice(context.Background(), limited)
+		assert.NoError(t, err)
+		assert.Equal(t, []int{1, 2, 3}, items)
+	})
+
+	t.Run("limit greater than total returns all items", func(t *testing.T) {
+		rrs := []requestResponse{
+			{req: "page1", page: map[string][]int{"page": {1, 2}}},
+			{req: "page2", page: map[string][]int{"page": {3, 4}}},
+		}
+		iterator := makeIterator(rrs)
+		limited := listing.NewLimitIterator[int](iterator, 100)
+
+		items, err := listing.ToSlice(context.Background(), limited)
+		assert.NoError(t, err)
+		assert.Equal(t, []int{1, 2, 3, 4}, items)
+	})
+
+	t.Run("limit of one", func(t *testing.T) {
+		rrs := []requestResponse{
+			{req: "page1", page: map[string][]int{"page": {1, 2}}},
+			{req: "page2", page: map[string][]int{"page": {3, 4}}},
+		}
+		iterator := makeIterator(rrs)
+		limited := listing.NewLimitIterator[int](iterator, 1)
+
+		items, err := listing.ToSlice(context.Background(), limited)
+		assert.NoError(t, err)
+		assert.Equal(t, []int{1}, items)
+	})
+
+	t.Run("zero limit is no-op", func(t *testing.T) {
+		rrs := []requestResponse{
+			{req: "page1", page: map[string][]int{"page": {1, 2}}},
+			{req: "page2", page: map[string][]int{"page": {3, 4}}},
+		}
+		iterator := makeIterator(rrs)
+		limited := listing.NewLimitIterator[int](iterator, 0)
+
+		items, err := listing.ToSlice(context.Background(), limited)
+		assert.NoError(t, err)
+		assert.Equal(t, []int{1, 2, 3, 4}, items)
+	})
+
+	t.Run("negative limit is no-op", func(t *testing.T) {
+		rrs := []requestResponse{
+			{req: "page1", page: map[string][]int{"page": {1, 2}}},
+			{req: "page2", page: map[string][]int{"page": {3, 4}}},
+		}
+		iterator := makeIterator(rrs)
+		limited := listing.NewLimitIterator[int](iterator, -5)
+
+		items, err := listing.ToSlice(context.Background(), limited)
+		assert.NoError(t, err)
+		assert.Equal(t, []int{1, 2, 3, 4}, items)
+	})
+
+	t.Run("Next returns ErrNoMoreItems after limit exhausted", func(t *testing.T) {
+		iterator := listing.SliceIterator[int]([]int{1, 2, 3, 4, 5})
+		limited := listing.NewLimitIterator[int](&iterator, 2)
+
+		item1, err := limited.Next(context.Background())
+		assert.NoError(t, err)
+		assert.Equal(t, 1, item1)
+
+		item2, err := limited.Next(context.Background())
+		assert.NoError(t, err)
+		assert.Equal(t, 2, item2)
+
+		_, err = limited.Next(context.Background())
+		assert.ErrorIs(t, err, listing.ErrNoMoreItems)
+	})
+
+	t.Run("error propagation from inner iterator", func(t *testing.T) {
+		expectedErr := errors.New("some error")
+		iterator := listing.NewIterator[struct{}, []int, int](&struct{}{},
+			func(ctx context.Context, req struct{}) ([]int, error) {
+				return nil, expectedErr
+			},
+			nil,
+			nil,
+		)
+		limited := listing.NewLimitIterator[int](iterator, 5)
+
+		assert.True(t, limited.HasNext(context.Background()))
+		_, err := limited.Next(context.Background())
+		assert.ErrorIs(t, err, expectedErr)
+	})
+
+	t.Run("standard HasNext/Next loop", func(t *testing.T) {
+		rrs := []requestResponse{
+			{req: "page1", page: map[string][]int{"page": {1, 2}}},
+			{req: "page2", page: map[string][]int{"page": {3, 4}}},
+			{req: "page3", page: map[string][]int{"page": {5, 6}}},
+		}
+		iterator := makeIterator(rrs)
+		limited := listing.NewLimitIterator[int](iterator, 4)
+
+		values := make([]int, 0)
+		for limited.HasNext(context.Background()) {
+			v, err := limited.Next(context.Background())
+			assert.NoError(t, err)
+			values = append(values, v)
+		}
+		assert.Equal(t, []int{1, 2, 3, 4}, values)
+	})
+
+	t.Run("HasNext is idempotent", func(t *testing.T) {
+		iterator := listing.SliceIterator[int]([]int{1, 2, 3})
+		limited := listing.NewLimitIterator[int](&iterator, 2)
+
+		assert.True(t, limited.HasNext(context.Background()))
+		assert.True(t, limited.HasNext(context.Background()))
+
+		v, err := limited.Next(context.Background())
+		assert.NoError(t, err)
+		assert.Equal(t, 1, v)
+	})
+}