From fc123bd432e48ce75f6e6b1c20b0c95e626e5716 Mon Sep 17 00:00:00 2001
From: banerjeeadrish <adbanerj95@gmail.com>
Date: Fri, 19 Jun 2026 15:59:54 -0700
Subject: [PATCH] sync_diff_inspector: check rows.Err in TiDB iterator

---
 sync_diff_inspector/source/source_test.go | 21 +++++++++++++++++++++
 sync_diff_inspector/source/tidb.go        |  3 +++
 2 files changed, 24 insertions(+)

diff --git a/sync_diff_inspector/source/source_test.go b/sync_diff_inspector/source/source_test.go
index 6c6d3557..7bf89fae 100644
--- a/sync_diff_inspector/source/source_test.go
+++ b/sync_diff_inspector/source/source_test.go
@@ -17,6 +17,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"fmt"
 	"os"
 	"regexp"
@@ -260,6 +261,26 @@ func TestTiDBSource(t *testing.T) {
 
 	rowIter.Close()
 
+	// Test RowIterator returns the underlying database/sql row-stream error instead of treating it as clean EOF.
+	rowsErr := errors.New("tidb row stream cancelled")
+	errorRows := sqlmock.NewRows(tableCase.rowColumns).
+		AddRow(tableCase.rows[0]...).
+		AddRow(tableCase.rows[1]...).
+		RowError(1, rowsErr)
+	mock.ExpectQuery(tableCase.rowQuery).WillReturnRows(errorRows)
+	rowIter, err = tidb.GetRowsIterator(ctx, tableCase.rangeInfo)
+	require.NoError(t, err)
+
+	columns, err := rowIter.Next()
+	require.NoError(t, err)
+	require.NotNil(t, columns)
+
+	columns, err = rowIter.Next()
+	require.ErrorIs(t, err, rowsErr)
+	require.Nil(t, columns)
+
+	rowIter.Close()
+
 	analyze := tidb.GetTableAnalyzer()
 	statsRows := sqlmock.NewRows([]string{"is_index", "hist_id", "bucket_id", "count", "lower_bound", "upper_bound"})
 	for i := 0; i < 5; i++ {
diff --git a/sync_diff_inspector/source/tidb.go b/sync_diff_inspector/source/tidb.go
index 81c76f3b..a4fb0e51 100644
--- a/sync_diff_inspector/source/tidb.go
+++ b/sync_diff_inspector/source/tidb.go
@@ -79,6 +79,9 @@ func (s *TiDBRowsIterator) Next() (map[string]*dbutil.ColumnData, error) {
 	if s.rows.Next() {
 		return dbutil.ScanRow(s.rows)
 	}
+	if err := s.rows.Err(); err != nil {
+		return nil, errors.Trace(err)
+	}
 	return nil, nil
 }