sessionctx: support encoding resource group name into session states (#…

…42886) close #42884
pingcap · Apr 14, 2023 · 477fbaf · 477fbaf
1 parent d01a6fd
commit 477fbaf
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/sessionctx/sessionstates/session_states.go b/sessionctx/sessionstates/session_states.go
@@ -81,5 +81,6 @@ type SessionStates struct {
 	LastInsertID         uint64                       `json:"last-insert-id,omitempty"`
 	Warnings             []stmtctx.SQLWarn            `json:"warnings,omitempty"`
 	// Define it as string to avoid cycle import.
-	Bindings string `json:"bindings,omitempty"`
+	Bindings          string `json:"bindings,omitempty"`
+	ResourceGroupName string `json:"rs-group,omitempty"`
 }
diff --git a/sessionctx/sessionstates/session_states_test.go b/sessionctx/sessionstates/session_states_test.go
@@ -422,6 +422,19 @@ func TestSessionCtx(t *testing.T) {
 				tk.MustQuery("select @@last_plan_from_binding").Check(testkit.Rows("1"))
 			},
 		},
+		{
+			// check ResourceGroupName
+			setFunc: func(tk *testkit.TestKit) any {
+				tk.MustExec("SET GLOBAL tidb_enable_resource_control='on'")
+				tk.MustExec("CREATE RESOURCE GROUP rg1 ru_per_sec = 100")
+				tk.MustExec("SET RESOURCE GROUP `rg1`")
+				require.Equal(t, "rg1", tk.Session().GetSessionVars().ResourceGroupName)
+				return nil
+			},
+			checkFunc: func(tk *testkit.TestKit, param any) {
+				tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows("rg1"))
+			},
+		},
 	}
 
 	for _, tt := range tests {

diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go
@@ -2474,6 +2474,7 @@ func (s *SessionVars) EncodeSessionStates(_ context.Context, sessionStates *sess
 	sessionStates.SequenceLatestValues = s.SequenceState.GetAllStates()
 	sessionStates.FoundInPlanCache = s.PrevFoundInPlanCache
 	sessionStates.FoundInBinding = s.PrevFoundInBinding
+	sessionStates.ResourceGroupName = s.ResourceGroupName
 
 	// Encode StatementContext. We encode it here to avoid circle dependency.
 	sessionStates.LastAffectedRows = s.StmtCtx.PrevAffectedRows
@@ -2507,6 +2508,7 @@ func (s *SessionVars) DecodeSessionStates(_ context.Context, sessionStates *sess
 	s.SequenceState.SetAllStates(sessionStates.SequenceLatestValues)
 	s.FoundInPlanCache = sessionStates.FoundInPlanCache
 	s.FoundInBinding = sessionStates.FoundInBinding
+	s.ResourceGroupName = sessionStates.ResourceGroupName
 
 	// Decode StatementContext.
 	s.StmtCtx.SetAffectedRows(uint64(sessionStates.LastAffectedRows))