session: support set current session's resource group name

executor/reload_expr_pushdown_blacklist.go

@@ -257,6 +257,7 @@ var funcName2Alias = map[string]string{
 	"collation":                  ast.Collation,
 	"connection_id":              ast.ConnectionID,
 	"current_user":               ast.CurrentUser,
+	"current_resource_group":     ast.CurrentResourceGroup,
 	"current_role":               ast.CurrentRole,
 	"database":                   ast.Database,
 	"found_rows":                 ast.FoundRows,

executor/showtest/show_test.go

@@ -1522,7 +1522,7 @@ func TestShowBuiltin(t *testing.T) {
 	res := tk.MustQuery("show builtins;")
 	require.NotNil(t, res)
 	rows := res.Rows()
-	const builtinFuncNum = 285
+	const builtinFuncNum = 286
 	require.Equal(t, builtinFuncNum, len(rows))
 	require.Equal(t, rows[0][0].(string), "abs")
 	require.Equal(t, rows[builtinFuncNum-1][0].(string), "yearweek")

executor/simple.go

@@ -209,6 +209,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
@@ -2833,3 +2835,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,28 @@ 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.MustExec("ALTER USER `root` RESOURCE GROUP `rg1`")
+	tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows(""))
+	require.NoError(t, tk.Session().Auth(&auth.UserIdentity{Username: "root", Hostname: "%"}, nil, nil))
+	tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows("rg1"))
+
+	tk.MustExec("CREATE RESOURCE GROUP rg2 ru_per_sec = 200")
+	tk.MustExec("SET RESOURCE GROUP `rg2`")
+	tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows("rg2"))
+	tk.MustExec("SET RESOURCE GROUP ``")
+	tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows(""))
+
+	tk.RefreshSession()
+	require.NoError(t, tk.Session().Auth(&auth.UserIdentity{Username: "root", Hostname: "%"}, nil, nil))
+	tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows("rg1"))
+}

expression/builtin.go

