session: support set current session's resource group name

DEPS.bzl

@@ -4101,8 +4101,8 @@ def go_deps():
         name = "com_github_tikv_client_go_v2",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/client-go/v2",
-        sum = "h1:hZ0aZfGAaS3rM0hYiMq6j0xHpVFgoYw6KeqrniGb7Zo=",
-        version = "v2.0.6-0.20230227032358-40a82457ebaa",
+        sum = "h1:9exwxGqABskTj2pwtTqwziT/nGJ2N3Mh4gFAdP7GsrE=",
+        version = "v2.0.6-0.20230228091502-e2da5527026f",
     )
     go_repository(
         name = "com_github_tikv_pd",
@@ -4116,8 +4116,8 @@ def go_deps():
         name = "com_github_tikv_pd_client",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/pd/client",
-        sum = "h1:hauBQBHSyrUxAI0zvkTiBKd472c+Iy+aY0Jd+b9VOJ8=",
-        version = "v0.0.0-20230209034200-6d23a31c24be",
+        sum = "h1:wZVfR5IsFod5Lym9lgjAaZ9TFsOqZ1iUaxockPacai8=",
+        version = "v0.0.0-20230224101107-b2f2a8219a40",
     )
     go_repository(
         name = "com_github_timakin_bodyclose",

executor/simple.go

@@ -216,6 +216,8 @@ func (e *SimpleExec) Next(ctx context.Context, req *chunk.Chunk) (err error) {
 		err = e.executeShutdown(x)
 	case *ast.AdminStmt:
 		err = e.executeAdmin(x)
+	case *ast.SetResourceGroupStmt:
+		err = e.executeSetResourceGroupName(x)
 	}
 	e.done = true
 	return err
@@ -2840,3 +2842,14 @@ func (e *SimpleExec) executeAdminFlushPlanCache(s *ast.AdminStmt) error {
 	}
 	return nil
 }
+
+func (e *SimpleExec) executeSetResourceGroupName(s *ast.SetResourceGroupStmt) error {
+	if s.Name.L != "" {
+		if _, ok := e.is.ResourceGroupByName(s.Name); !ok {
+			return infoschema.ErrResourceGroupNotExists.GenWithStackByArgs(s.Name.O)
+		}
+	}
+
+	e.ctx.GetSessionVars().ResourceGroupName = s.Name.L
+	return nil
+}

executor/simple_test.go

@@ -133,3 +133,17 @@ func TestUserAttributes(t *testing.T) {
 	rootTK.MustExec("alter user usr1 resource group rg1")
 	rootTK.MustQuery("select user_attributes from mysql.user where user = 'usr1'").Check(testkit.Rows(`{"metadata": {"comment": "comment1"}, "resource_group": "rg1"}`))
 }
+
+func TestSetResourceGroup(t *testing.T) {
+	store, _ := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("SET GLOBAL tidb_enable_resource_control='on'")
+	tk.MustContainErrMsg("SET RESOURCE GROUP rg1", "Unknown resource group 'rg1'")
+	tk.MustExec("CREATE RESOURCE GROUP rg1 ru_per_sec = 100")
+	tk.MustQuery("SELECT RESOURCE_GROUP FROM information_schema.processlist").Check(testkit.Rows(""))
+	tk.MustExec("SET RESOURCE GROUP rg1")
+	tk.MustQuery("SELECT RESOURCE_GROUP FROM information_schema.processlist").Check(testkit.Rows("rg1"))
+	tk.MustExec("SET RESOURCE GROUP ``")
+	tk.MustQuery("SELECT RESOURCE_GROUP FROM information_schema.processlist").Check(testkit.Rows(""))
+}

go.mod

