From 0bf75ff87a44536e816df08e60cecae997d9b268 Mon Sep 17 00:00:00 2001
From: Michael Erickson <michae2@cockroachlabs.com>
Date: Wed, 25 Sep 2024 10:38:43 -0700
Subject: [PATCH] roachtest: set timeout in MR costfuzz,
 unoptimized-query-oracle

We weren't setting a couple of session variables on the second
connection opened for multi-region costfuzz and
unoptimized-query-oracle, including statement_timeout.

Fixes: #131147
Fixes: #131153
Fixes: #131166
Fixes: #131344

Release note: None
---
 .../roachtest/tests/query_comparison_util.go  | 39 +++++++++++--------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/pkg/cmd/roachtest/tests/query_comparison_util.go b/pkg/cmd/roachtest/tests/query_comparison_util.go
index 3c14050fd531..7d76d3b6c176 100644
--- a/pkg/cmd/roachtest/tests/query_comparison_util.go
+++ b/pkg/cmd/roachtest/tests/query_comparison_util.go
@@ -184,6 +184,25 @@ func runOneRoundQueryComparison(
 		fmt.Fprint(failureLog, "\n")
 	}
 
+	connSetup := func(conn *gosql.DB) {
+		setStmtTimeout := fmt.Sprintf("SET statement_timeout='%s';", statementTimeout.String())
+		t.Status("setting statement_timeout")
+		t.L().Printf("statement timeout:\n%s", setStmtTimeout)
+		if _, err := conn.Exec(setStmtTimeout); err != nil {
+			t.Fatal(err)
+		}
+		logStmt(setStmtTimeout)
+
+		setUnconstrainedStmt := "SET unconstrained_non_covering_index_scan_enabled = true;"
+		t.Status("setting unconstrained_non_covering_index_scan_enabled")
+		t.L().Printf("\n%s", setUnconstrainedStmt)
+		if _, err := conn.Exec(setUnconstrainedStmt); err != nil {
+			logStmt(setUnconstrainedStmt)
+			t.Fatal(err)
+		}
+		logStmt(setUnconstrainedStmt)
+	}
+
 	node := 1
 	conn := c.Conn(ctx, t.L(), node)
 
@@ -191,6 +210,8 @@ func runOneRoundQueryComparison(
 	t.L().Printf("seed: %d", seed)
 	t.L().Printf("setupName: %s", qct.setupName)
 
+	connSetup(conn)
+
 	if qct.setupName == "workload-replay" {
 
 		logTest := func(logStr string, logType string) {
@@ -309,23 +330,6 @@ func runOneRoundQueryComparison(
 			}
 		}
 
-		setStmtTimeout := fmt.Sprintf("SET statement_timeout='%s';", statementTimeout.String())
-		t.Status("setting statement_timeout")
-		t.L().Printf("statement timeout:\n%s", setStmtTimeout)
-		if _, err := conn.Exec(setStmtTimeout); err != nil {
-			t.Fatal(err)
-		}
-		logStmt(setStmtTimeout)
-
-		setUnconstrainedStmt := "SET unconstrained_non_covering_index_scan_enabled = true;"
-		t.Status("setting unconstrained_non_covering_index_scan_enabled")
-		t.L().Printf("\n%s", setUnconstrainedStmt)
-		if _, err := conn.Exec(setUnconstrainedStmt); err != nil {
-			logStmt(setUnconstrainedStmt)
-			t.Fatal(err)
-		}
-		logStmt(setUnconstrainedStmt)
-
 		conn2 := conn
 		node2 := 1
 		if qct.isMultiRegion {
@@ -336,6 +340,7 @@ func runOneRoundQueryComparison(
 			node2 = rnd.Intn(qct.nodeCount-1) + 2
 			t.Status(fmt.Sprintf("running some queries from node %d with conn1 and some queries from node %d with conn2", node, node2))
 			conn2 = c.Conn(ctx, t.L(), node2)
+			connSetup(conn2)
 		}
 
 		// Initialize a smither that generates only INSERT and UPDATE statements with