Merge cockroachdb#96388 cockroachdb#96452 cockroachdb#96469 cockroach…

…db#96500

96388: logictest: Skip if using cockroach-go/testserver r=renatolabs a=itsbilal

Skip of flaky test.

Informs: cockroachdb#96387

Epic: none

Release note: None

96452: roachtest: add aws weekly script r=rickystewart a=healthy-pod

This change adds the scripts to be used by the aws weekly roachtest teamcity job.

Release note: None
Epic: none

96469: logictest: refactor client connection creation logic r=rafiss a=andyyang890

This patch abstracts the logic for creating client connections and
persisting them into a separate function. As a side benefit, it also
fixes the problem of clients created in `execQuery` not being properly
tracked in the `clients` field of `logicTest`.

Informs cockroachdb#92342

Release note: None

96500: authors: add BabuSrithar to authors r=healthy-pod a=BabuSrithar

Release note: None
Epic: None

Co-authored-by: Bilal Akhtar <[email protected]>
Co-authored-by: healthy-pod <[email protected]>
Co-authored-by: Andy Yang <[email protected]>
Co-authored-by: babusrithar <[email protected]>

AUTHORS

@@ -88,6 +88,7 @@ Artem Ervits <[email protected]> <[email protected]>
 Arul Ajmani <[email protected]> <[email protected]>
 Asit Mahato <[email protected]>
 Austen McClernon <[email protected]> <[email protected]>
+BabuSrithar <[email protected]> <[email protected]>
 Barry He <[email protected]> <[email protected]>
 bc <[email protected]>
 Ben Bardin <[email protected]>

build/teamcity/cockroach/nightlies/roachtest_weekly_aws.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+set -exuo pipefail
+
+dir="$(dirname $(dirname $(dirname $(dirname "${0}"))))"
+
+source "$dir/teamcity-support.sh"  # For $root
+source "$dir/teamcity-bazel-support.sh"  # For run_bazel
+
+BAZEL_SUPPORT_EXTRA_DOCKER_ARGS="-e LITERAL_ARTIFACTS_DIR=$root/artifacts -e AWS_ACCESS_KEY_ID -e AWS_ACCESS_KEY_ID_ASSUME_ROLE -e AWS_KMS_KEY_ARN_A -e AWS_KMS_KEY_ARN_B -e AWS_KMS_REGION_A -e AWS_KMS_REGION_B -e AWS_ROLE_ARN -e AWS_SECRET_ACCESS_KEY -e AWS_SECRET_ACCESS_KEY_ASSUME_ROLE -e BUILD_TAG -e BUILD_VCS_NUMBER -e CLOUD -e COCKROACH_DEV_LICENSE -e TESTS -e COUNT -e GITHUB_API_TOKEN -e GITHUB_ORG -e GITHUB_REPO -e GOOGLE_EPHEMERAL_CREDENTIALS -e SLACK_TOKEN -e TC_BUILDTYPE_ID -e TC_BUILD_BRANCH -e TC_BUILD_ID -e TC_SERVER_URL" \
+			       run_bazel build/teamcity/cockroach/nightlies/roachtest_weekly_aws_impl.sh

build/teamcity/cockroach/nightlies/roachtest_weekly_aws_impl.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+set -exuo pipefail
+
+dir="$(dirname $(dirname $(dirname $(dirname "${0}"))))"
+
+source "$dir/teamcity-support.sh"
+
+if [[ ! -f ~/.ssh/id_rsa.pub ]]; then
+  ssh-keygen -q -C "roachtest-weekly-bazel $(date)" -N "" -f ~/.ssh/id_rsa
+fi
+
+source $root/build/teamcity/cockroach/nightlies/roachtest_compile_bits.sh
+
+artifacts=/artifacts
+source $root/build/teamcity/util/roachtest_util.sh
+
+build/teamcity-roachtest-invoke.sh \
+  tag:aws-weekly \
+  --cloud="${CLOUD}" \
+  --build-tag="${BUILD_TAG}" \
+  --cluster-id "${TC_BUILD_ID}" \
+  --cockroach "$PWD/bin/cockroach" \
+  --artifacts=/artifacts \
+  --artifacts-literal="${LITERAL_ARTIFACTS_DIR:-}" \
+  --slack-token="${SLACK_TOKEN}"

pkg/sql/logictest/logic.go

