diff --git a/ddl/placement/bundle.go b/ddl/placement/bundle.go index 1bb7539d79ad0..34e0e074ecda0 100644 --- a/ddl/placement/bundle.go +++ b/ddl/placement/bundle.go @@ -19,6 +19,7 @@ import ( "errors" "fmt" "strconv" + "strings" "github.com/pingcap/failpoint" "github.com/pingcap/parser/ast" @@ -201,3 +202,40 @@ func (b *Bundle) Clone() *Bundle { func (b *Bundle) IsEmpty() bool { return len(b.Rules) == 0 && b.Index == 0 && !b.Override } + +// ObjectID extracts the db/table/partition ID from the group ID +func (b *Bundle) ObjectID() (int64, error) { + // If the rule doesn't come from TiDB, skip it. + if !strings.HasPrefix(b.ID, BundleIDPrefix) { + return 0, ErrInvalidBundleIDFormat + } + id, err := strconv.ParseInt(b.ID[len(BundleIDPrefix):], 10, 64) + if err != nil { + return 0, fmt.Errorf("%w: %s", ErrInvalidBundleID, err) + } + if id <= 0 { + return 0, fmt.Errorf("%w: %s doesn't include an id", ErrInvalidBundleID, b.ID) + } + return id, nil +} + +func isValidLeaderRule(rule *Rule, dcLabelKey string) bool { + if rule.Role == Leader && rule.Count == 1 { + for _, con := range rule.Constraints { + if con.Op == In && con.Key == dcLabelKey && len(con.Values) == 1 { + return true + } + } + } + return false +} + +// GetLeaderDC returns the leader's DC by Bundle if found. +func (b *Bundle) GetLeaderDC(dcLabelKey string) (string, bool) { + for _, rule := range b.Rules { + if isValidLeaderRule(rule, dcLabelKey) { + return rule.Constraints[0].Values[0], true + } + } + return "", false +} diff --git a/ddl/placement/bundle_test.go b/ddl/placement/bundle_test.go index c63fbae74e670..98b989a21ee75 100644 --- a/ddl/placement/bundle_test.go +++ b/ddl/placement/bundle_test.go @@ -56,6 +56,290 @@ func (s *testBundleSuite) TestClone(c *C) { c.Assert(newBundle, DeepEquals, &Bundle{ID: GroupID(2), Rules: []*Rule{{ID: "121"}}}) } +func (s *testBundleSuite) TestObjectID(c *C) { + type TestCase struct { + name string + bundleID string + expectedID int64 + err error + } + tests := []TestCase{ + {"non tidb bundle", "pd", 0, ErrInvalidBundleIDFormat}, + {"id of words", "TiDB_DDL_foo", 0, ErrInvalidBundleID}, + {"id of words and nums", "TiDB_DDL_3x", 0, ErrInvalidBundleID}, + {"id of floats", "TiDB_DDL_3.0", 0, ErrInvalidBundleID}, + {"id of negatives", "TiDB_DDL_-10", 0, ErrInvalidBundleID}, + {"id of positive integer", "TiDB_DDL_10", 10, nil}, + } + for _, t := range tests { + comment := Commentf("%s", t.name) + bundle := Bundle{ID: t.bundleID} + id, err := bundle.ObjectID() + if t.err == nil { + c.Assert(err, IsNil, comment) + c.Assert(id, Equals, t.expectedID, comment) + } else { + c.Assert(errors.Is(err, t.err), IsTrue, comment) + } + } +} + +func (s *testBundleSuite) TestGetLeaderDCByBundle(c *C) { + testcases := []struct { + name string + bundle *Bundle + expectedDC string + }{ + { + name: "only leader", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "12", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"bj"}, + }, + }, + Count: 1, + }, + }, + }, + expectedDC: "bj", + }, + { + name: "no leader", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "12", + Role: Voter, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"bj"}, + }, + }, + Count: 3, + }, + }, + }, + expectedDC: "", + }, + { + name: "voter and leader", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "11", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"sh"}, + }, + }, + Count: 1, + }, + { + ID: "12", + Role: Voter, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"bj"}, + }, + }, + Count: 3, + }, + }, + }, + expectedDC: "sh", + }, + { + name: "wrong label key", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "11", + Role: Leader, + Constraints: Constraints{ + { + Key: "fake", + Op: In, + Values: []string{"sh"}, + }, + }, + Count: 1, + }, + }, + }, + expectedDC: "", + }, + { + name: "wrong operator", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "11", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: NotIn, + Values: []string{"sh"}, + }, + }, + Count: 1, + }, + }, + }, + expectedDC: "", + }, + { + name: "leader have multi values", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "11", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"sh", "bj"}, + }, + }, + Count: 1, + }, + }, + }, + expectedDC: "", + }, + { + name: "irrelvant rules", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "15", + Role: Leader, + Constraints: Constraints{ + { + Key: EngineLabelKey, + Op: NotIn, + Values: []string{EngineLabelTiFlash}, + }, + }, + Count: 1, + }, + { + ID: "14", + Role: Leader, + Constraints: Constraints{ + { + Key: "disk", + Op: NotIn, + Values: []string{"ssd", "hdd"}, + }, + }, + Count: 1, + }, + { + ID: "13", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"bj"}, + }, + }, + Count: 1, + }, + }, + }, + expectedDC: "bj", + }, + { + name: "multi leaders 1", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "16", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"sh"}, + }, + }, + Count: 2, + }, + }, + }, + expectedDC: "", + }, + { + name: "multi leaders 2", + bundle: &Bundle{ + ID: GroupID(1), + Rules: []*Rule{ + { + ID: "17", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"sh"}, + }, + }, + Count: 1, + }, + { + ID: "18", + Role: Leader, + Constraints: Constraints{ + { + Key: "zone", + Op: In, + Values: []string{"bj"}, + }, + }, + Count: 1, + }, + }, + }, + expectedDC: "sh", + }, + } + for _, testcase := range testcases { + comment := Commentf("%s", testcase.name) + result, ok := testcase.bundle.GetLeaderDC("zone") + if len(testcase.expectedDC) > 0 { + c.Assert(ok, Equals, true, comment) + } else { + c.Assert(ok, Equals, false, comment) + } + c.Assert(result, Equals, testcase.expectedDC, comment) + } +} + func (s *testBundleSuite) TestApplyPlacmentSpec(c *C) { type TestCase struct { name string diff --git a/ddl/placement/utils.go b/ddl/placement/utils.go deleted file mode 100644 index 5b12f10e2d243..0000000000000 --- a/ddl/placement/utils.go +++ /dev/null @@ -1,79 +0,0 @@ -// Copyright 2020 PingCAP, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// See the License for the specific language governing permissions and -// limitations under the License. - -package placement - -import ( - "encoding/hex" - "strconv" - "strings" - - "github.com/pingcap/errors" - "github.com/pingcap/tidb/tablecodec" - "github.com/pingcap/tidb/util/codec" -) - -// ObjectIDFromGroupID extracts the db/table/partition ID from the group ID -func ObjectIDFromGroupID(groupID string) (int64, error) { - // If the rule doesn't come from TiDB, skip it. - if !strings.HasPrefix(groupID, BundleIDPrefix) { - return 0, nil - } - id, err := strconv.ParseInt(groupID[len(BundleIDPrefix):], 10, 64) - if err != nil || id <= 0 { - return 0, errors.Errorf("Rule %s doesn't include an id", groupID) - } - return id, nil -} - -// BuildPlacementDropBundle builds the bundle to drop placement rules. -func BuildPlacementDropBundle(partitionID int64) *Bundle { - return &Bundle{ - ID: GroupID(partitionID), - } -} - -// BuildPlacementCopyBundle copies a new bundle from the old, with a new name and a new key range. -func BuildPlacementCopyBundle(oldBundle *Bundle, newID int64) *Bundle { - newBundle := oldBundle.Clone() - newBundle.ID = GroupID(newID) - startKey := hex.EncodeToString(codec.EncodeBytes(nil, tablecodec.GenTableRecordPrefix(newID))) - endKey := hex.EncodeToString(codec.EncodeBytes(nil, tablecodec.GenTableRecordPrefix(newID+1))) - for _, rule := range newBundle.Rules { - rule.GroupID = newBundle.ID - rule.StartKeyHex = startKey - rule.EndKeyHex = endKey - } - return newBundle -} - -// GetLeaderDCByBundle returns the leader's DC by Bundle if found -func GetLeaderDCByBundle(bundle *Bundle, dcLabelKey string) (string, bool) { - for _, rule := range bundle.Rules { - if isValidLeaderRule(rule, dcLabelKey) { - return rule.Constraints[0].Values[0], true - } - } - return "", false -} - -func isValidLeaderRule(rule *Rule, dcLabelKey string) bool { - if rule.Role == Leader && rule.Count == 1 { - for _, con := range rule.Constraints { - if con.Op == In && con.Key == dcLabelKey && len(con.Values) == 1 { - return true - } - } - } - return false -} diff --git a/ddl/placement/utils_test.go b/ddl/placement/utils_test.go deleted file mode 100644 index 10941e0663455..0000000000000 --- a/ddl/placement/utils_test.go +++ /dev/null @@ -1,237 +0,0 @@ -// Copyright 2020 PingCAP, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// See the License for the specific language governing permissions and -// limitations under the License. - -package placement - -import ( - . "github.com/pingcap/check" -) - -var _ = Suite(&testUtilsSuite{}) - -type testUtilsSuite struct{} - -func (t *testUtilsSuite) TestObjectIDFromGroupID(c *C) { - testCases := []struct { - bundleID string - expectedID int64 - expectErr bool - }{ - {"pd", 0, false}, - {"TiDB_DDL_foo", 0, true}, - {"TiDB_DDL_3x", 0, true}, - {"TiDB_DDL_3.0", 0, true}, - {"TiDB_DDL_-10", 0, true}, - {"TiDB_DDL_10", 10, false}, - } - for _, testCase := range testCases { - id, err := ObjectIDFromGroupID(testCase.bundleID) - if testCase.expectErr { - c.Assert(err, NotNil) - } else { - c.Assert(id, Equals, testCase.expectedID) - } - } -} - -func (t *testUtilsSuite) TestGetLeaderDCByBundle(c *C) { - testcases := []struct { - name string - bundle *Bundle - expectedDC string - }{ - { - name: "only leader", - bundle: &Bundle{ - ID: GroupID(1), - Rules: []*Rule{ - { - ID: "12", - Role: Leader, - Constraints: []Constraint{ - { - Key: "zone", - Op: In, - Values: []string{"bj"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 1, - }, - }, - }, - expectedDC: "bj", - }, - { - name: "no leader", - bundle: &Bundle{ - ID: GroupID(1), - Rules: []*Rule{ - { - ID: "12", - Role: Voter, - Constraints: []Constraint{ - { - Key: "zone", - Op: In, - Values: []string{"bj"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 3, - }, - }, - }, - expectedDC: "", - }, - { - name: "voter and leader", - bundle: &Bundle{ - ID: GroupID(1), - Rules: []*Rule{ - { - ID: "11", - Role: Leader, - Constraints: []Constraint{ - { - Key: "zone", - Op: In, - Values: []string{"sh"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 1, - }, - { - ID: "12", - Role: Voter, - Constraints: []Constraint{ - { - Key: "zone", - Op: In, - Values: []string{"bj"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 3, - }, - }, - }, - expectedDC: "sh", - }, - { - name: "wrong label key", - bundle: &Bundle{ - ID: GroupID(1), - Rules: []*Rule{ - { - ID: "11", - Role: Leader, - Constraints: []Constraint{ - { - Key: "fake", - Op: In, - Values: []string{"sh"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 1, - }, - }, - }, - expectedDC: "", - }, - { - name: "wrong operator", - bundle: &Bundle{ - ID: GroupID(1), - Rules: []*Rule{ - { - ID: "11", - Role: Leader, - Constraints: []Constraint{ - { - Key: "zone", - Op: NotIn, - Values: []string{"sh"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 1, - }, - }, - }, - expectedDC: "", - }, - { - name: "leader have multi values", - bundle: &Bundle{ - ID: GroupID(1), - Rules: []*Rule{ - { - ID: "11", - Role: Leader, - Constraints: []Constraint{ - { - Key: "zone", - Op: In, - Values: []string{"sh", "bj"}, - }, - { - Key: EngineLabelKey, - Op: NotIn, - Values: []string{EngineLabelTiFlash}, - }, - }, - Count: 1, - }, - }, - }, - expectedDC: "", - }, - } - for _, testcase := range testcases { - c.Log(testcase.name) - result, ok := GetLeaderDCByBundle(testcase.bundle, "zone") - if len(testcase.expectedDC) > 0 { - c.Assert(ok, Equals, true) - } else { - c.Assert(ok, Equals, false) - } - c.Assert(result, Equals, testcase.expectedDC) - } -} diff --git a/distsql/request_builder.go b/distsql/request_builder.go index c71d9c27d97ff..94f3debb325f4 100644 --- a/distsql/request_builder.go +++ b/distsql/request_builder.go @@ -634,7 +634,7 @@ func VerifyTxnScope(txnScope string, physicalTableID int64, is infoschema.InfoSc if !ok { return true } - leaderDC, ok := placement.GetLeaderDCByBundle(bundle, placement.DCLabelKey) + leaderDC, ok := bundle.GetLeaderDC(placement.DCLabelKey) if !ok { return true } diff --git a/executor/infoschema_reader.go b/executor/infoschema_reader.go index 9937a3ffcb719..e087ccc757ad0 100644 --- a/executor/infoschema_reader.go +++ b/executor/infoschema_reader.go @@ -1921,13 +1921,13 @@ func (e *memtableRetriever) setDataForPlacementPolicy(ctx sessionctx.Context) er is := ctx.GetInfoSchema().(infoschema.InfoSchema) var rows [][]types.Datum for _, bundle := range is.RuleBundles() { - id, err := placement.ObjectIDFromGroupID(bundle.ID) + id, err := bundle.ObjectID() if err != nil { + if err == placement.ErrInvalidBundleIDFormat { + continue + } return errors.Wrapf(err, "Restore bundle %s failed", bundle.ID) } - if id == 0 { - continue - } // Currently, only partitions have placement rules. var tbName, dbName, ptName string skip := true diff --git a/session/session.go b/session/session.go index 3b0465e6a24ae..6e455e2436c22 100644 --- a/session/session.go +++ b/session/session.go @@ -2882,7 +2882,7 @@ func (s *session) checkPlacementPolicyBeforeCommit() error { err = ddl.ErrInvalidPlacementPolicyCheck.GenWithStackByArgs(errMsg) break } - dcLocation, ok := placement.GetLeaderDCByBundle(bundle, placement.DCLabelKey) + dcLocation, ok := bundle.GetLeaderDC(placement.DCLabelKey) if !ok { errMsg := fmt.Sprintf("table %v's leader placement policy is not defined", tableName) if len(partitionName) > 0 {