diff --git a/dumpling/export/ir_impl.go b/dumpling/export/ir_impl.go index 4afd2c29e7e14..c99191319bc01 100644 --- a/dumpling/export/ir_impl.go +++ b/dumpling/export/ir_impl.go @@ -210,6 +210,7 @@ func (td *tableData) Start(tctx *tcontext.Context, conn *sql.Conn) error { if err = rows.Err(); err != nil { return errors.Annotatef(err, "sql: %s", td.query) } + td.SQLRowIter = nil td.rows = rows if td.needColTypes { ns, err := rows.Columns() @@ -226,17 +227,27 @@ func (td *tableData) Start(tctx *tcontext.Context, conn *sql.Conn) error { td.colTypes = append(td.colTypes, c.DatabaseTypeName()) } } - td.SQLRowIter = newRowIter(rows, td.colLen) return nil } func (td *tableData) Rows() SQLRowIter { + // should be initialized lazily since it calls rows.Next() which might close the rows when + // there's nothing to read, causes code which relies on rows not closed to fail. + if td.SQLRowIter == nil { + td.SQLRowIter = newRowIter(td.rows, td.colLen) + } return td.SQLRowIter } func (td *tableData) Close() error { - return td.SQLRowIter.Close() + if td.SQLRowIter != nil { + // will close td.rows internally + return td.SQLRowIter.Close() + } else if td.rows != nil { + return td.rows.Close() + } + return nil } func (td *tableData) RawRows() *sql.Rows { @@ -351,16 +362,21 @@ func newMultiQueriesChunk(queries []string, colLength int) *multiQueriesChunk { func (td *multiQueriesChunk) Start(tctx *tcontext.Context, conn *sql.Conn) error { td.tctx = tctx td.conn = conn - td.SQLRowIter = newMultiQueryChunkIter(td.tctx, td.conn, td.queries, td.colLen) return nil } func (td *multiQueriesChunk) Rows() SQLRowIter { + if td.SQLRowIter == nil { + td.SQLRowIter = newMultiQueryChunkIter(td.tctx, td.conn, td.queries, td.colLen) + } return td.SQLRowIter } func (td *multiQueriesChunk) Close() error { - return td.SQLRowIter.Close() + if td.SQLRowIter != nil { + return td.SQLRowIter.Close() + } + return nil } func (*multiQueriesChunk) RawRows() *sql.Rows { diff --git a/dumpling/tests/basic/run.sh b/dumpling/tests/basic/run.sh index 4eae4d9b90d80..42b71689d9ee6 100644 --- a/dumpling/tests/basic/run.sh +++ b/dumpling/tests/basic/run.sh @@ -154,3 +154,11 @@ export GO_FAILPOINTS="github.com/pingcap/tidb/dumpling/export/SetIOTotalBytes=re run_dumpling -B "test_db" -L ${DUMPLING_OUTPUT_DIR}/dumpling.log cnt=$(grep "IOTotalBytes=" ${DUMPLING_OUTPUT_DIR}/dumpling.log | grep -v "IOTotalBytes=0" | wc -l) [ "$cnt" -ge 1 ] + +export GO_FAILPOINTS="" +export DUMPLING_TEST_PORT=4000 +echo "Test for empty query result, should success." +run_sql "drop database if exists test_db;" +run_sql "create database test_db;" +run_sql "create table test_db.test_table (a int primary key);" +run_dumpling --sql "select * from test_db.test_table" --filetype csv > ${DUMPLING_OUTPUT_DIR}/dumpling.log