Repository: pingcap/dumpling
Pull request #246 (2021-03-22 12:42:45)

*: add logs and try fixing error log
<!--
Thank you for contributing to Dumpling! Please read the [CONTRIBUTING](https://github.com/pingcap/dumpling/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.

If you want to open the **Challenge Program** pull request, please use the following template:
https://raw.githubusercontent.com/pingcap/.github/master/.github/PULL_REQUEST_TEMPLATE/challenge-program.md

You can use it with query parameters: https://github.com/pingcap/dumpling/compare/master...${you branch}?template=challenge-program.md
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
1. Dumpling's log is not clear when some bugs occurred.
2. We didn't check `FileRowIter.Error` when there is no next. However, we check this when we start this `rowIter`. But this is safer.

### What is changed and how it works?
1. Add more log.
2. Check Error.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test

Related changes

 - Need to cherry-pick to the release branch
 
### Release note

<!-- bugfixes or new feature need a release note, must in the form of a list, such as

- support -T/--tables-list argument

or if no need to be included in the release note, just add the following line

- No release note
-->
- No release note


hi-rustin: @ti-chi-bot Could you please mark this pr?
ti-chi-bot: @lichunzhu: PR needs rebase.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
hi-rustin: > @lichunzhu: PR needs rebase.

These reply and label will be automatically deleted after you resolve the conflict.

Patch:
diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index 09a6caf..224046e 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -3,11 +3,11 @@ on:
   push:
     branches:
       - master
-      - 'release-.*'
+      - release-*
   pull_request:
     branches:
       - master
-      - 'release-.*'
+      - release-*
 
 jobs:
   unit-test:
diff --git a/Makefile b/Makefile
index 3d5d932..b8ac63f 100644
--- a/Makefile
+++ b/Makefile
@@ -27,7 +27,7 @@ test: failpoint-enable
 	$(GOTEST) $(RACEFLAG) -coverprofile=coverage.txt -covermode=atomic -tags leak ./... || ( make failpoint-disable && exit 1 )
 	@make failpoint-disable
 
-integration_test: failpoint-enable bin/dumpling
+integration_test: bins failpoint-enable bin/dumpling
 	@make failpoint-disable
 	./tests/run.sh
 
@@ -115,4 +115,10 @@ tidy:
 	cd tools && GO111MODULE=on go mod tidy
 	git diff --exit-code go.mod go.sum tools/go.mod tools/go.sum
 
-.PHONY: build test integration_test tools failpoint-enable failpoint-disable check static lint tidy
+bins:
+	@which bin/tidb-server
+	@which bin/minio
+	@which bin/tidb-lightning
+	@which bin/sync_diff_inspector
+
+.PHONY: build test integration_test tools failpoint-enable failpoint-disable check static lint tidy bins
diff --git a/v4/export/dump.go b/v4/export/dump.go
index 9bfd089..eec79fd 100755
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -210,6 +210,12 @@ func (d *Dumper) Dump() (dumpErr error) {
 		}
 	})
 
