lightning: support extracting source/schema/table name to specific column

br/pkg/lightning/config/config.go

@@ -454,6 +454,7 @@ type MydumperRuntime struct {
 	ReadBlockSize    ByteSize         `toml:"read-block-size" json:"read-block-size"`
 	BatchSize        ByteSize         `toml:"batch-size" json:"batch-size"`
 	BatchImportRatio float64          `toml:"batch-import-ratio" json:"batch-import-ratio"`
+	SourceID         string           `toml:"source-id" json:"source-id"`
 	SourceDir        string           `toml:"data-source-dir" json:"data-source-dir"`
 	CharacterSet     string           `toml:"character-set" json:"character-set"`
 	CSV              CSVConfig        `toml:"csv" json:"csv"`

br/pkg/lightning/mydump/loader.go

@@ -82,6 +82,7 @@ type SourceFileMeta struct {
 	Compression Compression
 	SortKey     string
 	FileSize    int64
+	ExtendData  ExtendColumnData
 }
 
 // NewMDTableMeta creates an Mydumper table meta with specified character set.
@@ -167,6 +168,7 @@ type MDLoader struct {
 }
 
 type mdLoaderSetup struct {
+	sourceID      string
 	loader        *MDLoader
 	dbSchemas     []FileInfo
 	tableSchemas  []FileInfo
@@ -245,6 +247,7 @@ func NewMyDumpLoaderWithStore(ctx context.Context, cfg *config.Config, store sto
 	}
 
 	setup := mdLoaderSetup{
+		sourceID:      cfg.Mydumper.SourceID,
 		loader:        mdl,
 		dbIndexMap:    make(map[string]int),
 		tableIndexMap: make(map[filter.Table]int),
@@ -289,6 +292,12 @@ type FileInfo struct {
 	FileMeta  SourceFileMeta
 }
 
+// ExtendColumnData contains the extended column names and values information for a table.
+type ExtendColumnData struct {
+	Columns []string
+	Values  []string
+}
+
 // setup the `s.loader.dbs` slice by scanning all *.sql files inside `dir`.
 //
 // The database and tables are inserted in a consistent order, so creating an
@@ -488,6 +497,13 @@ func (s *mdLoaderSetup) route() error {
 				knownDBNames[targetDB] = newInfo
 			}
 			arr[i].TableName = filter.Table{Schema: targetDB, Name: targetTable}
+			extendCols, extendVals := r.FetchExtendColumn(rawDB, rawTable, s.sourceID)
+			if len(extendCols) > 0 {
+				arr[i].FileMeta.ExtendData = ExtendColumnData{
+					Columns: extendCols,
+					Values:  extendVals,
+				}
+			}
 		}
 		return nil
 	}

br/pkg/lightning/mydump/loader_test.go

@@ -990,3 +990,66 @@ func TestMaxScanFilesOption(t *testing.T) {
 	tbl = dbMeta.Tables[0]
 	require.Equal(t, maxScanFilesCount-2, len(tbl.DataFiles))
 }
+
+func TestExternalDataRoutes(t *testing.T) {
+	s := newTestMydumpLoaderSuite(t)
+
+	s.touch(t, "test_1-schema-create.sql")
+	s.touch(t, "test_1.t1-schema.sql")
+	s.touch(t, "test_1.t1.sql")
+	s.touch(t, "test_2-schema-create.sql")
+	s.touch(t, "test_2.t2-schema.sql")
+	s.touch(t, "test_2.t2.sql")
+	s.touch(t, "test_3-schema-create.sql")
+	s.touch(t, "test_3.t1-schema.sql")
+	s.touch(t, "test_3.t1.sql")
+	s.touch(t, "test_3.t3-schema.sql")
+	s.touch(t, "test_3.t3.sql")
+
+	s.cfg.Mydumper.SourceID = "mysql-01"
+	s.cfg.Routes = []*router.TableRule{
+		{
+			TableExtractor: &router.TableExtractor{
+				TargetColumn: "c_table",
+				TableRegexp:  "t(.*)",
+			},
+			SchemaExtractor: &router.SchemaExtractor{
+				TargetColumn: "c_schema",
+				SchemaRegexp: "test_(.*)",
+			},
+			SourceExtractor: &router.SourceExtractor{
+				TargetColumn: "c_source",
+				SourceRegexp: "mysql-(.*)",
+			},
+			SchemaPattern: "test_*",
+			TablePattern:  "t*",
+			TargetSchema:  "test",
+			TargetTable:   "t",
+		},
+	}
+
+	mdl, err := md.NewMyDumpLoader(context.Background(), s.cfg)
+
+	require.NoError(t, err)
+	var database *md.MDDatabaseMeta
+	for _, db := range mdl.GetDatabases() {
+		if db.Name == "test" {
+			require.Nil(t, database)
+			database = db
+		}
+	}
+	require.NotNil(t, database)
+	require.Len(t, database.Tables, 1)
+	require.Len(t, database.Tables[0].DataFiles, 4)
+	expectExtendCols := []string{"c_table", "c_schema", "c_source"}
+	expectedExtendVals := [][]string{
+		{"1", "1", "01"},
+		{"2", "2", "01"},
+		{"1", "3", "01"},
+		{"3", "3", "01"},
+	}
+	for i, fileInfo := range database.Tables[0].DataFiles {
+		require.Equal(t, expectExtendCols, fileInfo.FileMeta.ExtendData.Columns)
+		require.Equal(t, expectedExtendVals[i], fileInfo.FileMeta.ExtendData.Values)
+	}
+}

br/pkg/lightning/mydump/region.go

@@ -40,9 +40,10 @@ const (
 type TableRegion struct {
 	EngineID int32
 
-	DB       string
-	Table    string
-	FileMeta SourceFileMeta
+	DB         string
+	Table      string
+	FileMeta   SourceFileMeta
+	ExtendData ExtendColumnData
 
 	Chunk Chunk
 }

br/pkg/lightning/mydump/router.go

@@ -184,6 +184,11 @@ func NewFileRouter(cfg []*config.FileRouteRule, logger log.Logger) (FileRouter,
 	return chainRouters(res), nil
 }
 
+// NewDefaultFileRouter creates a new file router with the default file route rules.
+func NewDefaultFileRouter(logger log.Logger) (FileRouter, error) {
+	return NewFileRouter(defaultFileRouteRules, logger)
+}
+
 // RegexRouter is a `FileRouter` implement that apply specific regex pattern to filepath.
 // if regex pattern match, then each extractors with capture the matched regexp pattern and
 // set value to target field in `RouteResult`

br/pkg/lightning/restore/BUILD.bazel

@@ -62,6 +62,8 @@ go_library(
         "//util/engine",
         "//util/mathutil",
         "//util/mock",
+        "//util/regexpr-router",
+        "//util/set",
         "@com_github_coreos_go_semver//semver",
         "@com_github_docker_go_units//:go-units",
         "@com_github_go_sql_driver_mysql//:mysql",
@@ -140,6 +142,7 @@ go_test(
         "//util/mock",
         "//util/promutil",
         "//util/table-filter",
+        "//util/table-router",
         "@com_github_data_dog_go_sqlmock//:go-sqlmock",
         "@com_github_docker_go_units//:go-units",
         "@com_github_go_sql_driver_mysql//:mysql",

br/pkg/lightning/restore/chunk_restore_test.go

@@ -36,7 +36,12 @@ import (
 	"github.com/pingcap/tidb/br/pkg/lightning/worker"
 	"github.com/pingcap/tidb/br/pkg/mock"
 	"github.com/pingcap/tidb/br/pkg/storage"
+	"github.com/pingcap/tidb/ddl"
+	"github.com/pingcap/tidb/parser"
+	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/types"
+	tmock "github.com/pingcap/tidb/util/mock"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 )
@@ -273,6 +278,65 @@ func (s *chunkRestoreSuite) TestEncodeLoop() {
 	require.Equal(s.T(), s.cr.chunk.Chunk.EndOffset, kvs[0].offset)
 }
 
+func (s *chunkRestoreSuite) TestEncodeLoopWithExtendData() {
+	ctx := context.Background()
+	kvsCh := make(chan []deliveredKVs, 2)
+	deliverCompleteCh := make(chan deliverResult)
+
+	p := parser.New()
+	se := tmock.NewContext()
+
+	lastTi := s.tr.tableInfo
+	defer func() {
+		s.tr.tableInfo = lastTi
+	}()
+
+	node, err := p.ParseOneStmt("CREATE TABLE `t1` (`c1` varchar(5) NOT NULL, `c_table` varchar(5), `c_schema` varchar(5), `c_source` varchar(5))", "utf8mb4", "utf8mb4_bin")
+	require.NoError(s.T(), err)
+	tableInfo, err := ddl.MockTableInfo(se, node.(*ast.CreateTableStmt), int64(1))
+	require.NoError(s.T(), err)
+	tableInfo.State = model.StatePublic
+
+	schema := "test_1"
+	tb := "t1"
+	ti := &checkpoints.TidbTableInfo{
+		ID:   tableInfo.ID,
+		DB:   schema,
+		Name: tb,
+		Core: tableInfo,
+	}
+	s.tr.tableInfo = ti
+	s.cr.chunk.FileMeta.ExtendData = mydump.ExtendColumnData{
+		Columns: []string{"c_table", "c_schema", "c_source"},
+		Values:  []string{"1", "1", "01"},
+	}
+	defer func() {
+		s.cr.chunk.FileMeta.ExtendData = mydump.ExtendColumnData{}
+	}()
+
+	kvEncoder, err := kv.NewTableKVEncoder(s.tr.encTable, &kv.SessionOptions{
+		SQLMode:   s.cfg.TiDB.SQLMode,
+		Timestamp: 1234567895,
+	}, nil, log.L())
+	require.NoError(s.T(), err)
+	cfg := config.NewConfig()
+	rc := &Controller{pauser: DeliverPauser, cfg: cfg}
+	_, _, err = s.cr.encodeLoop(ctx, kvsCh, s.tr, s.tr.logger, kvEncoder, deliverCompleteCh, rc)
+	require.NoError(s.T(), err)
+	require.Len(s.T(), kvsCh, 2)
+
+	kvs := <-kvsCh
+	require.Len(s.T(), kvs, 1)
+	require.Equal(s.T(), int64(19), kvs[0].rowID)
+	require.Equal(s.T(), int64(36), kvs[0].offset)
+	require.Equal(s.T(), []string{"c1", "c_table", "c_schema", "c_source"}, kvs[0].columns)
+
+	kvs = <-kvsCh
+	require.Equal(s.T(), 1, len(kvs))
+	require.Nil(s.T(), kvs[0].kvs)
+	require.Equal(s.T(), s.cr.chunk.Chunk.EndOffset, kvs[0].offset)
+}
+
 func (s *chunkRestoreSuite) TestEncodeLoopCanceled() {
 	ctx, cancel := context.WithCancel(context.Background())
 	kvsCh := make(chan []deliveredKVs)

br/pkg/lightning/restore/get_pre_info.go

@@ -648,9 +648,12 @@ func (p *PreRestoreInfoGetterImpl) sampleDataFromTable(
 	}
 
 	initializedColumns := false
-	var columnPermutation []int
-	var kvSize uint64 = 0
-	var rowSize uint64 = 0
+	var (
+		columnPermutation []int
+		kvSize            uint64 = 0
+		rowSize           uint64 = 0
+		extendVals        []types.Datum
+	)
 	rowCount := 0
 	dataKVs := p.encBuilder.MakeEmptyRows()
 	indexKVs := p.encBuilder.MakeEmptyRows()
@@ -665,17 +668,32 @@ outloop:
 		switch errors.Cause(err) {
 		case nil:
 			if !initializedColumns {
+				ignoreColsMap := igCols.ColumnsMap()
 				if len(columnPermutation) == 0 {
 					columnPermutation, err = createColumnPermutation(
 						columnNames,
-						igCols.ColumnsMap(),
+						ignoreColsMap,
 						tableInfo,
 						log.FromContext(ctx))
 					if err != nil {
 						return 0.0, false, errors.Trace(err)
 					}
 				}
+				if len(sampleFile.ExtendData.Columns) > 0 {
+					_, extendVals = filterColumns(columnNames, sampleFile.ExtendData, ignoreColsMap, tableInfo)
+				}
 				initializedColumns = true
+				lastRow := parser.LastRow()
+				lastRowLen := len(lastRow.Row)
+				extendColsMap := make(map[string]int)
+				for i, c := range sampleFile.ExtendData.Columns {
+					extendColsMap[c] = lastRowLen + i
+				}
+				for i, col := range tableInfo.Columns {
+					if p, ok := extendColsMap[col.Name.O]; ok {
+						columnPermutation[i] = p
+					}
+				}
 			}
 		case io.EOF:
 			break outloop
@@ -685,6 +703,7 @@ outloop:
 		}
 		lastRow := parser.LastRow()
 		rowCount++
+		lastRow.Row = append(lastRow.Row, extendVals...)
 
 		var dataChecksum, indexChecksum verification.KVChecksum
 		kvs, encodeErr := kvEncoder.Encode(logTask.Logger, lastRow.Row, lastRow.RowID, columnPermutation, sampleFile.Path, offset)

br/pkg/lightning/restore/precheck_impl.go

@@ -39,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/engine"
 	"github.com/pingcap/tidb/util/mathutil"
+	"github.com/pingcap/tidb/util/set"
 	"go.uber.org/zap"
 	"golang.org/x/exp/slices"
 	"golang.org/x/sync/errgroup"
@@ -749,15 +750,73 @@ func (ci *schemaCheckItem) SchemaIsValid(ctx context.Context, tableInfo *mydump.
 	}
 	igCols := igCol.ColumnsMap()
 
+	fullExtendColsSet := make(set.StringSet)
+	for _, fileInfo := range tableInfo.DataFiles {
+		for _, col := range fileInfo.FileMeta.ExtendData.Columns {
+			if _, ok = igCols[col]; ok {
+				msgs = append(msgs, fmt.Sprintf("extend column %s is also assigned in ignore-column for table `%s`.`%s`, "+
+					"please keep only either one of them", col, tableInfo.DB, tableInfo.Name))
+			}
+			fullExtendColsSet.Insert(col)
+		}
+	}
+	if len(msgs) > 0 {
+		return msgs, nil
+	}
+
 	colCountFromTiDB := len(info.Core.Columns)
+	if len(fullExtendColsSet) > 0 {
+		log.FromContext(ctx).Info("check extend column count through data files", zap.String("db", tableInfo.DB),
+			zap.String("table", tableInfo.Name))
+		igColCnt := 0
+		for _, col := range info.Core.Columns {
+			if _, ok = igCols[col.Name.L]; ok {
+				igColCnt++
+			}
+		}
+		for _, f := range tableInfo.DataFiles {
+			cols, previewRows, err := ci.preInfoGetter.ReadFirstNRowsByFileMeta(ctx, f.FileMeta, 1)
+			if err != nil {
+				return nil, errors.Trace(err)
+			}
+			if len(cols) > 0 {
+				colsSet := set.NewStringSet(cols...)
+				for _, extendCol := range f.FileMeta.ExtendData.Columns {
+					if colsSet.Exist(strings.ToLower(extendCol)) {
+						msgs = append(msgs, fmt.Sprintf("extend column %s is contained in table `%s`.`%s`'s header, "+
+							"please remove this column in data or remove this extend rule", extendCol, tableInfo.DB, tableInfo.Name))
+					}
+				}
+			} else if len(previewRows) > 0 && len(previewRows[0])+len(f.FileMeta.ExtendData.Columns) > colCountFromTiDB+igColCnt {
+				msgs = append(msgs, fmt.Sprintf("row count %d adding with extend column length %d is larger than columnCount %d plus ignore column count %d for table `%s`.`%s`, "+
+					"please make sure your source data don't have extend columns and target schema has all of them", len(previewRows[0]), len(f.FileMeta.ExtendData.Columns), colCountFromTiDB, igColCnt, tableInfo.DB, tableInfo.Name))
+			}
+		}
+	}
+	if len(msgs) > 0 {
+		return msgs, nil
+	}
+
 	core := info.Core
 	defaultCols := make(map[string]struct{})
 	for _, col := range core.Columns {
-		if hasDefault(col) || (info.Core.ContainsAutoRandomBits() && mysql.HasPriKeyFlag(col.GetFlag())) {
+		// we can extend column the same with columns with default values
+		if _, isExtendCol := fullExtendColsSet[col.Name.O]; isExtendCol || hasDefault(col) || (info.Core.ContainsAutoRandomBits() && mysql.HasPriKeyFlag(col.GetFlag())) {
 			// this column has default value or it's auto random id, so we can ignore it
 			defaultCols[col.Name.L] = struct{}{}
 		}
+		delete(fullExtendColsSet, col.Name.O)
 	}
+	if len(fullExtendColsSet) > 0 {
+		extendCols := make([]string, 0, len(fullExtendColsSet))
+		for col := range fullExtendColsSet {
+			extendCols = append(extendCols, col)
+		}
+		msgs = append(msgs, fmt.Sprintf("extend column [%s] don't exist in target table `%s`.`%s` schema, "+
+			"please add these extend columns manually in downstream database/schema file", strings.Join(extendCols, ","), tableInfo.DB, tableInfo.Name))
+		return msgs, nil
+	}
+
 	// tidb_rowid have a default value.
 	defaultCols[model.ExtraHandleName.String()] = struct{}{}