[DNM]dumpling: Support check keyspace in dumpling

dumpling/export/config.go

@@ -76,6 +76,7 @@ const (
 	flagTransactionalConsistency = "transactional-consistency"
 	flagCompress                 = "compress"
 	flagCsvOutputDialect         = "csv-output-dialect"
+	flagKeyspaceName             = "keyspace-name"
 
 	// FlagHelp represents the help flag
 	FlagHelp = "help"
@@ -188,6 +189,8 @@ type Config struct {
 
 	IOTotalBytes *atomic.Uint64
 	Net          string
+
+	KeyspaceName string
 }
 
 // ServerInfoUnknown is the unknown database type to dumpling
@@ -357,6 +360,7 @@ func (*Config) DefineFlags(flags *pflag.FlagSet) {
 	_ = flags.MarkHidden(flagTransactionalConsistency)
 	flags.StringP(flagCompress, "c", "", "Compress output file type, support 'gzip', 'snappy', 'zstd', 'no-compression' now")
 	flags.String(flagCsvOutputDialect, "", "The dialect of output CSV file, support 'snowflake', 'redshift', 'bigquery' now")
+	flags.String(flagKeyspaceName, "", "The keyspace name used to check the data to be accessed is the same as the keyspace of the target TiDB instance")
 }
 
 // ParseFromFlags parses dumpling's export.Config from flags
@@ -611,6 +615,11 @@ func (conf *Config) ParseFromFlags(flags *pflag.FlagSet) error {
 		return errors.Trace(err)
 	}
 
+	conf.KeyspaceName, err = flags.GetString(flagKeyspaceName)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
 	return nil
 }
 

dumpling/export/dump.go

