Merge cockroachdb#107626

107626: logictest: reset role before doing post-test checks r=rafiss a=rafiss

This makes sure we run the checks with the root user.

fixes cockroachdb#107561
Release note: None

Co-authored-by: Rafi Shamim <[email protected]>

Author: craig[bot]

pkg/sql/logictest/logic.go

@@ -1001,7 +1001,7 @@ type logicTest struct {
 	clients map[string]map[int]*gosql.DB
 	// client currently in use. This can change during processing
 	// of a test input file when encountering the "user" directive.
-	// see setUser() for details.
+	// see setSessionUser() for details.
 	user string
 	db   *gosql.DB
 	// clusterCleanupFuncs contains the cleanup methods that are specific to a
@@ -1169,9 +1169,9 @@ func (t *logicTest) outf(format string, args ...interface{}) {
 	fmt.Printf("[%s] %s\n", now, msg)
 }
 
-// setUser sets the DB client to the specified user and connects
+// setSessionUser sets the DB client to the specified user and connects
 // to the node in the cluster at index nodeIdx.
-func (t *logicTest) setUser(user string, nodeIdx int) {
+func (t *logicTest) setSessionUser(user string, nodeIdx int) {
 	db := t.getOrOpenClient(user, nodeIdx)
 	t.db = db
 	t.user = user
@@ -1309,7 +1309,7 @@ func (t *logicTest) newTestServerCluster(bootstrapBinaryPath, upgradeBinaryPath
 	t.testserverCluster = ts
 	t.clusterCleanupFuncs = append(t.clusterCleanupFuncs, ts.Stop, cleanupLogsDir)
 
-	t.setUser(username.RootUser, 0 /* nodeIdx */)
+	t.setSessionUser(username.RootUser, 0 /* nodeIdx */)
 }
 
 // newCluster creates a new cluster. It should be called after the logic tests's
@@ -1753,7 +1753,7 @@ func (t *logicTest) newCluster(
 		)
 	}
 
-	t.setUser(username.RootUser, 0 /* nodeIdx */)
+	t.setSessionUser(username.RootUser, 0 /* nodeIdx */)
 }
 
 // waitForTenantReadOnlyClusterSettingToTakeEffectOrFatal waits until all tenant
@@ -2245,7 +2245,7 @@ func (t *logicTest) maybeBackupRestore(
 	oldUser := t.user
 	oldNodeIdx := t.nodeIdx
 	defer func() {
-		t.setUser(oldUser, oldNodeIdx)
+		t.setSessionUser(oldUser, oldNodeIdx)
 	}()
 
 	log.Info(context.Background(), "Running cluster backup and restore")
@@ -2265,7 +2265,7 @@ func (t *logicTest) maybeBackupRestore(
 		userToSessionVars[user] = make(map[int]map[string]string)
 		for nodeIdx := range userClients {
 			users[user] = append(users[user], nodeIdx)
-			t.setUser(user, nodeIdx)
+			t.setSessionUser(user, nodeIdx)
 
 			// Serialize session variables.
 			var userSession string
@@ -2307,7 +2307,7 @@ func (t *logicTest) maybeBackupRestore(
 		bucket, strconv.FormatInt(timeutil.Now().UnixNano(), 10))
 
 	// Perform the backup and restore as root.
-	t.setUser(username.RootUser, 0 /* nodeIdx */)
+	t.setSessionUser(username.RootUser, 0 /* nodeIdx */)
 
 	if _, err := t.db.Exec(fmt.Sprintf("BACKUP INTO '%s'", backupLocation)); err != nil {
 		return errors.Wrap(err, "backing up cluster")
@@ -2318,7 +2318,7 @@ func (t *logicTest) maybeBackupRestore(
 	t.resetCluster()
 
 	// Run the restore as root.
-	t.setUser(username.RootUser, 0 /* nodeIdx */)
+	t.setSessionUser(username.RootUser, 0 /* nodeIdx */)
 	if _, err := t.db.Exec(fmt.Sprintf("RESTORE FROM LATEST IN '%s'", backupLocation)); err != nil {
 		return errors.Wrap(err, "restoring cluster")
 	}
@@ -2330,7 +2330,7 @@ func (t *logicTest) maybeBackupRestore(
 	for user, userNodeIdxs := range users {
 		for _, nodeIdx := range userNodeIdxs {
 			// Call setUser for every user to create the connection for that user.
-			t.setUser(user, nodeIdx)
+			t.setSessionUser(user, nodeIdx)
 
 			if userSession, ok := userToHexSession[user][nodeIdx]; ok {
 				if _, err := t.db.Exec(fmt.Sprintf(`SELECT crdb_internal.deserialize_session(decode('%s', 'hex'))`, userSession)); err != nil {
@@ -2977,7 +2977,7 @@ func (t *logicTest) processSubtest(
 					nodeIdx = int(idx)
 				}
 			}
-			t.setUser(fields[1], nodeIdx)
+			t.setSessionUser(fields[1], nodeIdx)
 			// In multi-tenant tests, we may need to also create database test when
 			// we switch to a different tenant.
 			//
@@ -3107,7 +3107,7 @@ func (t *logicTest) processSubtest(
 			// If we upgraded the node we are currently on, we need to open a new
 			// connection since the previous one might now be invalid.
 			if t.nodeIdx == nodeIdx {
-				t.setUser(t.user, nodeIdx)
+				t.setSessionUser(t.user, nodeIdx)
 			}
 		default:
 			return errors.Errorf("%s:%d: unknown command: %s",
@@ -3806,17 +3806,19 @@ func (t *logicTest) validateAfterTestCompletion() error {
 		}
 		delete(t.clients, user)
 	}
-	t.setUser(username.RootUser, 0 /* nodeIdx */)
+	t.setSessionUser(username.RootUser, 0 /* nodeIdx */)
 
 	// Some cleanup to make sure the following validation queries can run
 	// successfully. First we rollback in case the logic test had an uncommitted
 	// txn and second we reset vectorize mode in case it was switched to
 	// `experimental_always`.
 	_, _ = t.db.Exec("ROLLBACK")
-	_, err := t.db.Exec("RESET vectorize")
-	if err != nil {
+	if _, err := t.db.Exec("RESET vectorize"); err != nil {
 		t.Fatal(errors.Wrap(err, "could not reset vectorize mode"))
 	}
+	if _, err := t.db.Exec("RESET ROLE"); err != nil {
+		t.Fatal(errors.Wrap(err, "could not reset role"))
+	}
 
 	validate := func() (string, error) {
 		rows, err := t.db.Query(`SELECT * FROM "".crdb_internal.invalid_objects ORDER BY id`)