+	// get estimate total count
+	err = d.getEstimateTotalRowsCount(tctx, metaConn)
+	if err != nil {
+		tctx.L().Error("fail to get estimate total count", zap.Error(err))
+	}
+
 	if conf.SQL == "" {
 		if err = d.dumpDatabases(metaConn, taskChan); err != nil {
 			return err
@@ -241,15 +247,24 @@ func (d *Dumper) startWriters(tctx *tcontext.Context, wg *errgroup.Group, taskCh
 		writer := NewWriter(tctx, int64(i), conf, conn, d.extStore)
 		writer.rebuildConnFn = rebuildConnFn
 		writer.setFinishTableCallBack(func(task Task) {
-			if td, ok := task.(*TaskTableData); ok {
+			if _, ok := task.(*TaskTableData); ok {
 				IncCounter(finishedTablesCounter, conf.Labels)
-				tctx.L().Debug("finished dumping table data",
-					zap.String("database", td.Meta.DatabaseName()),
-					zap.String("table", td.Meta.TableName()))
+				// FIXME: actually finishing the last chunk doesn't means this table is 'finished'.
+				//  We can call this table is 'finished' if all its chunks are finished.
+				//  Comment this log now to avoid ambiguity.
+				// tctx.L().Debug("finished dumping table data",
+				//	zap.String("database", td.Meta.DatabaseName()),
+				//	zap.String("table", td.Meta.TableName()))
 			}
 		})
 		writer.setFinishTaskCallBack(func(task Task) {
 			IncGauge(taskChannelCapacity, conf.Labels)
+			if td, ok := task.(*TaskTableData); ok {
+				tctx.L().Debug("finish dumping table data task",
+					zap.String("database", td.Meta.DatabaseName()),
+					zap.String("table", td.Meta.TableName()),
+					zap.Int("chunkIdx", td.ChunkIndex))
+			}
 		})
 		wg.Go(func() error {
 			return writer.run(taskChan)
@@ -375,7 +390,7 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 	cutoff := new(big.Int).Set(min)
 	totalChunks := estimatedChunks
 	if estimatedStep == 1 {
-		totalChunks = new(big.Int).Sub(max, min).Uint64()
+		totalChunks = new(big.Int).Sub(max, min).Uint64() + 1
 	}
 
 	selectField, selectLen, err := buildSelectField(conn, db, tbl, conf.CompleteInsert)
@@ -389,7 +404,10 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 	}
 
 	chunkIndex := 0
-	nullValueCondition := fmt.Sprintf("`%s` IS NULL OR ", escapeString(field))
+	nullValueCondition := ""
+	if conf.Where == "" {
+		nullValueCondition = fmt.Sprintf("`%s` IS NULL OR ", escapeString(field))
+	}
 	for max.Cmp(cutoff) >= 0 {
 		nextCutOff := new(big.Int).Add(cutoff, bigEstimatedStep)
 		where := fmt.Sprintf("%s(`%s` >= %d AND `%s` < %d)", nullValueCondition, escapeString(field), cutoff, escapeString(field), nextCutOff)
diff --git a/v4/export/ir.go b/v4/export/ir.go
index c776bcf..2b1d8d0 100644
--- a/v4/export/ir.go
+++ b/v4/export/ir.go
@@ -4,17 +4,18 @@ package export
 
 import (
 	"bytes"
-	"context"
 	"database/sql"
 	"strings"
 
 	"github.com/pingcap/errors"
+
+	tcontext "github.com/pingcap/dumpling/v4/context"
 )
 
 // TableDataIR is table data intermediate representation.
 // A table may be split into multiple TableDataIRs.
 type TableDataIR interface {
-	Start(context.Context, *sql.Conn) error
+	Start(*tcontext.Context, *sql.Conn) error
 	Rows() SQLRowIter
 	Close() error
 	RawRows() *sql.Rows
diff --git a/v4/export/ir_impl.go b/v4/export/ir_impl.go
index 207aac1..c9c216f 100644
--- a/v4/export/ir_impl.go
+++ b/v4/export/ir_impl.go
@@ -3,12 +3,14 @@
 package export
 
 import (
-	"context"
 	"database/sql"
 	"fmt"
 	"strings"
 
 	"github.com/pingcap/errors"
+	"go.uber.org/zap"
+
+	tcontext "github.com/pingcap/dumpling/v4/context"
 )
 
 // rowIter implements the SQLRowIter interface.
@@ -38,7 +40,7 @@ func (iter *rowIter) Decode(row RowReceiver) error {
 }
 
 func (iter *rowIter) Error() error {
-	return iter.rows.Err()
+	return errors.Trace(iter.rows.Err())
 }
 
 func (iter *rowIter) Next() {
@@ -91,21 +93,22 @@ func newTableData(query string, colLength int, needColTypes bool) *tableData {
 	}
 }
 
-func (td *tableData) Start(ctx context.Context, conn *sql.Conn) error {
-	rows, err := conn.QueryContext(ctx, td.query)
+func (td *tableData) Start(tctx *tcontext.Context, conn *sql.Conn) error {
+	tctx.L().Debug("try to start tableData", zap.String("query", td.query))
+	rows, err := conn.QueryContext(tctx, td.query)
 	if err != nil {
-		return errors.Trace(err)
+		return errors.Annotatef(err, "sql: %s", td.query)
 	}
 	if err = rows.Err(); err != nil {
-		return errors.Trace(err)
+		return errors.Annotatef(err, "sql: %s", td.query)
 	}
 	td.SQLRowIter = nil
 	td.rows = rows
-	ns, err := rows.Columns()
-	if err != nil {
-		return errors.Trace(err)
-	}
 	if td.needColTypes {
+		ns, err := rows.Columns()
+		if err != nil {
+			return errors.Trace(err)
+		}
 		td.colLen = len(ns)
 		td.colTypes = make([]string, 0, td.colLen)
 		colTps, err := rows.ColumnTypes()
diff --git a/v4/export/metrics.go b/v4/export/metrics.go
index b3713bf..649a5ee 100644
--- a/v4/export/metrics.go
+++ b/v4/export/metrics.go
@@ -13,6 +13,7 @@ var (
 	finishedSizeCounter            *prometheus.CounterVec
 	finishedRowsCounter            *prometheus.CounterVec
 	finishedTablesCounter          *prometheus.CounterVec
+	estimateTotalRowsCounter       *prometheus.CounterVec
 	writeTimeHistogram             *prometheus.HistogramVec
 	receiveWriteChunkTimeHistogram *prometheus.HistogramVec
 	errorCount                     *prometheus.CounterVec
@@ -33,6 +34,13 @@ func InitMetricsVector(labels prometheus.Labels) {
 			Name:      "finished_size",
 			Help:      "counter for dumpling finished file size",
 		}, labelNames)
+	estimateTotalRowsCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: "dumpling",
+			Subsystem: "dump",
+			Name:      "estimate_total_rows",
+			Help:      "estimate total rows for dumpling tables",
+		}, []string{})
 	finishedRowsCounter = prometheus.NewCounterVec(
 		prometheus.CounterOpts{
 			Namespace: "dumpling",
@@ -86,6 +94,7 @@ func RegisterMetrics(registry *prometheus.Registry) {
 	}
 	registry.MustRegister(finishedSizeCounter)
 	registry.MustRegister(finishedRowsCounter)
+	registry.MustRegister(estimateTotalRowsCounter)
 	registry.MustRegister(finishedTablesCounter)
 	registry.MustRegister(writeTimeHistogram)
 	registry.MustRegister(receiveWriteChunkTimeHistogram)
@@ -100,6 +109,7 @@ func RemoveLabelValuesWithTaskInMetrics(labels prometheus.Labels) {
 	}
 	finishedSizeCounter.Delete(labels)
 	finishedRowsCounter.Delete(labels)
+	estimateTotalRowsCounter.Delete(labels)
 	finishedTablesCounter.Delete(labels)
 	writeTimeHistogram.Delete(labels)
 	receiveWriteChunkTimeHistogram.Delete(labels)
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 398c492..7c93d1b 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -693,7 +693,12 @@ func pickupPossibleField(dbName, tableName string, db *sql.Conn, conf *Config) (
 }
 
 func estimateCount(tctx *tcontext.Context, dbName, tableName string, db *sql.Conn, field string, conf *Config) uint64 {
-	query := fmt.Sprintf("EXPLAIN SELECT `%s` FROM `%s`.`%s`", escapeString(field), escapeString(dbName), escapeString(tableName))
+	var query string
+	if strings.TrimSpace(field) == "*" || strings.TrimSpace(field) == "" {
+		query = fmt.Sprintf("EXPLAIN SELECT * FROM `%s`.`%s`", escapeString(dbName), escapeString(tableName))
+	} else {
+		query = fmt.Sprintf("EXPLAIN SELECT `%s` FROM `%s`.`%s`", escapeString(field), escapeString(dbName), escapeString(tableName))
+	}
 
 	if conf.Where != "" {
 		query += " WHERE "
diff --git a/v4/export/status.go b/v4/export/status.go
index 26048f4..3f93f7b 100644
--- a/v4/export/status.go
+++ b/v4/export/status.go
@@ -3,6 +3,7 @@
 package export
 
 import (
+	"database/sql"
 	"fmt"
 	"time"
 
@@ -32,10 +33,12 @@ func (d *Dumper) runLogProgress(tctx *tcontext.Context) {
 			completedTables := ReadCounter(finishedTablesCounter, conf.Labels)
 			finishedBytes := ReadCounter(finishedSizeCounter, conf.Labels)
 			finishedRows := ReadCounter(finishedRowsCounter, conf.Labels)
+			estimateTotalRows := ReadCounter(estimateTotalRowsCounter, conf.Labels)
 
 			tctx.L().Info("progress",
 				zap.String("tables", fmt.Sprintf("%.0f/%.0f (%.1f%%)", completedTables, totalTables, completedTables/totalTables*100)),
 				zap.String("finished rows", fmt.Sprintf("%.0f", finishedRows)),
+				zap.String("estimate total rows", fmt.Sprintf("%.0f", estimateTotalRows)),
 				zap.String("finished size", units.HumanSize(finishedBytes)),
 				zap.Float64("average speed(MiB/s)", (finishedBytes-lastBytes)/(1048576e-9*nanoseconds)),
 			)
@@ -57,3 +60,23 @@ func calculateTableCount(m DatabaseTables) int {
 	}
 	return cnt
 }
+
+func (d *Dumper) getEstimateTotalRowsCount(tctx *tcontext.Context, conn *sql.Conn) error {
+	conf := d.conf
+	var totalCount uint64
+	for db, tables := range conf.Tables {
+		for _, m := range tables {
+			if m.Type == TableTypeBase {
+				// get pk or uk for explain
+				field, err := pickupPossibleField(db, m.Name, conn, conf)
+				if err != nil {
+					return err
+				}
+				c := estimateCount(tctx, db, m.Name, conn, field, conf)
+				totalCount += c
+			}
+		}
+	}
+	estimateTotalRowsCounter.With(conf.Labels).Add(float64(totalCount))
+	return nil
+}
diff --git a/v4/export/writer.go b/v4/export/writer.go
index 1387251..7822df3 100644
--- a/v4/export/writer.go
+++ b/v4/export/writer.go
@@ -198,9 +198,10 @@ func (w *Writer) tryToWriteTableData(tctx *tcontext.Context, meta TableMeta, ir
 		return err
 	}
 
+	somethingIsWritten := false
 	for {
 		fileWriter, tearDown := buildInterceptFileWriter(tctx, w.extStorage, fileName, conf.CompressType)
-		err = format.WriteInsert(tctx, conf, meta, ir, fileWriter)
+		n, err := format.WriteInsert(tctx, conf, meta, ir, fileWriter)
 		tearDown(tctx)
 		if err != nil {
 			return err
@@ -210,6 +211,13 @@ func (w *Writer) tryToWriteTableData(tctx *tcontext.Context, meta TableMeta, ir
 			break
 		}
 
+		tctx.L().Debug("finish dumping table(chunk)",
+			zap.String("database", meta.DatabaseName()),
+			zap.String("table", meta.TableName()),
+			zap.Int("chunkIdx", curChkIdx),
+			zap.Uint64("total rows", n))
+		somethingIsWritten = true
+
 		if conf.FileSize == UnspecifiedSize {
 			break
 		}
@@ -218,6 +226,12 @@ func (w *Writer) tryToWriteTableData(tctx *tcontext.Context, meta TableMeta, ir
 			return err
 		}
 	}
+	if !somethingIsWritten {
+		tctx.L().Warn("no data written in table chunk",
+			zap.String("database", meta.DatabaseName()),
+			zap.String("table", meta.TableName()),
+			zap.Int("chunkIdx", curChkIdx))
+	}
 	return nil
 }
 
diff --git a/v4/export/writer_util.go b/v4/export/writer_util.go
index 0697f57..7a9f7f6 100755
--- a/v4/export/writer_util.go
+++ b/v4/export/writer_util.go
@@ -134,10 +134,10 @@ func WriteMeta(tctx *tcontext.Context, meta MetaIR, w storage.ExternalFileWriter
 }
 
 // WriteInsert writes TableDataIR to a storage.ExternalFileWriter in sql type
-func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR TableDataIR, w storage.ExternalFileWriter) error {
+func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR TableDataIR, w storage.ExternalFileWriter) (n uint64, err error) {
 	fileRowIter := tblIR.Rows()
 	if !fileRowIter.HasNext() {
-		return nil
+		return 0, fileRowIter.Error()
 	}
 
 	bf := pool.Get().(*bytes.Buffer)
@@ -173,7 +173,6 @@ func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR Tabl
 		counter               uint64
 		lastCounter           uint64
 		escapeBackslash       = cfg.EscapeBackslash
-		err                   error
 	)
 
 	selectedField := meta.SelectedField()
@@ -198,7 +197,7 @@ func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR Tabl
 			if selectedField != "" {
 				if err = fileRowIter.Decode(row); err != nil {
 					pCtx.L().Error("fail to scan from sql.Row", zap.Error(err))
-					return errors.Trace(err)
+					return counter, errors.Trace(err)
 				}
 				row.WriteToBuffer(bf, escapeBackslash)
 			} else {
@@ -207,7 +206,7 @@ func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR Tabl
 			counter++
 			wp.AddFileSize(uint64(bf.Len()-lastBfSize) + 2) // 2 is for ",\n" and ";\n"
 			failpoint.Inject("ChaosBrokenMySQLConn", func(_ failpoint.Value) {
-				failpoint.Return(errors.New("connection is closed"))
+				failpoint.Return(0, errors.New("connection is closed"))
 			})
 
 			fileRowIter.Next()
@@ -220,9 +219,9 @@ func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR Tabl
 			if bf.Len() >= lengthLimit {
 				select {
 				case <-pCtx.Done():
-					return pCtx.Err()
+					return counter, pCtx.Err()
 				case err = <-wp.errCh:
-					return err
+					return counter, err
 				case wp.input <- bf:
 					bf = pool.Get().(*bytes.Buffer)
 					if bfCap := bf.Cap(); bfCap < lengthLimit {
@@ -250,20 +249,24 @@ func WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR Tabl
 	}
 	close(wp.input)
 	<-wp.closed
-	summary.CollectSuccessUnit(summary.TotalBytes, 1, wp.finishedFileSize)
-	summary.CollectSuccessUnit("total rows", 1, counter)
 	AddCounter(finishedRowsCounter, cfg.Labels, float64(counter-lastCounter))
+	defer func() {
+		if err == nil {
+			summary.CollectSuccessUnit(summary.TotalBytes, 1, wp.finishedFileSize)
+			summary.CollectSuccessUnit("total rows", 1, counter)
+		}
+	}()
 	if err = fileRowIter.Error(); err != nil {
-		return errors.Trace(err)
+		return counter, errors.Trace(err)
 	}
-	return wp.Error()
+	return counter, wp.Error()
 }
 
 // WriteInsertInCsv writes TableDataIR to a storage.ExternalFileWriter in csv type
-func WriteInsertInCsv(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR TableDataIR, w storage.ExternalFileWriter) error {
+func WriteInsertInCsv(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR TableDataIR, w storage.ExternalFileWriter) (n uint64, err error) {
 	fileRowIter := tblIR.Rows()
 	if !fileRowIter.HasNext() {
-		return nil
+		return 0, fileRowIter.Error()
 	}
 
 	bf := pool.Get().(*bytes.Buffer)
@@ -297,7 +300,6 @@ func WriteInsertInCsv(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR
 		lastCounter     uint64
 		escapeBackslash = cfg.EscapeBackslash
 		selectedFields  = meta.SelectedField()
-		err             error
 	)
 
 	if !cfg.NoHeader && len(meta.ColumnNames()) != 0 && selectedFields != "" {
@@ -318,7 +320,7 @@ func WriteInsertInCsv(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR
 		if selectedFields != "" {
 			if err = fileRowIter.Decode(row); err != nil {
 				pCtx.L().Error("fail to scan from sql.Row", zap.Error(err))
-				return errors.Trace(err)
+				return counter, errors.Trace(err)
 			}
 			row.WriteToBufferInCsv(bf, escapeBackslash, opt)
 		}
@@ -329,9 +331,9 @@ func WriteInsertInCsv(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR
 		if bf.Len() >= lengthLimit {
 			select {
 			case <-pCtx.Done():
-				return pCtx.Err()
+				return counter, pCtx.Err()
 			case err = <-wp.errCh:
-				return err
+				return counter, err
 			case wp.input <- bf:
 				bf = pool.Get().(*bytes.Buffer)
 				if bfCap := bf.Cap(); bfCap < lengthLimit {
@@ -357,13 +359,17 @@ func WriteInsertInCsv(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR
 	}
 	close(wp.input)
 	<-wp.closed
-	summary.CollectSuccessUnit(summary.TotalBytes, 1, wp.finishedFileSize)
-	summary.CollectSuccessUnit("total rows", 1, counter)
+	defer func() {
+		if err == nil {
+			summary.CollectSuccessUnit(summary.TotalBytes, 1, wp.finishedFileSize)
+			summary.CollectSuccessUnit("total rows", 1, counter)
+		}
+	}()
 	AddCounter(finishedRowsCounter, cfg.Labels, float64(counter-lastCounter))
 	if err = fileRowIter.Error(); err != nil {
-		return errors.Trace(err)
+		return counter, errors.Trace(err)
 	}
-	return wp.Error()
+	return counter, wp.Error()
 }
 
 func write(tctx *tcontext.Context, writer storage.ExternalFileWriter, str string) error {
@@ -588,13 +594,13 @@ func (f FileFormat) Extension() string {
 }
 
 // WriteInsert writes TableDataIR to a storage.ExternalFileWriter in sql/csv type
-func (f FileFormat) WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR TableDataIR, w storage.ExternalFileWriter) error {
+func (f FileFormat) WriteInsert(pCtx *tcontext.Context, cfg *Config, meta TableMeta, tblIR TableDataIR, w storage.ExternalFileWriter) (uint64, error) {
 	switch f {
 	case FileFormatSQLText:
 		return WriteInsert(pCtx, cfg, meta, tblIR, w)
 	case FileFormatCSV:
 		return WriteInsertInCsv(pCtx, cfg, meta, tblIR, w)
 	default:
-		return errors.Errorf("unknown file format")
+		return 0, errors.Errorf("unknown file format")
 	}
 }


Test patch:
diff --git a/v4/export/test_util.go b/v4/export/test_util.go
index eca9bbf..62b249e 100644
--- a/v4/export/test_util.go
+++ b/v4/export/test_util.go
@@ -8,6 +8,8 @@ import (
 	"database/sql/driver"
 	"fmt"
 
+	tcontext "github.com/pingcap/dumpling/v4/context"
+
 	"github.com/DATA-DOG/go-sqlmock"
 )
 
@@ -82,7 +84,7 @@ func (m *mockTableIR) ShowCreateView() string {
 	return ""
 }
 
-func (m *mockTableIR) Start(_ context.Context, conn *sql.Conn) error {
+func (m *mockTableIR) Start(_ *tcontext.Context, conn *sql.Conn) error {
 	return nil
 }
 
diff --git a/v4/export/writer_util_test.go b/v4/export/writer_util_test.go
index d3b6920..ae6453e 100644
--- a/v4/export/writer_util_test.go
+++ b/v4/export/writer_util_test.go
@@ -65,7 +65,8 @@ func (s *testUtilSuite) TestWriteInsert(c *C) {
 	bf := storage.NewBufferWriter()
 
 	conf := configForWriteSQL(UnspecifiedSize, UnspecifiedSize)
-	err := WriteInsert(tcontext.Background(), conf, tableIR, tableIR, bf)
+	n, err := WriteInsert(tcontext.Background(), conf, tableIR, tableIR, bf)
+	c.Assert(n, Equals, uint64(4))
 	c.Assert(err, IsNil)
 	expected := "/*!40101 SET NAMES binary*/;\n" +
 		"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n" +
@@ -96,7 +97,8 @@ func (s *testUtilSuite) TestWriteInsertReturnsError(c *C) {
 	bf := storage.NewBufferWriter()
 
 	conf := configForWriteSQL(UnspecifiedSize, UnspecifiedSize)
-	err := WriteInsert(tcontext.Background(), conf, tableIR, tableIR, bf)
+	n, err := WriteInsert(tcontext.Background(), conf, tableIR, tableIR, bf)
+	c.Assert(n, Equals, uint64(3))
 	c.Assert(err, Equals, rowErr)
 	expected := "/*!40101 SET NAMES binary*/;\n" +
 		"/*!40014 SET FOREIGN_KEY_CHECKS=0*/;\n" +
@@ -121,7 +123,8 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 	// test nullValue
 	opt := &csvOption{separator: []byte(","), delimiter: doubleQuotationMark, nullValue: "\\N"}
 	conf := configForWriteCSV(true, opt)
-	err := WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	n, err := WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	c.Assert(n, Equals, uint64(4))
 	c.Assert(err, IsNil)
 	expected := "1,\"male\",\"bob@mail.com\",\"020-1234\",\\N\n" +
 		"2,\"female\",\"sarah@mail.com\",\"020-1253\",\"healthy\"\n" +
@@ -134,7 +137,8 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 	opt.delimiter = quotationMark
 	tableIR = newMockTableIR("test", "employee", data, nil, colTypes)
 	conf = configForWriteCSV(true, opt)
-	err = WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	n, err = WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	c.Assert(n, Equals, uint64(4))
 	c.Assert(err, IsNil)
 	expected = "1,'male','bob@mail.com','020-1234',\\N\n" +
 		"2,'female','sarah@mail.com','020-1253','healthy'\n" +
@@ -147,7 +151,8 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 	opt.separator = []byte(";")
 	tableIR = newMockTableIR("test", "employee", data, nil, colTypes)
 	conf = configForWriteCSV(true, opt)
-	err = WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	n, err = WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	c.Assert(n, Equals, uint64(4))
 	c.Assert(err, IsNil)
 	expected = "1;'male';'bob@mail.com';'020-1234';\\N\n" +
 		"2;'female';'sarah@mail.com';'020-1253';'healthy'\n" +
@@ -162,7 +167,8 @@ func (s *testUtilSuite) TestWriteInsertInCsv(c *C) {
 	tableIR = newMockTableIR("test", "employee", data, nil, colTypes)
 	tableIR.colNames = []string{"id", "gender", "email", "phone_number", "status"}
 	conf = configForWriteCSV(false, opt)
-	err = WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	n, err = WriteInsertInCsv(tcontext.Background(), conf, tableIR, tableIR, bf)
+	c.Assert(n, Equals, uint64(4))
 	c.Assert(err, IsNil)
 	expected = "maidma&;,?magenderma&;,?maemamailma&;,?maphone_numberma&;,?mastatusma\n" +
 		"1&;,?mamamalema&;,?mabob@mamail.comma&;,?ma020-1234ma&;,?\\N\n" +
@@ -188,7 +194,8 @@ func (s *testUtilSuite) TestSQLDataTypes(c *C) {
 		bf := storage.NewBufferWriter()
 
 		conf := configForWriteSQL(UnspecifiedSize, UnspecifiedSize)
-		err := WriteInsert(tcontext.Background(), conf, tableIR, tableIR, bf)
+		n, err := WriteInsert(tcontext.Background(), conf, tableIR, tableIR, bf)
+		c.Assert(n, Equals, uint64(1))
 		c.Assert(err, IsNil)
 		lines := strings.Split(bf.String(), "\n")
 		c.Assert(len(lines), Equals, 3)