Repository: pingcap/tidb-lightning
Pull request #491 (2020-11-24 12:05:52)

*: add method to check whether need local SST of table
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
in DBaaS environment resuming TiDB can't assure existence of local SST, so we should add a checking method.

### What is changed and how it works?
status between "chunk position > chunk start offset" and "index engine imported" have unrecoverable data in local SST, so check it

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Integration test

Side effects

 - Breaking backward compatibility (add method to interface)

Related changes

lance6716: ping @glorv @kennytm 
lance6716: /run-all-tests
lance6716: /run-all-tests

```
config 'tikv-importer.backend' value 'local' different from checkpoint value 'importer', please destroy all checkpoints and remove all restored tables and try again
```

don't know why test `local_backend` will generate importer backend checkpoint
lance6716: /run-all-tests 

tidb lightning encountered error:  restore table `cpeng`.`b` failed: pebble: could not open manifest file "MANIFEST-000001" for DB "/tmp/lightning_test_result/sorted/7553a9f1-d15d-58e9-9d34-99cd0ed0a85e": open /tmp/lightning_test_result/sorted/7553a9f1-d15d-58e9-9d34-99cd0ed0a85e/MANIFEST-000001: no such file or directory
lance6716: ping @glorv 
glorv: LGTM
kennytm: LGTM

Patch:
diff --git a/cmd/tidb-lightning-ctl/main.go b/cmd/tidb-lightning-ctl/main.go
index ceda277..b49bdcb 100644
--- a/cmd/tidb-lightning-ctl/main.go
+++ b/cmd/tidb-lightning-ctl/main.go
@@ -45,6 +45,7 @@ func run() error {
 		compact, flagFetchMode                      *bool
 		mode, flagImportEngine, flagCleanupEngine   *string
 		cpRemove, cpErrIgnore, cpErrDestroy, cpDump *string
+		localStoringTables                          *bool
 
 		fsUsage func()
 	)
@@ -69,6 +70,8 @@ func run() error {
 		cpErrDestroy = fs.String("checkpoint-error-destroy", "", "deletes imported data with table which has an error before (value can be 'all' or '`db`.`table`')")
 		cpDump = fs.String("checkpoint-dump", "", "dump the checkpoint information as two CSV files in the given folder")
 
+		localStoringTables = fs.Bool("check-local-storage", false, "show tables that are missing local intermediate files (value can be 'all' or '`db`.`table`')")
+
 		fsUsage = fs.Usage
 	}))
 
@@ -118,6 +121,9 @@ func run() error {
 	if len(*cpDump) != 0 {
 		return errors.Trace(checkpointDump(ctx, cfg, *cpDump))
 	}
+	if *localStoringTables {
+		return errors.Trace(getLocalStoringTables(ctx, cfg))
+	}
 
 	fsUsage()
 	return nil
@@ -308,6 +314,48 @@ func checkpointDump(ctx context.Context, cfg *config.Config, dumpFolder string)
 	return nil
 }
 
