*: make table attribute labels keyspace-aware

Author: ChangRui-Ryan

br/pkg/backup/BUILD.bazel

@@ -50,6 +50,7 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/metapb",
         "@com_github_pingcap_log//:log",
         "@com_github_tikv_client_go_v2//oracle",
+        "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//txnkv/txnlock",
         "@com_github_tikv_client_go_v2//util",
         "@com_github_tikv_pd_client//:client",

br/pkg/backup/schema.go

@@ -5,7 +5,6 @@ package backup
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"maps"
 	"slices"
 	"time"
@@ -29,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/pkg/statistics/handle"
 	"github.com/pingcap/tidb/pkg/statistics/util"
 	tidbutil "github.com/pingcap/tidb/pkg/util"
+	"github.com/tikv/client-go/v2/tikv"
 	kvutil "github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
@@ -170,7 +170,7 @@ func (ss *Schemas) BackupSchemas(
 					}
 				}
 				// Check merge option allowed (network I/O)
-				isMergeOptionAllowed, partitionMergeOptionAllowed, err := schema.checkMergeOptionAllowed(ectx)
+				isMergeOptionAllowed, partitionMergeOptionAllowed, err := schema.checkMergeOptionAllowed(ectx, store.GetCodec())
 				if err != nil {
 					logger.Warn("failed to check merge_option for table", logutil.ShortError(err))
 				} else {
@@ -325,20 +325,20 @@ func (s *schemaInfo) encodeToSchema() (*backuppb.Schema, error) {
 // Returns:
 //   - tableMergeOptionAllowed: whether merge_option=allow is set for the table itself
 //   - partitionMergeOptionAllowed: a map from partition name to whether merge_option=allow is set for that partition
-func (s *schemaInfo) checkMergeOptionAllowed(ctx context.Context) (bool, map[string]bool, error) {
+func (s *schemaInfo) checkMergeOptionAllowed(ctx context.Context, codec tikv.Codec) (bool, map[string]bool, error) {
 	partitionMergeOptionAllowed := make(map[string]bool)
 
 	// Construct the rule ID for this table using the same format as ddl/label
 	// Use .L (lowercase) to match DDL behavior for case-insensitive matching
 	dbName := s.dbInfo.Name.L
 	tableName := s.tableInfo.Name.L
-	ruleID := fmt.Sprintf(label.TableIDFormat, label.IDPrefix, dbName, tableName)
+	ruleID := label.NewRuleID(codec, dbName, tableName, "")
 
 	// Collect all rule IDs to check (table + partitions if any)
 	ruleIDs := []string{ruleID}
 	if s.tableInfo.Partition != nil && len(s.tableInfo.Partition.Definitions) > 0 {
 		for _, def := range s.tableInfo.Partition.Definitions {
-			partitionRuleID := fmt.Sprintf(label.PartitionIDFormat, label.IDPrefix, dbName, tableName, def.Name.L)
+			partitionRuleID := label.NewRuleID(codec, dbName, tableName, def.Name.L)
 			ruleIDs = append(ruleIDs, partitionRuleID)
 		}
 	}
@@ -370,7 +370,7 @@ func (s *schemaInfo) checkMergeOptionAllowed(ctx context.Context) (bool, map[str
 	// Check partition rules if this is a partitioned table
 	if s.tableInfo.Partition != nil && len(s.tableInfo.Partition.Definitions) > 0 {
 		for _, def := range s.tableInfo.Partition.Definitions {
-			partitionRuleID := fmt.Sprintf(label.PartitionIDFormat, label.IDPrefix, dbName, tableName, def.Name.L)
+			partitionRuleID := label.NewRuleID(codec, dbName, tableName, def.Name.L)
 			if rule, exists := rules[partitionRuleID]; exists {
 				for _, label := range rule.Labels {
 					if label.Key == "merge_option" && label.Value == "allow" {

br/pkg/backup/schema_merge_option_test.go

@@ -4,7 +4,6 @@ package backup_test
 
 import (
 	"context"
-	"fmt"
 	"math"
 	"testing"
 
@@ -17,34 +16,41 @@ import (
 	"github.com/pingcap/tidb/pkg/testkit"
 	filter "github.com/pingcap/tidb/pkg/util/table-filter"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/tikv"
 )
 
 // TestBackupSchemaMergeOptionRuleIDFormat tests the rule ID format generation
 func TestBackupSchemaMergeOptionRuleIDFormat(t *testing.T) {
+	codec := tikv.NewCodecV1(tikv.ModeTxn)
+
 	t.Run("normal table rule ID format", func(t *testing.T) {
 		dbName := "test"
 		tableName := "t1"
-		expectedRuleID := fmt.Sprintf(label.TableIDFormat, label.IDPrefix, dbName, tableName)
+		expectedRuleID := label.NewRuleID(codec, dbName, tableName, "")
 		require.Equal(t, "schema/test/t1", expectedRuleID)
 	})
 
 	t.Run("partition table rule ID format", func(t *testing.T) {
 		dbName := "test"
 		tableName := "pt1"
 		partitionName := "p0"
-		expectedRuleID := fmt.Sprintf(label.PartitionIDFormat, label.IDPrefix, dbName, tableName, partitionName)
+		expectedRuleID := label.NewRuleID(codec, dbName, tableName, partitionName)
 		require.Equal(t, "schema/test/pt1/p0", expectedRuleID)
 	})
 
 	t.Run("multiple partitions rule ID format", func(t *testing.T) {
 		dbName := "test"
 		tableName := "pt2"
 		partitions := []string{"p0", "p1", "p2"}
+		expected := []string{
+			"schema/test/pt2/p0",
+			"schema/test/pt2/p1",
+			"schema/test/pt2/p2",
+		}
 
-		for _, partName := range partitions {
-			ruleID := fmt.Sprintf(label.PartitionIDFormat, label.IDPrefix, dbName, tableName, partName)
-			expected := fmt.Sprintf("schema/%s/%s/%s", dbName, tableName, partName)
-			require.Equal(t, expected, ruleID, "partition %s rule ID should match", partName)
+		for i, partName := range partitions {
+			ruleID := label.NewRuleID(codec, dbName, tableName, partName)
+			require.Equal(t, expected[i], ruleID, "partition %s rule ID should match", partName)
 		}
 	})
 
@@ -56,13 +62,13 @@ func TestBackupSchemaMergeOptionRuleIDFormat(t *testing.T) {
 		// Simulate DDL behavior: uses .L (lowercase)
 		// DDL would use: schema.Name.L, meta.Name.L, spec.PartitionNames[0].L
 		// For "TestDB"/"TestTable"/"Partition0", .L would be: "testdb"/"testtable"/"partition0"
-		ddlRuleID := fmt.Sprintf(label.TableIDFormat, label.IDPrefix, "testdb", "testtable")
-		ddlPartitionRuleID := fmt.Sprintf(label.PartitionIDFormat, label.IDPrefix, "testdb", "testtable", "partition0")
+		ddlRuleID := label.NewRuleID(codec, "testdb", "testtable", "")
+		ddlPartitionRuleID := label.NewRuleID(codec, "testdb", "testtable", "partition0")
 
 		// BR should also use .L (lowercase) - simulate BR behavior
 		// In real code: table.DB.Name.L, table.Info.Name.L, def.Name.L
-		brRuleID := fmt.Sprintf(label.TableIDFormat, label.IDPrefix, "testdb", "testtable")
-		brPartitionRuleID := fmt.Sprintf(label.PartitionIDFormat, label.IDPrefix, "testdb", "testtable", "partition0")
+		brRuleID := label.NewRuleID(codec, "testdb", "testtable", "")
+		brPartitionRuleID := label.NewRuleID(codec, "testdb", "testtable", "partition0")
 
 		// Both should generate the same rule ID
 		require.Equal(t, ddlRuleID, brRuleID, "DDL and BR should generate same table rule ID")

br/pkg/restore/snap_client/client.go

@@ -1199,7 +1199,7 @@ func (rc *SnapClient) setMergeOptionForTables(ctx context.Context, createdTables
 			}
 			// Use Reset() to set ID, RuleType, Data, Index, and add/update db/table labels
 			// Reset() uses the NEW table ID (after restore)
-			rule.Reset(dbName, tableName, "", newTableInfo.ID)
+			rule.Reset(rc.dom.Store().GetCodec(), dbName, tableName, "", newTableInfo.ID)
 
 			rulesToSet = append(rulesToSet, rule)
 		}
@@ -1237,7 +1237,7 @@ func (rc *SnapClient) setMergeOptionForTables(ctx context.Context, createdTables
 					}
 					// Use Reset() to set ID, RuleType, Data, Index, and add/update db/table/partition labels
 					// Reset() uses the NEW partition ID (after restore)
-					rule.Reset(dbName, tableName, partitionName, newDef.ID)
+					rule.Reset(rc.dom.Store().GetCodec(), dbName, tableName, partitionName, newDef.ID)
 
 					rulesToSet = append(rulesToSet, rule)
 				}

pkg/ddl/executor.go

@@ -6101,7 +6101,7 @@ func (e *executor) AlterTableAttributes(ctx sessionctx.Context, ident ast.Ident,
 		return dbterror.ErrInvalidAttributesSpec.GenWithStackByArgs(err)
 	}
 	ids := getIDs([]*model.TableInfo{meta})
-	rule.Reset(schema.Name.L, meta.Name.L, "", ids...)
+	rule.Reset(e.store.GetCodec(), schema.Name.L, meta.Name.L, "", ids...)
 
 	job := &model.Job{
 		Version:        model.GetJobVerInUse(),
@@ -6146,7 +6146,7 @@ func (e *executor) AlterTablePartitionAttributes(ctx sessionctx.Context, ident a
 	if err != nil {
 		return dbterror.ErrInvalidAttributesSpec.GenWithStackByArgs(err)
 	}
-	rule.Reset(schema.Name.L, meta.Name.L, spec.PartitionNames[0].L, partitionID)
+	rule.Reset(e.store.GetCodec(), schema.Name.L, meta.Name.L, spec.PartitionNames[0].L, partitionID)
 
 	pdLabelRule := (*pdhttp.LabelRule)(rule)
 	job := &model.Job{

pkg/ddl/label/BUILD.bazel

@@ -10,9 +10,11 @@ go_library(
     importpath = "github.com/pingcap/tidb/pkg/ddl/label",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/config/kerneltype",
         "//pkg/parser/ast",
         "//pkg/tablecodec",
         "//pkg/util/codec",
+        "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_pd_client//http",
         "@in_gopkg_yaml_v2//:yaml_v2",
     ],
@@ -28,11 +30,15 @@ go_test(
     ],
     embed = [":label"],
     flaky = True,
-    shard_count = 8,
+    shard_count = 9,
     deps = [
+        "//pkg/config/kerneltype",
         "//pkg/parser/ast",
+        "//pkg/tablecodec",
         "//pkg/testkit/testsetup",
+        "@com_github_pingcap_kvproto//pkg/keyspacepb",
         "@com_github_stretchr_testify//require",
+        "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_pd_client//http",
         "@org_uber_go_goleak//:goleak",
     ],

pkg/ddl/label/attributes.go

@@ -18,10 +18,12 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/pingcap/tidb/pkg/config/kerneltype"
 	pd "github.com/tikv/pd/client/http"
 )
 
 const (
+	keyspaceKey  = "keyspace"
 	dbKey        = "db"
 	tableKey     = "table"
 	partitionKey = "partition"
@@ -103,6 +105,10 @@ func RestoreRegionLabels(labels *[]pd.RegionLabel) string {
 		switch label.Key {
 		case dbKey, tableKey, partitionKey:
 			continue
+		case keyspaceKey:
+			if kerneltype.IsNextGen() {
+				continue
+			}
 		default:
 		}
 

pkg/ddl/label/attributes_test.go

@@ -17,6 +17,7 @@ package label
 import (
 	"testing"
 
+	"github.com/pingcap/tidb/pkg/config/kerneltype"
 	"github.com/stretchr/testify/require"
 	pd "github.com/tikv/pd/client/http"
 )
@@ -200,6 +201,8 @@ func TestRestoreLabels(t *testing.T) {
 	require.NoError(t, err)
 	input5, err := NewLabel("partition=p1")
 	require.NoError(t, err)
+	input6, err := NewLabel("keyspace=42")
+	require.NoError(t, err)
 
 	tests := []TestCase{
 		{
@@ -230,4 +233,12 @@ func TestRestoreLabels(t *testing.T) {
 			require.Equal(t, test.output, output)
 		})
 	}
+
+	if kerneltype.IsNextGen() {
+		output := RestoreRegionLabels(&[]pd.RegionLabel{input1, input6})
+		require.Equal(t, `"merge_option=allow"`, output)
+	} else {
+		output := RestoreRegionLabels(&[]pd.RegionLabel{input1, input6})
+		require.Equal(t, `"merge_option=allow","keyspace=42"`, output)
+	}
 }

pkg/ddl/label/rule.go

@@ -19,18 +19,23 @@ import (
 	"encoding/json"
 	"fmt"
 	"slices"
+	"strconv"
 
+	"github.com/pingcap/tidb/pkg/config/kerneltype"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/tablecodec"
 	"github.com/pingcap/tidb/pkg/util/codec"
+	"github.com/tikv/client-go/v2/tikv"
 	pd "github.com/tikv/pd/client/http"
 	"gopkg.in/yaml.v2"
 )
 
 const (
 	// IDPrefix is the prefix for label rule ID.
 	IDPrefix = "schema"
-	ruleType = "key-range"
+	// KeyspacePrefix is the prefix for keyspace in label rule ID.
+	KeyspacePrefix = "keyspace"
+	ruleType       = "key-range"
 )
 
 const (
@@ -94,20 +99,43 @@ func (r *Rule) Clone() *Rule {
 	return newRule
 }
 
-// Reset will reset the label rule for a table/partition with a given ID and names.
-func (r *Rule) Reset(dbName, tableName, partName string, ids ...int64) *Rule {
+// UseKeyspaceAwareRules returns true when table attribute label rules should be
+// scoped by keyspace in NextGen deployments.
+func UseKeyspaceAwareRules(tikvCodec tikv.Codec) bool {
+	return kerneltype.IsNextGen() && tikvCodec != nil && tikvCodec.GetKeyspaceMeta() != nil
+}
+
+// NewRuleID generates a new rule ID for a table or partition.
+func NewRuleID(tikvCodec tikv.Codec, dbName, tableName, partName string) string {
+	var id string
 	isPartition := partName != ""
 	if isPartition {
-		r.ID = fmt.Sprintf(PartitionIDFormat, IDPrefix, dbName, tableName, partName)
+		id = fmt.Sprintf(PartitionIDFormat, IDPrefix, dbName, tableName, partName)
 	} else {
-		r.ID = fmt.Sprintf(TableIDFormat, IDPrefix, dbName, tableName)
+		id = fmt.Sprintf(TableIDFormat, IDPrefix, dbName, tableName)
 	}
+	if UseKeyspaceAwareRules(tikvCodec) {
+		id = fmt.Sprintf("%s/%d/%s", KeyspacePrefix, tikvCodec.GetKeyspaceID(), id)
+	}
+	return id
+}
+
+// Reset will reset the label rule for a table/partition with a given ID and names.
+func (r *Rule) Reset(tikvCodec tikv.Codec, dbName, tableName, partName string, ids ...int64) *Rule {
+	isPartition := partName != ""
+	useKeyspace := UseKeyspaceAwareRules(tikvCodec)
+	r.ID = NewRuleID(tikvCodec, dbName, tableName, partName)
 	if len(r.Labels) == 0 {
 		return r
 	}
-	var hasDBKey, hasTableKey, hasPartitionKey bool
+	var hasKeyspaceKey, hasDBKey, hasTableKey, hasPartitionKey bool
 	for i := range r.Labels {
 		switch r.Labels[i].Key {
+		case keyspaceKey:
+			if useKeyspace {
+				r.Labels[i].Value = strconv.FormatInt(int64(tikvCodec.GetKeyspaceID()), 10)
+				hasKeyspaceKey = true
+			}
 		case dbKey:
 			r.Labels[i].Value = dbName
 			hasDBKey = true
@@ -123,6 +151,10 @@ func (r *Rule) Reset(dbName, tableName, partName string, ids ...int64) *Rule {
 		}
 	}
 
+	if useKeyspace && !hasKeyspaceKey {
+		r.Labels = append(r.Labels, pd.RegionLabel{Key: keyspaceKey, Value: strconv.FormatInt(int64(tikvCodec.GetKeyspaceID()), 10)})
+	}
+
 	if !hasDBKey {
 		r.Labels = append(r.Labels, pd.RegionLabel{Key: dbKey, Value: dbName})
 	}
@@ -138,9 +170,16 @@ func (r *Rule) Reset(dbName, tableName, partName string, ids ...int64) *Rule {
 	dataSlice := make([]any, 0, len(ids))
 	slices.Sort(ids)
 	for i := range ids {
+		startKey := codec.EncodeBytes(nil, tablecodec.GenTablePrefix(ids[i]))
+		endKey := codec.EncodeBytes(nil, tablecodec.GenTablePrefix(ids[i]+1))
+		if useKeyspace {
+			// Label rules are consumed as region boundary keys, so V2 must encode
+			// the whole outer key instead of prefixing a mem-encoded table key.
+			startKey, endKey = tikvCodec.EncodeRegionRange(tablecodec.GenTablePrefix(ids[i]), tablecodec.GenTablePrefix(ids[i]+1))
+		}
 		data := map[string]string{
-			"start_key": hex.EncodeToString(codec.EncodeBytes(nil, tablecodec.GenTablePrefix(ids[i]))),
-			"end_key":   hex.EncodeToString(codec.EncodeBytes(nil, tablecodec.GenTablePrefix(ids[i]+1))),
+			"start_key": hex.EncodeToString(startKey),
+			"end_key":   hex.EncodeToString(endKey),
 		}
 		dataSlice = append(dataSlice, data)
 	}