@@ -30,13 +30,15 @@ import (
 	"github.com/pingcap/tidb/dumpling/cli"
 	tcontext "github.com/pingcap/tidb/dumpling/context"
 	"github.com/pingcap/tidb/dumpling/log"
+	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/format"
 	"github.com/pingcap/tidb/pkg/store/helper"
 	"github.com/pingcap/tidb/pkg/tablecodec"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/codec"
+	"github.com/pingcap/tidb/pkg/util/dbutil"
 	pd "github.com/tikv/pd/client"
 	gatomic "go.uber.org/atomic"
 	"go.uber.org/zap"
@@ -336,6 +338,27 @@ func (d *Dumper) Dump() (dumpErr error) {
 	return nil
 }
 
+func (d *Dumper) updateTiDBGlobalConfigKeyspaceName() {
+	d.tctx.L().Info("updateTiDBGlobalConfigKeyspaceName.", zap.Int("d.conf.ServerInfo.ServerType", int(d.conf.ServerInfo.ServerType)))
+	if d.conf.ServerInfo.ServerType == version.ServerTypeTiDB || d.conf.ServerInfo.ServerType == version.ServerTypeUnknown {
+		keyspaceNameInTiDB, err := dbutil.GetKeyspaceNameFromTiDB(d.dbHandle)
+		if err != nil {
+			panic(err)
+		}
+
+		if d.conf.KeyspaceName != keyspaceNameInTiDB {
+			panic("the keyspace name in command line is different from keyspace name in TiDB.")
+		}
+
+		if keyspaceNameInTiDB != "" {
+			config.UpdateGlobal(func(conf *config.Config) {
+				conf.KeyspaceName = keyspaceNameInTiDB
+			})
+			d.tctx.L().Info("dumpling using keyspace.", zap.String("keyspaceName", keyspaceNameInTiDB))
+		}
+	}
+}
+
 func (d *Dumper) startWriters(tctx *tcontext.Context, wg *errgroup.Group, taskChan <-chan Task,
 	rebuildConnFn func(*sql.Conn, bool) (*sql.Conn, error)) ([]*Writer, func(), error) {
 	conf, pool := d.conf, d.dbHandle
@@ -1369,6 +1392,9 @@ func openSQLDB(d *Dumper) error {
 		return errors.Trace(err)
 	}
 	d.dbHandle = sql.OpenDB(c)
+
+	// Check that the command line keyspace name is the same as the target tidb keyspace name.
+	d.updateTiDBGlobalConfigKeyspaceName()
 	return nil
 }
 

pkg/executor/builder.go

@@ -2025,6 +2025,7 @@ func (b *executorBuilder) buildMemTable(v *plannercore.PhysicalMemTable) exec.Ex
 			strings.ToLower(infoschema.ClusterTableMemoryUsage),
 			strings.ToLower(infoschema.ClusterTableMemoryUsageOpsHistory),
 			strings.ToLower(infoschema.TableResourceGroups),
+			strings.ToLower(infoschema.TableServerlessMeta),
 			strings.ToLower(infoschema.TableRunawayWatches),
 			strings.ToLower(infoschema.TableCheckConstraints),
 			strings.ToLower(infoschema.TableTiDBCheckConstraints),

pkg/executor/infoschema_reader.go

@@ -198,6 +198,8 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex
 			err = e.setDataForClusterMemoryUsageOpsHistory(sctx)
 		case infoschema.TableResourceGroups:
 			err = e.setDataFromResourceGroups()
+		case infoschema.TableServerlessMeta:
+			e.setKeyspaceMeta(sctx)
 		case infoschema.TableRunawayWatches:
 			err = e.setDataFromRunawayWatches(sctx)
 		case infoschema.TableCheckConstraints:
@@ -3631,6 +3633,27 @@ func (e *memtableRetriever) setDataForClusterIndexUsage(ctx sessionctx.Context,
 	return nil
 }
 
+func (e *memtableRetriever) setKeyspaceMeta(sctx sessionctx.Context) {
+	var (
+		keyspaceName string
+		keyspaceID   string
+
+		keyspaceMeta = sctx.GetStore().GetCodec().GetKeyspaceMeta()
+	)
+
+	e.rows = make([][]types.Datum, 1)
+
+	if keyspaceMeta != nil {
+		keyspaceName = keyspaceMeta.Name
+		keyspaceID = fmt.Sprintf("%d", keyspaceMeta.Id)
+	}
+
+	e.rows[0] = types.MakeDatums(
+		keyspaceName,
+		keyspaceID,
+	)
+}
+
 func checkRule(rule *label.Rule) (dbName, tableName string, partitionName string, err error) {
 	s := strings.Split(rule.ID, "/")
 	if len(s) < 3 {

pkg/infoschema/tables.go

@@ -205,6 +205,8 @@ const (
 	TableMemoryUsageOpsHistory = "MEMORY_USAGE_OPS_HISTORY"
 	// TableResourceGroups is the metadata of resource groups.
 	TableResourceGroups = "RESOURCE_GROUPS"
+	// TableServerlessMeta is the metadata of the current serverless cluster.
+	TableServerlessMeta = "KEYSPACE_META"
 	// TableRunawayWatches is the query list of runaway watch.
 	TableRunawayWatches = "RUNAWAY_WATCHES"
 	// TableCheckConstraints is the list of CHECK constraints.
@@ -336,6 +338,7 @@ var tableIDMap = map[string]int64{
 	TableKeywords:                        autoid.InformationSchemaDBID + 92,
 	TableTiDBIndexUsage:                  autoid.InformationSchemaDBID + 93,
 	ClusterTableTiDBIndexUsage:           autoid.InformationSchemaDBID + 94,
+	TableServerlessMeta:                  autoid.InformationSchemaDBID + 95,
 }
 
 // columnInfo represents the basic column information of all kinds of INFORMATION_SCHEMA tables
@@ -1651,6 +1654,11 @@ var tableResourceGroupsCols = []columnInfo{
 	{name: "BACKGROUND", tp: mysql.TypeVarchar, size: 256},
 }
 
+var tableServerlessMetaCols = []columnInfo{
+	{name: "KEYSPACE_NAME", tp: mysql.TypeVarchar, size: 64},
+	{name: "KEYSPACE_ID", tp: mysql.TypeVarchar, size: 64},
+}
+
 var tableRunawayWatchListCols = []columnInfo{
 	{name: "ID", tp: mysql.TypeLonglong, size: 64, flag: mysql.NotNullFlag},
 	{name: "RESOURCE_GROUP_NAME", tp: mysql.TypeVarchar, size: resourcegroup.MaxGroupNameLength, flag: mysql.NotNullFlag},
@@ -2348,6 +2356,7 @@ var tableNameToColumns = map[string][]columnInfo{
 	TableTiDBCheckConstraints:               tableTiDBCheckConstraintsCols,
 	TableKeywords:                           tableKeywords,
 	TableTiDBIndexUsage:                     tableTiDBIndexUsage,
+	TableServerlessMeta:                     tableServerlessMetaCols,
 }
 
 func createInfoSchemaTable(_ autoid.Allocators, meta *model.TableInfo) (table.Table, error) {

pkg/util/dbutil/common.go

@@ -887,3 +887,44 @@ func GetParserForDB(ctx context.Context, db QueryExecutor) (*parser.Parser, erro
 	parser2.SetSQLMode(mode)
 	return parser2, nil
 }
+
+// GetKeyspaceNameFromTiDB get keyspace name from TiDB.
+func GetKeyspaceNameFromTiDB(db *sql.DB) (string, error) {
+	if db == nil {
+		return "", nil
+	}
+
+	getConfigTableSQL := "select  TABLE_NAME  from  INFORMATION_SCHEMA . TABLES  where TABLE_SCHEMA ='INFORMATION_SCHEMA' and  TABLE_NAME ='KEYSPACE_META'"
+	rows, err := db.Query(getConfigTableSQL)
+	if err != nil {
+		return "", err
+	}
+	defer rows.Close()
+
+	if rows == nil {
+		return "", nil
+	}
+
+	if !rows.Next() {
+		return "", rows.Err()
+	}
+
+	sql := "select * from INFORMATION_SCHEMA.KEYSPACE_META"
+	rows, err = db.Query(sql)
+	if err != nil {
+		return "", err
+	}
+	var (
+		keyspaceName string
+		_keyspaceID  string
+	)
+	if rows.Next() {
+		err = rows.Scan(&keyspaceName, &_keyspaceID)
+		if err != nil {
+			return "", err
+		}
+		log.Info("get keyspace from TiDB", zap.String("keyspace-name", keyspaceName))
+		return keyspaceName, rows.Err()
+	}
+	return "", nil
+}