@@ -1161,35 +1161,35 @@ func (t *logicTest) outf(format string, args ...interface{}) {
 
 // setUser sets the DB client to the specified user and connects
 // to the node in the cluster at index nodeIdx.
-// It returns a cleanup function to be run when the credentials
-// are no longer needed.
-func (t *logicTest) setUser(user string, nodeIdx int) func() {
-	if db, ok := t.clients[user][nodeIdx]; ok {
-		t.db = db
-		t.user = user
+func (t *logicTest) setUser(user string, nodeIdx int) {
+	db := t.getOrOpenClient(user, nodeIdx)
+	t.db = db
+	t.user = user
+	t.nodeIdx = nodeIdx
+}
 
-		// No cleanup necessary, but return a no-op func to avoid nil pointer dereference.
-		return func() {}
+// getOrOpenClient returns the existing client for the given user and nodeIdx,
+// if one exists. Otherwise, it opens and returns a new client.
+func (t *logicTest) getOrOpenClient(user string, nodeIdx int) *gosql.DB {
+	if db, ok := t.clients[user][nodeIdx]; ok {
+		return db
 	}
 
-	var addr string
 	var pgURL url.URL
-	var pgUser string
-	var cleanupFunc func()
-	pgUser = strings.TrimPrefix(user, "host-cluster-")
+	pgUser := strings.TrimPrefix(user, "host-cluster-")
 	if t.cfg.UseCockroachGoTestserver {
 		pgURL = *t.testserverCluster.PGURLForNode(nodeIdx)
 		pgURL.User = url.User(pgUser)
-		pgURL.Path = "test"
-		cleanupFunc = func() {}
 	} else {
-		addr = t.cluster.Server(nodeIdx).ServingSQLAddr()
+		addr := t.cluster.Server(nodeIdx).ServingSQLAddr()
 		if len(t.tenantAddrs) > 0 && !strings.HasPrefix(user, "host-cluster-") {
 			addr = t.tenantAddrs[nodeIdx]
 		}
+		var cleanupFunc func()
 		pgURL, cleanupFunc = sqlutils.PGUrl(t.rootT, addr, "TestLogic", url.User(pgUser))
-		pgURL.Path = "test"
+		t.clusterCleanupFuncs = append(t.clusterCleanupFuncs, cleanupFunc)
 	}
+	pgURL.Path = "test"
 
 	db := t.openDB(pgURL)
 
@@ -1213,11 +1213,8 @@ func (t *logicTest) setUser(user string, nodeIdx int) func() {
 		t.clients[user] = make(map[int]*gosql.DB)
 	}
 	t.clients[user][nodeIdx] = db
-	t.db = db
-	t.user = pgUser
-	t.nodeIdx = nodeIdx
 
-	return cleanupFunc
+	return db
 }
 
 func (t *logicTest) openDB(pgURL url.URL) *gosql.DB {
@@ -1306,8 +1303,9 @@ func (t *logicTest) newTestServerCluster(bootstrapBinaryPath string, upgradeBina
 	}
 
 	t.testserverCluster = ts
-	t.clusterCleanupFuncs = append(t.clusterCleanupFuncs, t.setUser(username.RootUser, 0 /* nodeIdx */))
 	t.clusterCleanupFuncs = append(t.clusterCleanupFuncs, ts.Stop)
+
+	t.setUser(username.RootUser, 0 /* nodeIdx */)
 }
 
 // newCluster creates a new cluster. It should be called after the logic tests's
@@ -1721,9 +1719,7 @@ func (t *logicTest) newCluster(
 		)
 	}
 
-	// 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(username.RootUser, 0 /* nodeIdx */))
+	t.setUser(username.RootUser, 0 /* nodeIdx */)
 }
 
 // waitForTenantReadOnlyClusterSettingToTakeEffectOrFatal waits until all tenant
@@ -1823,6 +1819,7 @@ func (t *logicTest) setup(
 	t.testCleanupFuncs = append(t.testCleanupFuncs, tempExternalIODirCleanup)
 
 	if cfg.UseCockroachGoTestserver {
+		skip.WithIssue(t.t(), 96387, "flakes with panic over connection string")
 		skip.UnderStress(t.t(), "test takes a long time and downloads release artifacts")
 		if runtime.GOARCH == "arm64" && strings.HasPrefix(cfg.CockroachGoBootstrapVersion, "v22.1") {
 			skip.IgnoreLint(t.t(), "Skip under ARM64. There are no ARM release artifacts for v22.1.")
@@ -2960,15 +2957,14 @@ func (t *logicTest) processSubtest(
 					nodeIdx = int(idx)
 				}
 			}
-			cleanupUserFunc := t.setUser(fields[1], nodeIdx)
+			t.setUser(fields[1], nodeIdx)
 			// In multi-tenant tests, we may need to also create database test when
 			// we switch to a different tenant.
 			if t.cfg.UseTenant && strings.HasPrefix(fields[1], "host-cluster-") {
 				if _, err := t.db.Exec("CREATE DATABASE IF NOT EXISTS test; USE test;"); err != nil {
 					return errors.Wrapf(err, "error creating database on admin tenant")
 				}
 			}
-			defer cleanupUserFunc()
 
 		case "skip":
 			reason := "skipped"
@@ -3311,30 +3307,8 @@ func (t *logicTest) execQuery(query logicQuery) error {
 	t.noticeBuffer = nil
 
 	db := t.db
-	var closeDB func()
 	if query.nodeIdx != t.nodeIdx {
-		var pgURL url.URL
-		if t.testserverCluster != nil {
-			pgURL = *t.testserverCluster.PGURLForNode(query.nodeIdx)
-			pgURL.User = url.User(t.user)
-			pgURL.Path = "test"
-		} else {
-			addr := t.cluster.Server(query.nodeIdx).ServingSQLAddr()
-			if len(t.tenantAddrs) > 0 {
-				addr = t.tenantAddrs[query.nodeIdx]
-			}
-			var cleanupFunc func()
-			pgURL, cleanupFunc = sqlutils.PGUrl(t.rootT, addr, "TestLogic", url.User(t.user))
-			defer cleanupFunc()
-			pgURL.Path = "test"
-		}
-
-		db = t.openDB(pgURL)
-		closeDB = func() {
-			if err := db.Close(); err != nil {
-				t.Fatal(err)
-			}
-		}
+		db = t.getOrOpenClient(t.user, query.nodeIdx)
 	}
 
 	if query.expectAsync {
@@ -3354,18 +3328,13 @@ func (t *logicTest) execQuery(query logicQuery) error {
 
 		startedChan := make(chan struct{})
 		go func() {
-			if closeDB != nil {
-				defer closeDB()
-			}
 			startedChan <- struct{}{}
 			rows, err := db.Query(query.sql)
 			pending.resultChan <- pendingQueryResult{rows, err}
 		}()
 
 		<-startedChan
 		return nil
-	} else if closeDB != nil {
-		defer closeDB()
 	}
 
 	rows, err := db.Query(query.sql)