session: add check resource group hint existence (#45296)

close #45178

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;")
 }

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)