diff --git a/session/resourcegrouptest/resource_group_test.go b/session/resourcegrouptest/resource_group_test.go index 208775f822a7f..cd31663228e45 100644 --- a/session/resourcegrouptest/resource_group_test.go +++ b/session/resourcegrouptest/resource_group_test.go @@ -45,4 +45,15 @@ func TestResourceGroupHintInTxn(t *testing.T) { // for final prewrite/commit the resource group should be rg2 tk.MustExec("update /*+ RESOURCE_GROUP(rg2) */ t set val = val + 1 where id = 3;") tk.MustExec("COMMIT;") + + tk.MustExec("SET @@autocommit=1;") + require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/kv/TxnResouceGroupChecker", `return("default")`)) + tk.MustExec("insert /*+ RESOURCE_GROUP(not_exist_group) */ into t values (4, 4);") + + tk.MustExec("BEGIN;") + // for pessimistic lock the resource group should be rg1 + tk.MustExec("insert /*+ RESOURCE_GROUP(unknown_1) */ into t values (5, 5);") + // for final prewrite/commit the resource group should be rg2 + tk.MustExec("update /*+ RESOURCE_GROUP(unknown_2) */ t set val = val + 1 where id = 5;") + tk.MustExec("COMMIT;") } diff --git a/session/session.go b/session/session.go index a368797a4a6d7..1c307d16c9fc9 100644 --- a/session/session.go +++ b/session/session.go @@ -2226,10 +2226,20 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex // session resource-group might be changed by query hint, ensure restore it back when // the execution finished. if sessVars.ResourceGroupName != originalResourceGroup { - defer func() { - // Restore the resource group for the session + // if target resource group doesn't exist, fallback to the origin resource group. + if _, ok := domain.GetDomain(s).InfoSchema().ResourceGroupByName(model.NewCIStr(sessVars.ResourceGroupName)); !ok { + logutil.Logger(ctx).Warn("Unknown resource group from hint", zap.String("name", sessVars.ResourceGroupName)) sessVars.ResourceGroupName = originalResourceGroup - }() + // if we are in a txn, should also reset the txn resource group. + if txn, err := s.Txn(false); err == nil && txn != nil && txn.Valid() { + kv.SetTxnResourceGroup(txn, originalResourceGroup) + } + } else { + defer func() { + // Restore the resource group for the session + sessVars.ResourceGroupName = originalResourceGroup + }() + } } if err != nil { s.rollbackOnError(ctx)