@@ -746,10 +746,12 @@ var funcs = map[string]functionClass{
 	ast.WeightString:    &weightStringFunctionClass{baseFunctionClass{ast.WeightString, 1, 3}},
 
 	// information functions
-	ast.ConnectionID: &connectionIDFunctionClass{baseFunctionClass{ast.ConnectionID, 0, 0}},
-	ast.CurrentUser:  &currentUserFunctionClass{baseFunctionClass{ast.CurrentUser, 0, 0}},
-	ast.CurrentRole:  &currentRoleFunctionClass{baseFunctionClass{ast.CurrentRole, 0, 0}},
-	ast.Database:     &databaseFunctionClass{baseFunctionClass{ast.Database, 0, 0}},
+	ast.ConnectionID:         &connectionIDFunctionClass{baseFunctionClass{ast.ConnectionID, 0, 0}},
+	ast.CurrentUser:          &currentUserFunctionClass{baseFunctionClass{ast.CurrentUser, 0, 0}},
+	ast.CurrentRole:          &currentRoleFunctionClass{baseFunctionClass{ast.CurrentRole, 0, 0}},
+	ast.Database:             &databaseFunctionClass{baseFunctionClass{ast.Database, 0, 0}},
+	ast.CurrentResourceGroup: &currentResourceGroupFunctionClass{baseFunctionClass{ast.CurrentResourceGroup, 0, 0}},
+
 	// This function is a synonym for DATABASE().
 	// See http://dev.mysql.com/doc/refman/5.7/en/information-functions.html#function_schema
 	ast.Schema:       &databaseFunctionClass{baseFunctionClass{ast.Schema, 0, 0}},

expression/builtin_info.go

@@ -44,6 +44,7 @@ var (
 	_ functionClass = &foundRowsFunctionClass{}
 	_ functionClass = &currentUserFunctionClass{}
 	_ functionClass = &currentRoleFunctionClass{}
+	_ functionClass = &currentResourceGroupFunctionClass{}
 	_ functionClass = &userFunctionClass{}
 	_ functionClass = &connectionIDFunctionClass{}
 	_ functionClass = &lastInsertIDFunctionClass{}
@@ -69,6 +70,7 @@ var (
 	_ builtinFunc = &builtinDatabaseSig{}
 	_ builtinFunc = &builtinFoundRowsSig{}
 	_ builtinFunc = &builtinCurrentUserSig{}
+	_ builtinFunc = &builtinCurrentResourceGroupSig{}
 	_ builtinFunc = &builtinUserSig{}
 	_ builtinFunc = &builtinConnectionIDSig{}
 	_ builtinFunc = &builtinLastInsertIDSig{}
@@ -245,6 +247,41 @@ func (b *builtinCurrentRoleSig) evalString(row chunk.Row) (res string, isNull bo
 	return res, false, nil
 }
 
+type currentResourceGroupFunctionClass struct {
+	baseFunctionClass
+}
+
+func (c *currentResourceGroupFunctionClass) getFunction(ctx sessionctx.Context, args []Expression) (builtinFunc, error) {
+	if err := c.verifyArgs(args); err != nil {
+		return nil, err
+	}
+	bf, err := newBaseBuiltinFuncWithTp(ctx, c.funcName, args, types.ETString)
+	if err != nil {
+		return nil, err
+	}
+	bf.tp.SetFlen(64)
+	sig := &builtinCurrentResourceGroupSig{bf}
+	return sig, nil
+}
+
+type builtinCurrentResourceGroupSig struct {
+	baseBuiltinFunc
+}
+
+func (b *builtinCurrentResourceGroupSig) Clone() builtinFunc {
+	newSig := &builtinCurrentResourceGroupSig{}
+	newSig.cloneFrom(&b.baseBuiltinFunc)
+	return newSig
+}
+
+func (b *builtinCurrentResourceGroupSig) evalString(row chunk.Row) (res string, isNull bool, err error) {
+	data := b.ctx.GetSessionVars()
+	if data == nil {
+		return "", true, errors.Errorf("Missing session variable when eval builtin")
+	}
+	return data.ResourceGroupName, false, nil
+}
+
 type userFunctionClass struct {
 	baseFunctionClass
 }

expression/builtin_info_test.go

@@ -96,6 +96,20 @@ func TestCurrentUser(t *testing.T) {
 	require.Equal(t, f.PbCode(), f.Clone().PbCode())
 }
 
+func TestCurrentResourceGroup(t *testing.T) {
+	ctx := mock.NewContext()
+	sessionVars := ctx.GetSessionVars()
+	sessionVars.ResourceGroupName = "rg1"
+
+	fc := funcs[ast.CurrentResourceGroup]
+	f, err := fc.getFunction(ctx, nil)
+	require.NoError(t, err)
+	d, err := evalBuiltinFunc(f, chunk.Row{})
+	require.NoError(t, err)
+	require.Equal(t, "rg1", d.GetString())
+	require.Equal(t, f.PbCode(), f.Clone().PbCode())
+}
+
 func TestCurrentRole(t *testing.T) {
 	ctx := mock.NewContext()
 	fc := funcs[ast.CurrentRole]

expression/builtin_info_vec.go

@@ -119,6 +119,23 @@ func (b *builtinCurrentUserSig) vecEvalString(input *chunk.Chunk, result *chunk.
 	return nil
 }
 
+func (b *builtinCurrentResourceGroupSig) vectorized() bool {
+	return true
+}
+
+func (b *builtinCurrentResourceGroupSig) vecEvalString(input *chunk.Chunk, result *chunk.Column) error {
+	data := b.ctx.GetSessionVars()
+	if data == nil {
+		return errors.Errorf("Missing session variable when eval builtin")
+	}
+	n := input.NumRows()
+	result.ReserveString(n)
+	for i := 0; i < n; i++ {
+		result.AppendString(data.ResourceGroupName)
+	}
+	return nil
+}
+
 func (b *builtinCurrentRoleSig) vectorized() bool {
 	return true
 }

expression/builtin_info_vec_test.go

@@ -79,6 +79,9 @@ var vecBuiltinInfoCases = map[string][]vecExprBenchCase{
 	ast.CurrentRole: {
 		{retEvalType: types.ETString, childrenTypes: []types.EvalType{}},
 	},
+	ast.CurrentResourceGroup: {
+		{retEvalType: types.ETString, childrenTypes: []types.EvalType{}},
+	},
 	ast.TiDBIsDDLOwner: {
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{}},
 	},

expression/collation.go

@@ -274,7 +274,7 @@ func deriveCollation(ctx sessionctx.Context, funcName string, args []Expression,
 			}
 			return CheckAndDeriveCollationFromExprs(ctx, funcName, retType, fieldArgs...)
 		}
-	case ast.Database, ast.User, ast.CurrentUser, ast.Version, ast.CurrentRole, ast.TiDBVersion:
+	case ast.Database, ast.User, ast.CurrentUser, ast.Version, ast.CurrentRole, ast.TiDBVersion, ast.CurrentResourceGroup:
 		chs, coll := charset.GetDefaultCharsetAndCollate()
 		return &ExprCollation{CoercibilitySysconst, UNICODE, chs, coll}, nil
 	case ast.Format, ast.Space, ast.ToBase64, ast.UUID, ast.Hex, ast.MD5, ast.SHA, ast.SHA2, ast.SM3:

expression/collation_test.go

@@ -673,7 +673,7 @@ func TestDeriveCollation(t *testing.T) {
 		},
 		{
 			[]string{
-				ast.Database, ast.User, ast.CurrentUser, ast.Version, ast.CurrentRole, ast.TiDBVersion,
+				ast.Database, ast.User, ast.CurrentUser, ast.Version, ast.CurrentRole, ast.TiDBVersion, ast.CurrentResourceGroup,
 			},
 			[]Expression{},
 			[]types.EvalType{},

expression/function_traits.go

@@ -33,15 +33,16 @@ var NonPreparedPlanCacheableOp = map[string]struct{}{
 
 // UnCacheableFunctions stores functions which can not be cached to plan cache.
 var UnCacheableFunctions = map[string]struct{}{
-	ast.Database:     {},
-	ast.CurrentUser:  {},
-	ast.CurrentRole:  {},
-	ast.User:         {},
-	ast.ConnectionID: {},
-	ast.LastInsertId: {},
-	ast.RowCount:     {},
-	ast.Version:      {},
-	ast.Like:         {},
+	ast.Database:             {},
+	ast.CurrentUser:          {},
+	ast.CurrentRole:          {},
+	ast.CurrentResourceGroup: {},
+	ast.User:                 {},
+	ast.ConnectionID:         {},
+	ast.LastInsertId:         {},
+	ast.RowCount:             {},
+	ast.Version:              {},
+	ast.Like:                 {},
 }
 
 // unFoldableFunctions stores functions which can not be folded duration constant folding stage.

expression/typeinfer_test.go

@@ -881,6 +881,7 @@ func (s *InferTypeSuite) createTestCase4InfoFunc() []typeInferTestCase {
 		{"user()", mysql.TypeVarString, charset.CharsetUTF8MB4, 0, 64, types.UnspecifiedLength},
 		{"connection_id()", mysql.TypeLonglong, charset.CharsetBin, mysql.BinaryFlag | mysql.UnsignedFlag | mysql.NotNullFlag, mysql.MaxIntWidth, 0},
 		{"version()", mysql.TypeVarString, charset.CharsetUTF8MB4, mysql.NotNullFlag, 64, types.UnspecifiedLength},
+		{"current_resource_group()", mysql.TypeVarString, charset.CharsetUTF8MB4, mysql.NotNullFlag, 64, types.UnspecifiedLength},
 	}
 }
 

parser/ast/functions.go

@@ -266,6 +266,7 @@ const (
 	TiDBDecodeSQLDigests = "tidb_decode_sql_digests"
 	FormatBytes          = "format_bytes"
 	FormatNanoTime       = "format_nano_time"
+	CurrentResourceGroup = "current_resource_group"
 
 	// control functions
 	If     = "if"

parser/ast/misc.go

@@ -58,6 +58,7 @@ var (
 	_ StmtNode = &HelpStmt{}
 	_ StmtNode = &PlanReplayerStmt{}
 	_ StmtNode = &CompactTableStmt{}
+	_ StmtNode = &SetResourceGroupStmt{}
 
 	_ Node = &PrivElem{}
 	_ Node = &VariableAssignment{}
@@ -3862,3 +3863,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)
+}