@@ -92,8 +92,8 @@ require (
 	github.com/stretchr/testify v1.8.1
 	github.com/tdakkota/asciicheck v0.1.1
 	github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2
-	github.com/tikv/client-go/v2 v2.0.6-0.20230227032358-40a82457ebaa
-	github.com/tikv/pd/client v0.0.0-20230209034200-6d23a31c24be
+	github.com/tikv/client-go/v2 v2.0.6-0.20230228091502-e2da5527026f
+	github.com/tikv/pd/client v0.0.0-20230224101107-b2f2a8219a40
 	github.com/timakin/bodyclose v0.0.0-20221125081123-e39cf3fc478e
 	github.com/twmb/murmur3 v1.1.6
 	github.com/uber/jaeger-client-go v2.22.1+incompatible

go.sum

@@ -937,10 +937,10 @@ github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJf
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW93SG+q0F8KI+yFrcIDT4c/RNoc4=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
-github.com/tikv/client-go/v2 v2.0.6-0.20230227032358-40a82457ebaa h1:hZ0aZfGAaS3rM0hYiMq6j0xHpVFgoYw6KeqrniGb7Zo=
-github.com/tikv/client-go/v2 v2.0.6-0.20230227032358-40a82457ebaa/go.mod h1:Kw6+qOCoLSymJLtOPTym6wkH97Ej0Yna4pDv0ak3hMc=
-github.com/tikv/pd/client v0.0.0-20230209034200-6d23a31c24be h1:hauBQBHSyrUxAI0zvkTiBKd472c+Iy+aY0Jd+b9VOJ8=
-github.com/tikv/pd/client v0.0.0-20230209034200-6d23a31c24be/go.mod h1:ryhYHDwupsZHeOOF/N7So+1hbtAnuw0K2A+pKOElSVs=
+github.com/tikv/client-go/v2 v2.0.6-0.20230228091502-e2da5527026f h1:9exwxGqABskTj2pwtTqwziT/nGJ2N3Mh4gFAdP7GsrE=
+github.com/tikv/client-go/v2 v2.0.6-0.20230228091502-e2da5527026f/go.mod h1:a2N4+WwNsicV1E2+KQFg4CDqeoxvOyphZXmldX/TXDU=
+github.com/tikv/pd/client v0.0.0-20230224101107-b2f2a8219a40 h1:wZVfR5IsFod5Lym9lgjAaZ9TFsOqZ1iUaxockPacai8=
+github.com/tikv/pd/client v0.0.0-20230224101107-b2f2a8219a40/go.mod h1:j94ECCo0drzB/GsMFaUjESjLXNvIbed5PWJXcd8eggo=
 github.com/timakin/bodyclose v0.0.0-20221125081123-e39cf3fc478e h1:MV6KaVu/hzByHP0UvJ4HcMGE/8a6A4Rggc/0wx2AvJo=
 github.com/timakin/bodyclose v0.0.0-20221125081123-e39cf3fc478e/go.mod h1:27bSVNWSBOHm+qRp1T9qzaIpsWEP6TbUnei/43HK+PQ=
 github.com/tklauser/go-sysconf v0.3.9/go.mod h1:11DU/5sG7UexIrp/O6g35hrWzu0JxlwQ3LSFUzyeuhs=

parser/ast/misc.go

@@ -58,6 +58,7 @@ var (
 	_ StmtNode = &HelpStmt{}
 	_ StmtNode = &PlanReplayerStmt{}
 	_ StmtNode = &CompactTableStmt{}
+	_ StmtNode = &SetResourceGroupStmt{}
 
 	_ Node = &PrivElem{}
 	_ Node = &VariableAssignment{}
@@ -3860,3 +3861,25 @@ var NewHexLiteral func(string) (interface{}, error)
 
 // NewBitLiteral creates a types.BitLiteral value, it's provided by parser driver.
 var NewBitLiteral func(string) (interface{}, error)
+
+// SetResourceGroupStmt is a statement to set the resource group name for current session.
+type SetResourceGroupStmt struct {
+	stmtNode
+	Name model.CIStr
+}
+
+func (n *SetResourceGroupStmt) Restore(ctx *format.RestoreCtx) error {
+	ctx.WriteKeyWord("SET RESOURCE GROUP ")
+	ctx.WriteName(n.Name.O)
+	return nil
+}
+
+// Accept implements Node Accept interface.
+func (n *SetResourceGroupStmt) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	n = newNode.(*SetResourceGroupStmt)
+	return v.Leave(n)
+}