logictest: add support for per-user nodeidx overrides

While we have support for switching users, which each hold a distinct db
client handle, there was previously no way to set a specific node for
the SQL client to connect to other than the default used for the entire
test.  With this change, we are now able to override the node used for
the SQL client connection for a particular user with the following
statement in a test:
```
user testuser nodeidx=3
```
With this directive, subsequent statements will be made via a SQL
connection to `NodeID` 4 (as `NodeID`s are 1-indexed), using user
`testuser`.  When a subsequent switch to the user is made that does not
specify a nodeidx override, such as `user root`, subsequent statements
will revert to using the default node for the test (typically `NodeID`
1).

This enables support for testing Contention Events, where the node
initiating the contending transaction needs to be validated, among other
use cases.

Note that the connections are per-user, so opening multiple connections
to different nodes for a single user is not currently supported.

Release note: None

pkg/sql/logictest/logic.go

@@ -282,8 +282,11 @@ import (
 //  - sleep <duration>
 //    Introduces a sleep period. Example: sleep 2s
 //
-//  - user <username>
+//  - user <username> [nodeidx=N]
 //    Changes the user for subsequent statements or queries.
+//    If nodeidx is specified, this user will connect to the node
+//    in the cluster with index N (note this is 0-indexed, while
+//    node IDs themselves are 1-indexed).
 //
 //  - skipif <mysql/mssql/postgresql/cockroachdb/config CONFIG [ISSUE]>
 //    Skips the following `statement` or `query` if the argument is
@@ -1329,7 +1332,7 @@ func (t *logicTest) outf(format string, args ...interface{}) {
 // setUser sets the DB client to the specified user.
 // It returns a cleanup function to be run when the credentials
 // are no longer needed.
-func (t *logicTest) setUser(user string) func() {
+func (t *logicTest) setUser(user string, nodeIdxOverride int) func() {
 	if t.clients == nil {
 		t.clients = map[string]*gosql.DB{}
 	}
@@ -1341,9 +1344,14 @@ func (t *logicTest) setUser(user string) func() {
 		return func() {}
 	}
 
-	addr := t.cluster.Server(t.nodeIdx).ServingSQLAddr()
+	nodeIdx := t.nodeIdx
+	if nodeIdxOverride > 0 {
+		nodeIdx = nodeIdxOverride
+	}
+
+	addr := t.cluster.Server(nodeIdx).ServingSQLAddr()
 	if len(t.tenantAddrs) > 0 {
-		addr = t.tenantAddrs[t.nodeIdx]
+		addr = t.tenantAddrs[nodeIdx]
 	}
 	pgURL, cleanupFunc := sqlutils.PGUrl(t.rootT, addr, "TestLogic", url.User(user))
 	pgURL.Path = "test"
@@ -1671,7 +1679,7 @@ func (t *logicTest) newCluster(serverArgs TestServerArgs, opts []clusterOpt) {
 
 	// db may change over the lifetime of this function, with intermediate
 	// values cached in t.clients and finally closed in t.close().
-	t.clusterCleanupFuncs = append(t.clusterCleanupFuncs, t.setUser(security.RootUser))
+	t.clusterCleanupFuncs = append(t.clusterCleanupFuncs, t.setUser(security.RootUser, 0 /* nodeIdxOverride */))
 }
 
 // shutdownCluster performs the necessary cleanup to shutdown the current test
@@ -2474,13 +2482,23 @@ func (t *logicTest) processSubtest(subtest subtestDetails, path string) error {
 		case "halt", "hash-threshold":
 
 		case "user":
+			var nodeIdx int
 			if len(fields) < 2 {
 				return errors.Errorf("user command requires one argument, found: %v", fields)
 			}
 			if len(fields[1]) == 0 {
 				return errors.Errorf("user command requires a non-blank argument")
 			}
-			cleanupUserFunc := t.setUser(fields[1])
+			if len(fields) >= 3 {
+				if strings.HasPrefix(fields[2], "nodeidx=") {
+					idx, err := strconv.ParseInt(strings.SplitN(fields[2], "=", 2)[1], 10, 64)
+					if err != nil {
+						return errors.Wrapf(err, "error parsing nodeidx")
+					}
+					nodeIdx = int(idx)
+				}
+			}
+			cleanupUserFunc := t.setUser(fields[1], nodeIdx)
 			defer cleanupUserFunc()
 
 		case "skip":
@@ -3191,7 +3209,7 @@ func (t *logicTest) validateAfterTestCompletion() error {
 			t.Fatalf("failed to close connection for user %s: %v", username, err)
 		}
 	}
-	t.setUser("root")
+	t.setUser("root", 0 /* nodeIdxOverride */)
 
 	// Some cleanup to make sure the following validation queries can run
 	// successfully. First we rollback in case the logic test had an uncommitted

pkg/sql/logictest/testdata/logic_test/contention_event

@@ -16,7 +16,7 @@ query TT
 SELECT * FROM kv
 ----
 
-user testuser
+user testuser nodeidx=3
 
 statement ok
 BEGIN
@@ -38,7 +38,7 @@ BEGIN;
 SET TRANSACTION PRIORITY HIGH;
 SELECT * FROM kv ORDER BY k ASC
 
-user testuser
+user testuser nodeidx=3
 
 statement ok
 ROLLBACK
@@ -86,7 +86,7 @@ WITH payloads AS (
   FROM payloads
   WHERE payload_type = 'roachpb.ContentionEvent'
   AND crdb_internal.pretty_key(decode(payload_jsonb->>'key', 'base64'), 1) LIKE '/1/"k"/%'
-  AND (payload_jsonb->'txnMeta'->>'coordinatorNodeId')::INTEGER = 1
+  AND (payload_jsonb->'txnMeta'->>'coordinatorNodeId')::INTEGER = 4
 ----
 true