*: support resource_group sql hint (#41678)

* *: support resource_group sql hint Signed-off-by: BornChanger <[email protected]> * *: refine case Signed-off-by: BornChanger <[email protected]> * *: make case stable Signed-off-by: BornChanger <[email protected]> * *: code format Signed-off-by: BornChanger <[email protected]> * *: relocate the cases Signed-off-by: BornChanger <[email protected]> * *: fix bazel attribute Signed-off-by: BornChanger <[email protected]> * *: fix bazel Signed-off-by: BornChanger <[email protected]> * *: code fmt Signed-off-by: BornChanger <[email protected]> * *: code format Signed-off-by: BornChanger <[email protected]> * *: address comments Signed-off-by: BornChanger <[email protected]> * *: address comments for 2nd round Signed-off-by: BornChanger <[email protected]> * *: polish check logic Signed-off-by: BornChanger <[email protected]> --------- Signed-off-by: BornChanger <[email protected]>
pingcap · Mar 6, 2023 · 17bbc29 · 17bbc29
1 parent 9958a93
commit 17bbc29
Show file tree

Hide file tree

Showing 14 changed files with 223 additions and 168 deletions.
diff --git a/ddl/BUILD.bazel b/ddl/BUILD.bazel
@@ -198,7 +198,6 @@ go_test(
         "primary_key_handle_test.go",
         "reorg_partition_test.go",
         "repair_table_test.go",
-        "resource_group_test.go",
         "restart_test.go",
         "rollingback_test.go",
         "schema_test.go",

diff --git a/ddl/resourcegroup/BUILD.bazel b/ddl/resourcegroup/BUILD.bazel
@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
 
 go_library(
     name = "resourcegroup",
@@ -14,16 +14,3 @@ go_library(
         "@com_github_pingcap_kvproto//pkg/resource_manager",
     ],
 )
-
-go_test(
-    name = "resourcegroup_test",
-    timeout = "short",
-    srcs = ["group_test.go"],
-    embed = [":resourcegroup"],
-    flaky = True,
-    deps = [
-        "//parser/model",
-        "@com_github_pingcap_kvproto//pkg/resource_manager",
-        "@com_github_stretchr_testify//require",
-    ],
-)
diff --git a/ddl/resourcegroup/group_test.go b/ddl/resourcegroup/group_test.go
diff --git a/ddl/resourcegrouptest/BUILD.bazel b/ddl/resourcegrouptest/BUILD.bazel
@@ -0,0 +1,20 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_test")
+
+go_test(
+    name = "resourcegrouptest_test",
+    timeout = "short",
+    srcs = ["resource_group_test.go"],
+    flaky = True,
+    deps = [
+        "//ddl/internal/callback",
+        "//ddl/resourcegroup",
+        "//domain",
+        "//domain/infosync",
+        "//errno",
+        "//parser/model",
+        "//sessionctx",
+        "//testkit",
+        "@com_github_pingcap_kvproto//pkg/resource_manager",
+        "@com_github_stretchr_testify//require",
+    ],
+)
diff --git a/ddl/resource_group_test.go → ddl/resourcegrouptest/resource_group_test.go b/ddl/resource_group_test.go → ddl/resourcegrouptest/resource_group_test.go
@@ -12,13 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package ddl_test
+package resourcegrouptest_test
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
+	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
 	"github.com/pingcap/tidb/ddl/internal/callback"
+	"github.com/pingcap/tidb/ddl/resourcegroup"
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/domain/infosync"
 	mysql "github.com/pingcap/tidb/errno"
@@ -176,3 +179,143 @@ func testResourceGroupNameFromIS(t *testing.T, ctx sessionctx.Context, name stri
 	g, _ := dom.InfoSchema().ResourceGroupByName(model.NewCIStr(name))
 	return g
 }
+
+func TestResourceGroupHint(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("create table t1(c1 int)")
+	tk.MustExec("insert into t1 values(1)")
+
+	tk.MustExec("set global tidb_enable_resource_control='on'")
+	tk.MustExec("create resource group rg1 ru_per_sec=1000")
+	tk.MustQuery("select /*+ resource_group(default) */ * from t1")
+	tk.MustQuery("select /*+ resource_group(rg1) */ * from t1")
+	tk.MustQuery("select /*+ resource_group(rg1) resource_group(default) */ * from t1")
+	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 RESOURCE_GROUP() is defined more than once, only the last definition takes effect: RESOURCE_GROUP(default)"))
+	tk.MustQuery("select /*+ resource_group(rg1) */ DB, RESOURCE_GROUP from information_schema.processlist").Check(testkit.Rows("test rg1"))
+	tk.MustQuery("select DB, RESOURCE_GROUP from information_schema.processlist").Check(testkit.Rows("test "))
+	tk.MustExec("set global tidb_enable_resource_control='off'")
+	tk.MustQuery("select /*+ resource_group(rg1) */ DB, RESOURCE_GROUP from information_schema.processlist").Check(testkit.Rows("test "))
+	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 8250 Resource control feature is disabled. Run `SET GLOBAL tidb_enable_resource_control='on'` to enable the feature"))
+}
+
+func TestNewResourceGroupFromOptions(t *testing.T) {
+	type TestCase struct {
+		name      string
+		groupName string
+		input     *model.ResourceGroupSettings
+		output    *rmpb.ResourceGroup
+		err       error
+	}
+	var tests []TestCase
+	groupName := "test"
+	tests = append(tests, TestCase{
+		name:  "empty 1",
+		input: &model.ResourceGroupSettings{},
+		err:   resourcegroup.ErrUnknownResourceGroupMode,
+	})
+
+	tests = append(tests, TestCase{
+		name:  "empty 2",
+		input: nil,
+		err:   resourcegroup.ErrInvalidGroupSettings,
+	})
+
+	tests = append(tests, TestCase{
+		name: "normal case: ru case 1",
+		input: &model.ResourceGroupSettings{
+			RURate: 2000,
+		},
+		output: &rmpb.ResourceGroup{
+			Name: groupName,
+			Mode: rmpb.GroupMode_RUMode,
+			RUSettings: &rmpb.GroupRequestUnitSettings{
+				RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 2000}},
+			},
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name: "normal case: ru case 2",
+		input: &model.ResourceGroupSettings{
+			RURate: 5000,
+		},
+		output: &rmpb.ResourceGroup{
+			Name: groupName,
+			Mode: rmpb.GroupMode_RUMode,
+			RUSettings: &rmpb.GroupRequestUnitSettings{
+				RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 5000}},
+			},
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name: "error case: native case 1",
+		input: &model.ResourceGroupSettings{
+			CPULimiter:       "8",
+			IOReadBandwidth:  "3000MB/s",
+			IOWriteBandwidth: "3000Mi",
+		},
+		err: resourcegroup.ErrUnknownResourceGroupMode,
+	})
+
+	tests = append(tests, TestCase{
+		name: "error case: native case 2",
+		input: &model.ResourceGroupSettings{
+			CPULimiter:       "8c",
+			IOReadBandwidth:  "3000Mi",
+			IOWriteBandwidth: "3000Mi",
+		},
+		err: resourcegroup.ErrUnknownResourceGroupMode,
+	})
+
+	tests = append(tests, TestCase{
+		name: "error case: native case 3",
+		input: &model.ResourceGroupSettings{
+			CPULimiter:       "8",
+			IOReadBandwidth:  "3000G",
+			IOWriteBandwidth: "3000MB",
+		},
+		err: resourcegroup.ErrUnknownResourceGroupMode,
+	})
+
+	tests = append(tests, TestCase{
+		name: "error case: duplicated mode",
+		input: &model.ResourceGroupSettings{
+			CPULimiter:       "8",
+			IOReadBandwidth:  "3000Mi",
+			IOWriteBandwidth: "3000Mi",
+			RURate:           1000,
+		},
+		err: resourcegroup.ErrInvalidResourceGroupDuplicatedMode,
+	})
+
+	tests = append(tests, TestCase{
+		name:      "error case: duplicated mode",
+		groupName: "test_group_too_looooooooooooooooooooooooooooooooooooooooooooooooong",
+		input: &model.ResourceGroupSettings{
+			CPULimiter:       "8",
+			IOReadBandwidth:  "3000Mi",
+			IOWriteBandwidth: "3000Mi",
+			RURate:           1000,
+		},
+		err: resourcegroup.ErrTooLongResourceGroupName,
+	})
+
+	for _, test := range tests {
+		name := groupName
+		if len(test.groupName) > 0 {
+			name = test.groupName
+		}
+		group, err := resourcegroup.NewGroupFromOptions(name, test.input)
+		comment := fmt.Sprintf("[%s]\nerr1 %s\nerr2 %s", test.name, err, test.err)
+		if test.err != nil {
+			require.ErrorIs(t, err, test.err, comment)
+		} else {
+			require.NoError(t, err, comment)
+			require.Equal(t, test.output, group)
+		}
+	}
+}
diff --git a/parser/ast/misc.go b/parser/ast/misc.go
@@ -3766,6 +3766,8 @@ func (n *TableOptimizerHint) Restore(ctx *format.RestoreCtx) error {
 	switch n.HintName.L {
 	case "max_execution_time":
 		ctx.WritePlainf("%d", n.HintData.(uint64))
+	case "resource_group":
+		ctx.WriteName(n.HintData.(string))
 	case "nth_plan":
 		ctx.WritePlainf("%d", n.HintData.(int64))
 	case "tidb_hj", "tidb_smj", "tidb_inlj", "hash_join", "hash_join_build", "hash_join_probe", "merge_join", "inl_join", "broadcast_join", "shuffle_join", "inl_hash_join", "inl_merge_join", "leading":

diff --git a/parser/ast/misc_test.go b/parser/ast/misc_test.go
@@ -303,6 +303,8 @@ func TestTableOptimizerHintRestore(t *testing.T) {
 		{"READ_FROM_STORAGE(@sel TIFLASH[t1, t2])", "READ_FROM_STORAGE(@`sel` TIFLASH[`t1`, `t2`])"},
 		{"READ_FROM_STORAGE(@sel TIFLASH[t1 partition(p0)])", "READ_FROM_STORAGE(@`sel` TIFLASH[`t1` PARTITION(`p0`)])"},
 		{"TIME_RANGE('2020-02-02 10:10:10','2020-02-02 11:10:10')", "TIME_RANGE('2020-02-02 10:10:10', '2020-02-02 11:10:10')"},
+		{"RESOURCE_GROUP(rg1)", "RESOURCE_GROUP(`rg1`)"},
+		{"RESOURCE_GROUP(`default`)", "RESOURCE_GROUP(`default`)"},
 	}
 	extractNodeFunc := func(node ast.Node) ast.Node {
 		return node.(*ast.SelectStmt).TableHints[0]

diff --git a/parser/hintparser.go b/parser/hintparser.go