lightning: optimize a query to I_S.columns (#57117)

close #57115

lightning/pkg/importer/get_pre_info.go

@@ -94,7 +94,7 @@ type TargetInfoGetter interface {
 	// FetchRemoteDBModels fetches the database structures from the remote target.
 	FetchRemoteDBModels(ctx context.Context) ([]*model.DBInfo, error)
 	// FetchRemoteTableModels fetches the table structures from the remote target.
-	FetchRemoteTableModels(ctx context.Context, schemaName string) ([]*model.TableInfo, error)
+	FetchRemoteTableModels(ctx context.Context, schemaName string, tableNames []string) (map[string]*model.TableInfo, error)
 	// CheckVersionRequirements performs the check whether the target satisfies the version requirements.
 	CheckVersionRequirements(ctx context.Context) error
 	// IsTableEmpty checks whether the specified table on the target DB contains data or not.
@@ -162,8 +162,12 @@ func (g *TargetInfoGetterImpl) FetchRemoteDBModels(ctx context.Context) ([]*mode
 
 // FetchRemoteTableModels fetches the table structures from the remote target.
 // It implements the TargetInfoGetter interface.
-func (g *TargetInfoGetterImpl) FetchRemoteTableModels(ctx context.Context, schemaName string) ([]*model.TableInfo, error) {
-	return g.backend.FetchRemoteTableModels(ctx, schemaName)
+func (g *TargetInfoGetterImpl) FetchRemoteTableModels(
+	ctx context.Context,
+	schemaName string,
+	tableNames []string,
+) (map[string]*model.TableInfo, error) {
+	return g.backend.FetchRemoteTableModels(ctx, schemaName, tableNames)
 }
 
 // CheckVersionRequirements performs the check whether the target satisfies the version requirements.
@@ -365,6 +369,10 @@ func (p *PreImportInfoGetterImpl) GetAllTableStructures(ctx context.Context, opt
 
 func (p *PreImportInfoGetterImpl) getTableStructuresByFileMeta(ctx context.Context, dbSrcFileMeta *mydump.MDDatabaseMeta, getPreInfoCfg *ropts.GetPreInfoConfig) ([]*model.TableInfo, error) {
 	dbName := dbSrcFileMeta.Name
+	tableNames := make([]string, 0, len(dbSrcFileMeta.Tables))
+	for _, tableFileMeta := range dbSrcFileMeta.Tables {
+		tableNames = append(tableNames, tableFileMeta.Name)
+	}
 	failpoint.Inject(
 		"getTableStructuresByFileMeta_BeforeFetchRemoteTableModels",
 		func(v failpoint.Value) {
@@ -378,7 +386,7 @@ func (p *PreImportInfoGetterImpl) getTableStructuresByFileMeta(ctx context.Conte
 			failpoint.Enable("github.com/pingcap/tidb/pkg/lightning/backend/tidb/FetchRemoteTableModels_BeforeFetchTableAutoIDInfos", fmt.Sprintf("sleep(%d)", sleepMilliSeconds))
 		},
 	)
-	currentTableInfosFromDB, err := p.targetInfoGetter.FetchRemoteTableModels(ctx, dbName)
+	currentTableInfosMap, err := p.targetInfoGetter.FetchRemoteTableModels(ctx, dbName, tableNames)
 	if err != nil {
 		if getPreInfoCfg != nil && getPreInfoCfg.IgnoreDBNotExist {
 			dbNotExistErr := dbterror.ClassSchema.NewStd(errno.ErrBadDB).FastGenByArgs(dbName)
@@ -394,10 +402,6 @@ func (p *PreImportInfoGetterImpl) getTableStructuresByFileMeta(ctx context.Conte
 		return nil, errors.Trace(err)
 	}
 get_struct_from_src:
-	currentTableInfosMap := make(map[string]*model.TableInfo)
-	for _, tblInfo := range currentTableInfosFromDB {
-		currentTableInfosMap[tblInfo.Name.L] = tblInfo
-	}
 	resultInfos := make([]*model.TableInfo, len(dbSrcFileMeta.Tables))
 	for i, tableFileMeta := range dbSrcFileMeta.Tables {
 		if curTblInfo, ok := currentTableInfosMap[strings.ToLower(tableFileMeta.Name)]; ok {
@@ -804,8 +808,12 @@ func (p *PreImportInfoGetterImpl) FetchRemoteDBModels(ctx context.Context) ([]*m
 
 // FetchRemoteTableModels fetches the table structures from the remote target.
 // It implements the PreImportInfoGetter interface.
-func (p *PreImportInfoGetterImpl) FetchRemoteTableModels(ctx context.Context, schemaName string) ([]*model.TableInfo, error) {
-	return p.targetInfoGetter.FetchRemoteTableModels(ctx, schemaName)
+func (p *PreImportInfoGetterImpl) FetchRemoteTableModels(
+	ctx context.Context,
+	schemaName string,
+	tableNames []string,
+) (map[string]*model.TableInfo, error) {
+	return p.targetInfoGetter.FetchRemoteTableModels(ctx, schemaName, tableNames)
 }
 
 // CheckVersionRequirements performs the check whether the target satisfies the version requirements.

lightning/pkg/importer/mock/mock.go

@@ -225,17 +225,25 @@ func (t *TargetInfo) FetchRemoteDBModels(_ context.Context) ([]*model.DBInfo, er
 
 // FetchRemoteTableModels fetches the table structures from the remote target.
 // It implements the TargetInfoGetter interface.
-func (t *TargetInfo) FetchRemoteTableModels(_ context.Context, schemaName string) ([]*model.TableInfo, error) {
-	resultInfos := []*model.TableInfo{}
+func (t *TargetInfo) FetchRemoteTableModels(
+	_ context.Context,
+	schemaName string,
+	tableNames []string,
+) (map[string]*model.TableInfo, error) {
 	tblMap, ok := t.dbTblInfoMap[schemaName]
 	if !ok {
 		dbNotExistErr := dbterror.ClassSchema.NewStd(errno.ErrBadDB).FastGenByArgs(schemaName)
 		return nil, errors.Errorf("get xxxxxx http status code != 200, message %s", dbNotExistErr.Error())
 	}
-	for _, tblInfo := range tblMap {
-		resultInfos = append(resultInfos, tblInfo.TableModel)
+	ret := make(map[string]*model.TableInfo, len(tableNames))
+	for _, tableName := range tableNames {
+		tblInfo, ok := tblMap[tableName]
+		if !ok {
+			continue
+		}
+		ret[tableName] = tblInfo.TableModel
 	}
-	return resultInfos, nil
+	return ret, nil
 }
 
 // GetTargetSysVariablesForImport gets some important systam variables for importing on the target.

lightning/pkg/importer/mock/mock_test.go

@@ -185,7 +185,7 @@ func TestMockTargetInfoBasic(t *testing.T) {
 			RowCount: 100,
 		},
 	)
-	tblInfos, err := ti.FetchRemoteTableModels(ctx, "testdb")
+	tblInfos, err := ti.FetchRemoteTableModels(ctx, "testdb", []string{"testtbl1", "testtbl2"})
 	require.NoError(t, err)
 	require.Equal(t, 2, len(tblInfos))
 	for _, tblInfo := range tblInfos {

pkg/lightning/backend/backend.go

@@ -142,9 +142,13 @@ type TargetInfoGetter interface {
 	// the database name is filled.
 	FetchRemoteDBModels(ctx context.Context) ([]*model.DBInfo, error)
 
-	// FetchRemoteTableModels obtains the models of all tables given the schema
-	// name. The returned table info does not need to be precise if the encoder,
-	// is not requiring them, but must at least fill in the following fields for
+	// FetchRemoteTableModels obtains the TableInfo of given tables under the schema
+	// name. It returns a map whose key is the table name in lower case and value is
+	// the TableInfo. If the table does not exist, it will not be included in the
+	// map.
+	//
+	// The returned table info does not need to be precise if the encoder, is not
+	// requiring them, but must at least fill in the following fields for
 	// TablesFromMeta to succeed:
 	//  - Name
 	//  - State (must be model.StatePublic)
@@ -154,7 +158,7 @@ type TargetInfoGetter interface {
 	//     * State (must be model.StatePublic)
 	//     * Offset (must be 0, 1, 2, ...)
 	//  - PKIsHandle (true = do not generate _tidb_rowid)
-	FetchRemoteTableModels(ctx context.Context, schemaName string) ([]*model.TableInfo, error)
+	FetchRemoteTableModels(ctx context.Context, schemaName string, tableNames []string) (map[string]*model.TableInfo, error)
 
 	// CheckRequirements performs the check whether the backend satisfies the version requirements
 	CheckRequirements(ctx context.Context, checkCtx *CheckCtx) error

pkg/lightning/backend/local/local.go

@@ -286,8 +286,27 @@ func (g *targetInfoGetter) FetchRemoteDBModels(ctx context.Context) ([]*model.DB
 
 // FetchRemoteTableModels obtains the models of all tables given the schema name.
 // It implements the `TargetInfoGetter` interface.
-func (g *targetInfoGetter) FetchRemoteTableModels(ctx context.Context, schemaName string) ([]*model.TableInfo, error) {
-	return tikv.FetchRemoteTableModelsFromTLS(ctx, g.tls, schemaName)
+func (g *targetInfoGetter) FetchRemoteTableModels(
+	ctx context.Context,
+	schemaName string,
+	tableNames []string,
+) (map[string]*model.TableInfo, error) {
+	allTablesInDB, err := tikv.FetchRemoteTableModelsFromTLS(ctx, g.tls, schemaName)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+
+	tableNamesSet := make(map[string]struct{}, len(tableNames))
+	for _, name := range tableNames {
+		tableNamesSet[strings.ToLower(name)] = struct{}{}
+	}
+	ret := make(map[string]*model.TableInfo, len(tableNames))
+	for _, tbl := range allTablesInDB {
+		if _, ok := tableNamesSet[tbl.Name.L]; ok {
+			ret[tbl.Name.L] = tbl
+		}
+	}
+	return ret, nil
 }
 
 // CheckRequirements performs the check whether the backend satisfies the version requirements.

pkg/lightning/backend/tidb/BUILD.bazel

@@ -6,7 +6,6 @@ go_library(
     importpath = "github.com/pingcap/tidb/pkg/lightning/backend/tidb",
     visibility = ["//visibility:public"],
     deps = [
-        "//br/pkg/version",
         "//pkg/errno",
         "//pkg/lightning/backend",
         "//pkg/lightning/backend/encode",
@@ -21,6 +20,7 @@ go_library(
         "//pkg/parser/mysql",
         "//pkg/table",
         "//pkg/types",
+        "//pkg/util",
         "//pkg/util/dbutil",
         "//pkg/util/hack",
         "//pkg/util/kvcache",