+func getLocalStoringTables(ctx context.Context, cfg *config.Config) (err2 error) {
+	var (
+		tables []string
+	)
+	defer func() {
+		if err2 == nil {
+			if len(tables) == 0 {
+				fmt.Fprintln(os.Stderr, "No table has lost intermediate files according to given config")
+			} else {
+				fmt.Fprintln(os.Stderr, "These tables are missing intermediate files:", tables)
+			}
+		}
+	}()
+
+	if cfg.TikvImporter.Backend != config.BackendLocal {
+		return nil
+	}
+	exist, err := checkpoints.IsCheckpointsDBExists(ctx, cfg)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	if !exist {
+		return nil
+	}
+	cpdb, err := checkpoints.OpenCheckpointsDB(ctx, cfg)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	defer cpdb.Close()
+
+	tableWithEngine, err := cpdb.GetLocalStoringTables(ctx)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	tables = make([]string, 0, len(tableWithEngine))
+	for tableName := range tableWithEngine {
+		tables = append(tables, tableName)
+	}
+
+	return nil
+}
+
 func unsafeCloseEngine(ctx context.Context, importer kv.Backend, engine string) (*kv.ClosedEngine, error) {
 	if index := strings.LastIndexByte(engine, ':'); index >= 0 {
 		tableName := engine[:index]
diff --git a/lightning/backend/local.go b/lightning/backend/local.go
index 2369b3d..83170b6 100644
--- a/lightning/backend/local.go
+++ b/lightning/backend/local.go
@@ -141,6 +141,15 @@ func (e *LocalFile) Cleanup(dataDir string) error {
 	return os.RemoveAll(dbPath)
 }
 
+// Exist checks if db folder existing (meta sometimes won't flush before lightning exit)
+func (e *LocalFile) Exist(dataDir string) error {
+	dbPath := filepath.Join(dataDir, e.Uuid.String())
+	if _, err := os.Stat(dbPath); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (e *LocalFile) getSizeProperties() (*sizeProperties, error) {
 	sstables, err := e.db.SSTables(pebble.WithProperties())
 	if err != nil {
diff --git a/lightning/checkpoints/checkpoints.go b/lightning/checkpoints/checkpoints.go
index d740245..fdc6cea 100644
--- a/lightning/checkpoints/checkpoints.go
+++ b/lightning/checkpoints/checkpoints.go
@@ -26,14 +26,13 @@ import (
 	"strings"
 	"sync"
 
-	"github.com/pingcap/tidb-lightning/lightning/config"
-
 	"github.com/joho/sqltocsv"
 	"github.com/pingcap/errors"
 	"go.uber.org/zap"
 	"modernc.org/mathutil"
 
 	"github.com/pingcap/tidb-lightning/lightning/common"
+	"github.com/pingcap/tidb-lightning/lightning/config"
 	"github.com/pingcap/tidb-lightning/lightning/log"
 	"github.com/pingcap/tidb-lightning/lightning/mydump"
 	verify "github.com/pingcap/tidb-lightning/lightning/verification"
@@ -472,6 +471,9 @@ type CheckpointsDB interface {
 	// MoveCheckpoints renames the checkpoint schema to include a suffix
 	// including the taskID (e.g. `tidb_lightning_checkpoints.1234567890.bak`).
 	MoveCheckpoints(ctx context.Context, taskID int64) error
+	// GetLocalStoringTables returns a map containing tables have engine files stored on local disk.
+	// currently only meaningful for local backend
+	GetLocalStoringTables(ctx context.Context) (map[string][]int32, error)
 	IgnoreErrorCheckpoint(ctx context.Context, tableName string) error
 	DestroyErrorCheckpoint(ctx context.Context, tableName string) ([]DestroyedTableCheckpoint, error)
 	DumpTables(ctx context.Context, csv io.Writer) error
@@ -505,6 +507,39 @@ func OpenCheckpointsDB(ctx context.Context, cfg *config.Config) (CheckpointsDB,
 	}
 }
 
+func IsCheckpointsDBExists(ctx context.Context, cfg *config.Config) (bool, error) {
+	if !cfg.Checkpoint.Enable {
+		return false, nil
+	}
+	switch cfg.Checkpoint.Driver {
+	case config.CheckpointDriverMySQL:
+		db, err := sql.Open("mysql", cfg.Checkpoint.DSN)
+		if err != nil {
+			return false, errors.Trace(err)
+		}
+		defer db.Close()
+		checkSQL := "SHOW DATABASES WHERE `DATABASE` = ?"
+		rows, err := db.QueryContext(ctx, checkSQL, cfg.Checkpoint.Schema)
+		if err != nil {
+			return false, errors.Trace(err)
+		}
+		defer rows.Close()
+		return rows.Next(), nil
+
+	case config.CheckpointDriverFile:
+		_, err := os.Stat(cfg.Checkpoint.DSN)
+		if err == nil {
+			return true, err
+		} else if os.IsNotExist(err) {
+			return false, nil
+		}
+		return false, errors.Trace(err)
+
+	default:
+		return false, errors.Errorf("Unknown checkpoint driver %s", cfg.Checkpoint.Driver)
+	}
+}
+
 // NullCheckpointsDB is a checkpoints database with no checkpoints.
 type NullCheckpointsDB struct{}
 
@@ -1130,6 +1165,9 @@ func (*NullCheckpointsDB) RemoveCheckpoint(context.Context, string) error {
 func (*NullCheckpointsDB) MoveCheckpoints(context.Context, int64) error {
 	return errors.Trace(cannotManageNullDB)
 }
+func (*NullCheckpointsDB) GetLocalStoringTables(context.Context) (map[string][]int32, error) {
+	return nil, nil
+}
 func (*NullCheckpointsDB) IgnoreErrorCheckpoint(context.Context, string) error {
 	return errors.Trace(cannotManageNullDB)
 }
@@ -1199,6 +1237,54 @@ func (cpdb *MySQLCheckpointsDB) MoveCheckpoints(ctx context.Context, taskID int6
 	return nil
 }
 
+func (cpdb *MySQLCheckpointsDB) GetLocalStoringTables(ctx context.Context) (map[string][]int32, error) {
+	var targetTables map[string][]int32
+
+	// lightning didn't check CheckpointStatusMaxInvalid before this function is called, so we skip invalid ones
+	// engines should exist if
+	// 1. table status is earlier than CheckpointStatusIndexImported, and
+	// 2. engine status is earlier than CheckpointStatusImported, and
+	// 3. chunk has been read
+	query := fmt.Sprintf(`
+		SELECT DISTINCT t.table_name, c.engine_id 
+		FROM %s.%s t, %s.%s c, %s.%s e 
+		WHERE t.table_name = c.table_name AND t.table_name = e.table_name AND c.engine_id = e.engine_id 
+			AND %d < t.status AND t.status < %d 
+			AND %d < e.status AND e.status < %d
+			AND c.pos > c.offset;`,
+		cpdb.schema, CheckpointTableNameTable, cpdb.schema, CheckpointTableNameChunk, cpdb.schema, CheckpointTableNameEngine,
+		CheckpointStatusMaxInvalid, CheckpointStatusIndexImported,
+		CheckpointStatusMaxInvalid, CheckpointStatusImported)
+
+	err := common.Retry("get local storing tables", log.L(), func() error {
+		targetTables = make(map[string][]int32)
+		rows, err := cpdb.db.QueryContext(ctx, query)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		defer rows.Close()
+		for rows.Next() {
+			var (
+				tableName string
+				engineID  int32
+			)
+			if err := rows.Scan(&tableName, &engineID); err != nil {
+				return errors.Trace(err)
+			}
+			targetTables[tableName] = append(targetTables[tableName], engineID)
+		}
+		if err := rows.Err(); err != nil {
+			return errors.Trace(err)
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	return targetTables, err
+}
+
 func (cpdb *MySQLCheckpointsDB) IgnoreErrorCheckpoint(ctx context.Context, tableName string) error {
 	var colName string
 	if tableName == "all" {
@@ -1398,6 +1484,36 @@ func (cpdb *FileCheckpointsDB) MoveCheckpoints(ctx context.Context, taskID int64
 	return errors.Trace(os.Rename(cpdb.path, newPath))
 }
 
+func (cpdb *FileCheckpointsDB) GetLocalStoringTables(_ context.Context) (map[string][]int32, error) {
+	cpdb.lock.Lock()
+	defer cpdb.lock.Unlock()
+
+	targetTables := make(map[string][]int32)
+
+	for tableName, tableModel := range cpdb.checkpoints.Checkpoints {
+		if tableModel.Status <= uint32(CheckpointStatusMaxInvalid) ||
+			tableModel.Status >= uint32(CheckpointStatusIndexImported) {
+			continue
+		}
+		for engineID, engineModel := range tableModel.Engines {
+			if engineModel.Status <= uint32(CheckpointStatusMaxInvalid) ||
+				engineModel.Status >= uint32(CheckpointStatusImported) {
+				continue
+			}
+
+			for _, chunkModel := range engineModel.Chunks {
+				if chunkModel.Pos > chunkModel.Offset {
+					targetTables[tableName] = append(targetTables[tableName], engineID)
+					break
+				}
+			}
+
+		}
+	}
+
+	return targetTables, nil
+}
+
 func (cpdb *FileCheckpointsDB) IgnoreErrorCheckpoint(_ context.Context, targetTableName string) error {
 	cpdb.lock.Lock()
 	defer cpdb.lock.Unlock()
diff --git a/lightning/checkpoints/glue_checkpoint.go b/lightning/checkpoints/glue_checkpoint.go
index a35f9ba..bb17f0c 100644
--- a/lightning/checkpoints/glue_checkpoint.go
+++ b/lightning/checkpoints/glue_checkpoint.go
@@ -53,6 +53,8 @@ type GlueCheckpointsDB struct {
 	schema         string
 }
 
+var _ CheckpointsDB = (*GlueCheckpointsDB)(nil)
+
 func NewGlueCheckpointsDB(ctx context.Context, se Session, f func() (Session, error), schemaName string) (*GlueCheckpointsDB, error) {
 	var escapedSchemaName strings.Builder
 	common.WriteMySQLIdentifier(&escapedSchemaName, schemaName)
@@ -220,9 +222,7 @@ func (g GlueCheckpointsDB) Get(ctx context.Context, tableName string) (*TableChe
 	}
 	defer se.Close()
 
-	var tableNameBuilder strings.Builder
-	common.EscapeMySQLSingleQuote(&tableNameBuilder, tableName)
-	tableName = tableNameBuilder.String()
+	tableName = common.InterpolateMySQLString(tableName)
 	err = Transact(ctx, "read checkpoint", se, logger, func(c context.Context, s Session) error {
 		// 1. Populate the engines.
 		sql := fmt.Sprintf(ReadEngineTemplate, g.schema, CheckpointTableNameEngine)
@@ -510,9 +510,7 @@ func (g GlueCheckpointsDB) RemoveCheckpoint(ctx context.Context, tableName strin
 			return err
 		})
 	}
-	var tableNameBuilder strings.Builder
-	common.EscapeMySQLSingleQuote(&tableNameBuilder, tableName)
-	tableName = tableNameBuilder.String()
+	tableName = common.InterpolateMySQLString(tableName)
 	deleteChunkQuery := fmt.Sprintf(DeleteCheckpointRecordTemplate, g.schema, CheckpointTableNameChunk)
 	deleteChunkQuery = strings.ReplaceAll(deleteChunkQuery, "?", tableName)
 	deleteEngineQuery := fmt.Sprintf(DeleteCheckpointRecordTemplate, g.schema, CheckpointTableNameEngine)
@@ -564,6 +562,56 @@ func (g GlueCheckpointsDB) MoveCheckpoints(ctx context.Context, taskID int64) er
 	return nil
 }
 
+func (g GlueCheckpointsDB) GetLocalStoringTables(ctx context.Context) (map[string][]int32, error) {
+	se, err := g.getSessionFunc()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	defer se.Close()
+
+	var targetTables map[string][]int32
+
+	// lightning didn't check CheckpointStatusMaxInvalid before this function is called, so we skip invalid ones
+	// engines should exist if
+	// 1. table status is earlier than CheckpointStatusIndexImported, and
+	// 2. engine status is earlier than CheckpointStatusImported, and
+	// 3. chunk has been read
+	query := fmt.Sprintf(`
+		SELECT DISTINCT t.table_name, c.engine_id 
+		FROM %s.%s t, %s.%s c, %s.%s e 
+		WHERE t.table_name = c.table_name AND t.table_name = e.table_name AND c.engine_id = e.engine_id 
+			AND %d < t.status AND t.status < %d 
+			AND %d < e.status AND e.status < %d
+			AND c.pos > c.offset;`,
+		g.schema, CheckpointTableNameTable, g.schema, CheckpointTableNameChunk, g.schema, CheckpointTableNameEngine,
+		CheckpointStatusMaxInvalid, CheckpointStatusIndexImported,
+		CheckpointStatusMaxInvalid, CheckpointStatusImported)
+
+	err = common.Retry("get local storing tables", log.L(), func() error {
+		targetTables = make(map[string][]int32)
+		rs, err := se.Execute(ctx, query)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		rows, err := drainFirstRecordSet(ctx, rs)
+		if err != nil {
+			return errors.Trace(err)
+		}
+
+		for _, row := range rows {
+			tableName := row.GetString(0)
+			engineID := int32(row.GetInt64(1))
+			targetTables[tableName] = append(targetTables[tableName], engineID)
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	return targetTables, err
+}
+
 func (g GlueCheckpointsDB) IgnoreErrorCheckpoint(ctx context.Context, tableName string) error {
 	logger := log.With(zap.String("table", tableName))
 	se, err := g.getSessionFunc()
@@ -581,9 +629,7 @@ func (g GlueCheckpointsDB) IgnoreErrorCheckpoint(ctx context.Context, tableName
 		colName = "table_name"
 	}
 
-	var tableNameBuilder strings.Builder
-	common.EscapeMySQLSingleQuote(&tableNameBuilder, tableName)
-	tableName = tableNameBuilder.String()
+	tableName = common.InterpolateMySQLString(tableName)
 
 	engineQuery := fmt.Sprintf(`
 		UPDATE %s.%s SET status = %d WHERE %s = %s AND status <= %d;
@@ -622,9 +668,7 @@ func (g GlueCheckpointsDB) DestroyErrorCheckpoint(ctx context.Context, tableName
 		aliasedColName = "t.table_name"
 	}
 
-	var tableNameBuilder strings.Builder
-	common.EscapeMySQLSingleQuote(&tableNameBuilder, tableName)
-	tableName = tableNameBuilder.String()
+	tableName = common.InterpolateMySQLString(tableName)
 
 	selectQuery := fmt.Sprintf(`
 		SELECT
@@ -726,3 +770,25 @@ func Transact(ctx context.Context, purpose string, s Session, logger log.Logger,
 		return nil
 	})
 }
+
+// TODO: will use drainFirstRecordSet to reduce repeat in GlueCheckpointsDB later
+func drainFirstRecordSet(ctx context.Context, rss []sqlexec.RecordSet) ([]chunk.Row, error) {
+	if len(rss) != 1 {
+		return nil, errors.New("given result set doesn't have length 1")
+	}
+	rs := rss[0]
+	var rows []chunk.Row
+	req := rs.NewChunk()
+	for {
+		err := rs.Next(ctx, req)
+		if err != nil || req.NumRows() == 0 {
+			rs.Close()
+			return rows, err
+		}
+		iter := chunk.NewIterator4Chunk(req)
+		for r := iter.Begin(); r != iter.End(); r = iter.Next() {
+			rows = append(rows, r)
+		}
+		req = chunk.Renew(req, 1024)
+	}
+}
diff --git a/lightning/common/util.go b/lightning/common/util.go
index 3265990..bcd5232 100644
--- a/lightning/common/util.go
+++ b/lightning/common/util.go
@@ -269,7 +269,8 @@ func WriteMySQLIdentifier(builder *strings.Builder, identifier string) {
 	builder.WriteByte('`')
 }
 
-func EscapeMySQLSingleQuote(builder *strings.Builder, s string) {
+func InterpolateMySQLString(s string) string {
+	var builder strings.Builder
 	builder.Grow(len(s) + 2)
 	builder.WriteByte('\'')
 	for i := 0; i < len(s); i++ {
@@ -281,6 +282,7 @@ func EscapeMySQLSingleQuote(builder *strings.Builder, s string) {
 		}
 	}
 	builder.WriteByte('\'')
+	return builder.String()
 }
 
 // GetJSON fetches a page and parses it as JSON. The parsed result will be
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index b4c6602..dae4a4c 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -224,6 +224,10 @@ func NewRestoreControllerWithPauser(
 		if err != nil {
 			return nil, errors.Annotate(err, "build local backend failed")
 		}
+		err = verifyLocalFile(ctx, cpdb, cfg.TikvImporter.SortedKVDir)
+		if err != nil {
+			return nil, err
+		}
 	default:
 		return nil, errors.New("unknown backend: " + cfg.TikvImporter.Backend)
 	}
@@ -643,6 +647,28 @@ func verifyCheckpoint(cfg *config.Config, taskCp *TaskCheckpoint) error {
 	return nil
 }
 
+// for local backend, we should check if local SST exists in disk, otherwise we'll lost data
+func verifyLocalFile(ctx context.Context, cpdb CheckpointsDB, dir string) error {
+	targetTables, err := cpdb.GetLocalStoringTables(ctx)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	for tableName, engineIDs := range targetTables {
+		for _, engineID := range engineIDs {
+			_, eID := kv.MakeUUID(tableName, engineID)
+			file := kv.LocalFile{Uuid: eID}
+			err := file.Exist(dir)
+			if err != nil {
+				log.L().Error("can't find local file",
+					zap.String("table name", tableName),
+					zap.Int32("engine ID", engineID))
+				return errors.Trace(err)
+			}
+		}
+	}
+	return nil
+}
+
 func (rc *RestoreController) estimateChunkCountIntoMetrics(ctx context.Context) error {
 	estimatedChunkCount := 0.0
 	for _, dbMeta := range rc.dbMetas {
@@ -1242,11 +1268,10 @@ func (t *TableRestore) restoreEngines(ctx context.Context, rc *RestoreController
 			if !rc.isLocalBackend() {
 				rc.postProcessLock.Lock()
 			}
-			err = t.importKV(ctx, closedIndexEngine)
+			err = t.importKV(ctx, closedIndexEngine, rc, indexEngineID)
 			if !rc.isLocalBackend() {
 				rc.postProcessLock.Unlock()
 			}
-			rc.saveStatusCheckpoint(t.tableName, indexEngineID, err, CheckpointStatusImported)
 		}
 
 		failpoint.Inject("FailBeforeIndexEngineImported", func() {
@@ -1421,11 +1446,10 @@ func (t *TableRestore) importEngine(
 	if !rc.isLocalBackend() {
 		rc.postProcessLock.Lock()
 	}
-	err := t.importKV(ctx, closedEngine)
+	err := t.importKV(ctx, closedEngine, rc, engineID)
 	if !rc.isLocalBackend() {
 		rc.postProcessLock.Unlock()
 	}
-	rc.saveStatusCheckpoint(t.tableName, engineID, err, CheckpointStatusImported)
 	if err != nil {
 		return errors.Trace(err)
 	}
@@ -1928,10 +1952,16 @@ func getColumnNames(tableInfo *model.TableInfo, permutation []int) []string {
 	return names
 }
 
-func (tr *TableRestore) importKV(ctx context.Context, closedEngine *kv.ClosedEngine) error {
+func (tr *TableRestore) importKV(
+	ctx context.Context,
+	closedEngine *kv.ClosedEngine,
+	rc *RestoreController,
+	engineID int32,
+) error {
 	task := closedEngine.Logger().Begin(zap.InfoLevel, "import and cleanup engine")
 
 	err := closedEngine.Import(ctx)
+	rc.saveStatusCheckpoint(tr.tableName, engineID, err, CheckpointStatusImported)
 	if err == nil {
 		closedEngine.Cleanup(ctx)
 	}
@@ -2082,8 +2112,15 @@ func (cr *chunkRestore) deliverLoop(
 			panic("forcing failure due to FailAfterWriteRows")
 		})
 		// TODO: for local backend, we may save checkpoint more frequently, e.g. after writen
-		//10GB kv pairs to data engine, we can do a flush for both data & index engine, then we
+		// 10GB kv pairs to data engine, we can do a flush for both data & index engine, then we
 		// can safely update current checkpoint.
+
+		failpoint.Inject("LocalBackendSaveCheckpoint", func() {
+			if !rc.isLocalBackend() && (dataChecksum.SumKVS() != 0 || indexChecksum.SumKVS() != 0) {
+				// No need to save checkpoint if nothing was delivered.
+				saveCheckpoint(rc, t, engineID, cr.chunk)
+			}
+		})
 	}
 
 	return


Test patch:
diff --git a/lightning/common/util_test.go b/lightning/common/util_test.go
index ca022cf..de392d3 100644
--- a/lightning/common/util_test.go
+++ b/lightning/common/util_test.go
@@ -204,3 +204,9 @@ func (s *utilSuite) TestStringSliceEqual(c *C) {
 	c.Assert(common.StringSliceEqual([]string{"a", "b", "c"}, []string{"a", "b"}), IsFalse)
 	c.Assert(common.StringSliceEqual([]string{"a", "x", "y"}, []string{"a", "y", "x"}), IsFalse)
 }
+
+func (s *utilSuite) TestInterpolateMySQLString(c *C) {
+	c.Assert(common.InterpolateMySQLString("123"), Equals, "'123'")
+	c.Assert(common.InterpolateMySQLString("1'23"), Equals, "'1''23'")
+	c.Assert(common.InterpolateMySQLString("1'2''3"), Equals, "'1''2''''3'")
+}
diff --git a/lightning/restore/restore_test.go b/lightning/restore/restore_test.go
index c552c15..7ba75b6 100644
--- a/lightning/restore/restore_test.go
+++ b/lightning/restore/restore_test.go
@@ -724,6 +724,13 @@ func (s *tableRestoreSuite) TestImportKVSuccess(c *C) {
 	defer controller.Finish()
 	mockBackend := mock.NewMockBackend(controller)
 	importer := kv.MakeBackend(mockBackend)
+	chptCh := make(chan saveCp)
+	defer close(chptCh)
+	rc := &RestoreController{saveCpCh: chptCh}
+	go func() {
+		for range chptCh {
+		}
+	}()
 
 	ctx := context.Background()
 	engineUUID := uuid.New()
@@ -740,7 +747,7 @@ func (s *tableRestoreSuite) TestImportKVSuccess(c *C) {
 
 	closedEngine, err := importer.UnsafeCloseEngineWithUUID(ctx, "tag", engineUUID)
 	c.Assert(err, IsNil)
-	err = s.tr.importKV(ctx, closedEngine)
+	err = s.tr.importKV(ctx, closedEngine, rc, 1)
 	c.Assert(err, IsNil)
 }
 
@@ -749,6 +756,13 @@ func (s *tableRestoreSuite) TestImportKVFailure(c *C) {
 	defer controller.Finish()
 	mockBackend := mock.NewMockBackend(controller)
 	importer := kv.MakeBackend(mockBackend)
+	chptCh := make(chan saveCp)
+	defer close(chptCh)
+	rc := &RestoreController{saveCpCh: chptCh}
+	go func() {
+		for range chptCh {
+		}
+	}()
 
 	ctx := context.Background()
 	engineUUID := uuid.New()
@@ -762,7 +776,7 @@ func (s *tableRestoreSuite) TestImportKVFailure(c *C) {
 
 	closedEngine, err := importer.UnsafeCloseEngineWithUUID(ctx, "tag", engineUUID)
 	c.Assert(err, IsNil)
-	err = s.tr.importKV(ctx, closedEngine)
+	err = s.tr.importKV(ctx, closedEngine, rc, 1)
 	c.Assert(err, ErrorMatches, "fake import error.*")
 }
 
diff --git a/tests/_utils/run_lightning_ctl b/tests/_utils/run_lightning_ctl
index b972eef..8f081c0 100755
--- a/tests/_utils/run_lightning_ctl
+++ b/tests/_utils/run_lightning_ctl
@@ -25,6 +25,7 @@ bin/tidb-lightning-ctl.test -test.coverprofile="$TEST_DIR/cov.ctl.$TEST_NAME.$$.
     --pd-urls '127.0.0.1:2379' \
     -d "tests/$TEST_NAME/data" \
     --importer '127.0.0.1:8808' \
+    --sorted-kv-dir "$TEST_DIR/sorted" \
     --enable-checkpoint=0 \
     --check-requirements=0 \
     "$@"
diff --git a/tests/local_backend/config.toml b/tests/local_backend/config.toml
index eeed351..46ca06e 100644
--- a/tests/local_backend/config.toml
+++ b/tests/local_backend/config.toml
@@ -4,6 +4,7 @@ table-concurrency = 1
 [checkpoint]
 enable = true
 driver = "file"
+schema = "tidb_lightning_checkpoint_local_backend_test"
 
 [tikv-importer]
 send-kv-pairs = 2
diff --git a/tests/local_backend/file.toml b/tests/local_backend/file.toml
new file mode 100644
index 0000000..ff263c0
--- /dev/null
+++ b/tests/local_backend/file.toml
@@ -0,0 +1,17 @@
+[lightning]
+table-concurrency = 1
+
+[checkpoint]
+enable = true
+driver = "file"
+schema = "tidb_lightning_checkpoint_local_backend_test"
+
+[tikv-importer]
+send-kv-pairs = 2
+
+[mydumper]
+batch-size = 50 # force splitting the data into 4 batches
+
+[[black-white-list.do-tables]]
+db-name = "cpeng"
+tbl-name = "a"
\ No newline at end of file
diff --git a/tests/local_backend/mysql.toml b/tests/local_backend/mysql.toml
new file mode 100644
index 0000000..6e80272
--- /dev/null
+++ b/tests/local_backend/mysql.toml
@@ -0,0 +1,17 @@
+[lightning]
+table-concurrency = 1
+
+[checkpoint]
+enable = true
+driver = "mysql"
+schema = "tidb_lightning_checkpoint_local_backend_test"
+
+[tikv-importer]
+send-kv-pairs = 2
+
+[mydumper]
+batch-size = 50 # force splitting the data into 4 batches
+
+[[black-white-list.do-tables]]
+db-name = "cpeng"
+tbl-name = "a"
\ No newline at end of file
diff --git a/tests/local_backend/run.sh b/tests/local_backend/run.sh
index 569c1d9..96bb85d 100755
--- a/tests/local_backend/run.sh
+++ b/tests/local_backend/run.sh
@@ -21,7 +21,7 @@ ENGINE_COUNT=6
 
 # First, verify that inject with not leader error is fine.
 rm -f "$TEST_DIR/lightning-local.log"
-rm -f "/tmp/tidb_lightning_checkpoint.pb"
+rm -f "/tmp/tidb_lightning_checkpoint_local_backend_test.pb"
 run_sql 'DROP DATABASE IF EXISTS cpeng;'
 export GO_FAILPOINTS='github.com/pingcap/tidb-lightning/lightning/backend/local/FailIngestMeta=return("notleader")'
 
@@ -38,7 +38,7 @@ check_contains 'sum(c): 46'
 
 # Now, verify it works with epoch not match as well.
 run_sql 'DROP DATABASE cpeng;'
-rm -f "/tmp/tidb_lightning_checkpoint.pb"
+rm -f "/tmp/tidb_lightning_checkpoint_local_backend_test.pb"
 
 export GO_FAILPOINTS='github.com/pingcap/tidb-lightning/lightning/backend/local/FailIngestMeta=return("epochnotmatch")'
 
@@ -55,7 +55,7 @@ check_contains 'sum(c): 46'
 
 # Now, verify it works with checkpoints as well.
 run_sql 'DROP DATABASE cpeng;'
-rm -f "/tmp/tidb_lightning_checkpoint.pb"
+rm -f "/tmp/tidb_lightning_checkpoint_local_backend_test.pb"
 
 set +e
 export GO_FAILPOINTS='github.com/pingcap/tidb-lightning/lightning/restore/FailBeforeDataEngineImported=return'
@@ -77,3 +77,47 @@ check_contains 'sum(c): 10'
 run_sql 'SELECT count(*), sum(c) FROM cpeng.b'
 check_contains 'count(*): 4'
 check_contains 'sum(c): 46'
+
+# Verify GetLocalStoringTables works
+# failpoint works for per table not task, so we limit this test to task that allow one table
+for ckpt in mysql file; do
+  run_sql 'DROP DATABASE IF EXISTS cpeng;'
+  run_sql 'DROP DATABASE IF EXISTS tidb_lightning_checkpoint_local_backend_test'
+  rm -f "/tmp/tidb_lightning_checkpoint_local_backend_test.pb"
+  
+  # before chunk pos is updated, local files could handle lost
+  set +e
+  export GO_FAILPOINTS="github.com/pingcap/tidb-lightning/lightning/restore/FailAfterWriteRows=return"
+  run_lightning --backend local --enable-checkpoint=1 --log-file "$TEST_DIR/lightning-local.log" --config "tests/$TEST_NAME/$ckpt.toml"
+  set -e
+  run_lightning_ctl --check-local-storage \
+    --backend local \
+    --enable-checkpoint=1 \
+    --config=tests/$TEST_NAME/$ckpt.toml >$TEST_DIR/lightning_ctl.output 2>&1
+  grep -Fq "No table has lost intermediate files according to given config" $TEST_DIR/lightning_ctl.output
+  
+  # when position of chunk file doesn't equal to offset, intermediate file should exist
+  set +e
+  export GO_FAILPOINTS="github.com/pingcap/tidb-lightning/lightning/restore/LocalBackendSaveCheckpoint=return;github.com/pingcap/tidb-lightning/lightning/restore/FailIfImportedChunk=return(1)"
+  run_lightning --backend local --enable-checkpoint=1 --log-file "$TEST_DIR/lightning-local.log" --config "tests/$TEST_NAME/$ckpt.toml"
+  set -e
+  run_lightning_ctl --check-local-storage \
+    --backend local \
+    --enable-checkpoint=1 \
+    --config=tests/$TEST_NAME/$ckpt.toml >$TEST_DIR/lightning_ctl.output 2>&1
+  grep -Eq "These tables are missing intermediate files: \[.+\]" $TEST_DIR/lightning_ctl.output
+  # don't distinguish whole sort-kv directory missing and table's directory missing for now
+  ls -lA $TEST_DIR/sorted
+  
+  # after index engine is imported, local file could handle lost
+  set +e
+  export GO_FAILPOINTS="github.com/pingcap/tidb-lightning/lightning/restore/FailIfIndexEngineImported=return(1)"
+  run_lightning --backend local --enable-checkpoint=1 --log-file "$TEST_DIR/lightning-local.log" --config "tests/$TEST_NAME/$ckpt.toml"
+  set -e
+  run_lightning_ctl --check-local-storage \
+    --backend local \
+    --enable-checkpoint=1 \
+    --config=tests/$TEST_NAME/$ckpt.toml >$TEST_DIR/lightning_ctl.output 2>&1
+  grep -Fq "No table has lost intermediate files according to given config" $TEST_DIR/lightning_ctl.output
+done
+rm -r $TEST_DIR/sorted
\ No newline at end of file