address pr comments

Author: prnvbn

README.md

@@ -24,7 +24,7 @@ Both support TCP and HTTP transport. When in doubt, use the native interface.
 * Supports both native ClickHouse TCP and HTTP client-server protocols
 * Compatibility with [`database/sql`](#std-databasesql-interface) ([slower](#benchmark) than [native interface](#native-interface)!)
 * [`database/sql`](#std-databasesql-interface) supports both native TCP and HTTP protocols for transport.
-* Marshal rows into structs ([ScanStruct](examples/clickhouse_api/scan_struct.go), [Select](examples/clickhouse_api/select_struct.go))
+* Marshal rows into structs ([ScanStruct](examples/clickhouse_api/scan_struct.go), [Select](examples/clickhouse_api/select_struct.go), [StructIter](examples/clickhouse_api/iterators.go) for native `driver.Rows`)
 * Unmarshal struct to row ([AppendStruct](benchmark/v2/write-native-struct/main.go))
 * Connection pool (for both TCP-Native and HTTP)
 * Failover and load balancing

examples/clickhouse_api/iterators.go

@@ -7,15 +7,25 @@ import (
 	chdriver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 )
 
+<<<<<<< Updated upstream
 func Iterators() error {
+=======
+func Iterators() (err error) {
+>>>>>>> Stashed changes
 	conn, err := GetNativeConnection(nil, nil, nil)
 	if err != nil {
 		return err
 	}
 
 	ctx := context.Background()
 	defer func() {
+<<<<<<< Updated upstream
 		conn.Exec(ctx, "DROP TABLE example_iterators")
+=======
+		if dropErr := conn.Exec(ctx, "DROP TABLE example_iterators"); dropErr != nil && err == nil {
+			err = fmt.Errorf("drop example_iterators: %w", dropErr)
+		}
+>>>>>>> Stashed changes
 	}()
 
 	if err := conn.Exec(ctx, `DROP TABLE IF EXISTS example_iterators`); err != nil {

lib/driver/iter.go

@@ -1,24 +1,43 @@
 package driver
 
-import "iter"
+import (
+	"errors"
+	"iter"
+)
 
+// StructIter returns an iterator that scans each row into T with ScanStruct.
+// It works with native Rows, not database/sql.Rows.
 func StructIter[T any](rows Rows) iter.Seq2[T, error] {
 	return func(yield func(T, error) bool) {
-		defer rows.Close()
-
 		for rows.Next() {
 			var value T
 			if err := rows.ScanStruct(&value); err != nil {
 				var zero T
+				if closeErr := rows.Close(); closeErr != nil {
+					err = errors.Join(err, closeErr)
+				}
 				_ = yield(zero, err)
 				return
 			}
 			if !yield(value, nil) {
+				// The caller stopped iteration, so the protocol forbids yielding a close error.
+				if closeErr := rows.Close(); closeErr != nil {
+					return
+				}
 				return
 			}
 		}
 
 		if err := rows.Err(); err != nil {
+			var zero T
+			if closeErr := rows.Close(); closeErr != nil {
+				err = errors.Join(err, closeErr)
+			}
+			_ = yield(zero, err)
+			return
+		}
+
+		if err := rows.Close(); err != nil {
 			var zero T
 			_ = yield(zero, err)
 		}

lib/driver/iter_test.go

@@ -12,6 +12,7 @@ type testRows struct {
 	index           int
 	closeCalls      int
 	err             error
+	closeErr        error
 	scanStructErrAt int
 }
 
@@ -49,7 +50,7 @@ func (r *testRows) Columns() []string { return nil }
 
 func (r *testRows) Close() error {
 	r.closeCalls++
-	return nil
+	return r.closeErr
 }
 
 func (r *testRows) Err() error { return r.err }
@@ -128,3 +129,55 @@ func TestStructIterTerminalRowsError(t *testing.T) {
 		t.Fatalf("unexpected values before terminal error: %#v", got)
 	}
 }
+
+func TestStructIterCloseError(t *testing.T) {
+	type item struct {
+		Value int
+	}
+
+	rows := &testRows{values: []int{1}, closeErr: io.ErrClosedPipe}
+
+	var got []item
+	var gotErr error
+	for value, err := range StructIter[item](rows) {
+		if err != nil {
+			gotErr = err
+			break
+		}
+		got = append(got, value)
+	}
+
+	if !errors.Is(gotErr, io.ErrClosedPipe) {
+		t.Fatalf("unexpected close error: %v", gotErr)
+	}
+	if !reflect.DeepEqual(got, []item{{Value: 1}}) {
+		t.Fatalf("unexpected values before close error: %#v", got)
+	}
+	if rows.closeCalls != 1 {
+		t.Fatalf("unexpected close calls: %d", rows.closeCalls)
+	}
+}
+
+func TestStructIterStopsAfterCallerBreak(t *testing.T) {
+	type item struct {
+		Value int
+	}
+
+	rows := &testRows{values: []int{1, 2, 3}}
+
+	var got []item
+	for value, err := range StructIter[item](rows) {
+		if err != nil {
+			t.Fatalf("unexpected iter error: %v", err)
+		}
+		got = append(got, value)
+		break
+	}
+
+	if !reflect.DeepEqual(got, []item{{Value: 1}}) {
+		t.Fatalf("unexpected values before caller break: %#v", got)
+	}
+	if rows.closeCalls != 1 {
+		t.Fatalf("unexpected close calls: %d", rows.closeCalls)
+	}
+}

tests/iterator_test.go

@@ -0,0 +1,55 @@
+package tests
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ClickHouse/clickhouse-go/v2"
+	chdriver "github.com/ClickHouse/clickhouse-go/v2/lib/driver"
+)
+
+func TestStructIterProtocols(t *testing.T) {
+	TestProtocols(t, func(t *testing.T, protocol clickhouse.Protocol) {
+		conn, err := GetNativeConnection(t, protocol, nil, nil, nil)
+		require.NoError(t, err)
+
+		ctx := context.Background()
+		const table = "test_struct_iter"
+		require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS "+table))
+		t.Cleanup(func() {
+			require.NoError(t, conn.Exec(ctx, "DROP TABLE IF EXISTS "+table))
+		})
+		require.NoError(t, conn.Exec(ctx, `
+			CREATE TABLE test_struct_iter (
+				Col1 UInt8,
+				Col2 String
+			) ENGINE = Memory
+		`))
+
+		batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_struct_iter")
+		require.NoError(t, err)
+		require.NoError(t, batch.Append(uint8(1), "one"))
+		require.NoError(t, batch.Append(uint8(2), "two"))
+		require.NoError(t, batch.Send())
+
+		rows, err := conn.Query(ctx, "SELECT Col1, Col2 FROM test_struct_iter ORDER BY Col1")
+		require.NoError(t, err)
+
+		type result struct {
+			Col1 uint8
+			Col2 string
+		}
+		var got []result
+		for value, err := range chdriver.StructIter[result](rows) {
+			require.NoError(t, err)
+			got = append(got, value)
+		}
+
+		require.Equal(t, []result{
+			{Col1: 1, Col2: "one"},
+			{Col1: 2, Col2: "two"},
+		}, got)
+	})
+}