Allow for custom cockroach binary versions

Author: pawalt

testserver/binaries.go

@@ -62,7 +62,9 @@ var muslRE = regexp.MustCompile(`(?i)\bmusl\b`)
 // GetDownloadResponse return the http response of a CRDB download.
 // It creates the url for downloading a CRDB binary for current runtime OS,
 // makes a request to this url, and return the response.
-func GetDownloadResponse(nonStable bool) (*http.Response, string, error) {
+// nonStable should only be used if desiredVersion is not specified. If nonStable
+// is used, the latest cockroach binary will be used.
+func GetDownloadResponse(desiredVersion string, nonStable bool) (*http.Response, string, error) {
 	goos := runtime.GOOS
 	if goos == "linux" {
 		goos += func() string {
@@ -86,9 +88,10 @@ func GetDownloadResponse(nonStable bool) (*http.Response, string, error) {
 	var dbUrl string
 	var err error
 
-	latestStableVersion := ""
-	// For the latest (beta) CRDB, we use the `edge-binaries.cockroachdb.com` host.
-	if nonStable {
+	if desiredVersion != "" {
+		dbUrl = getDownloadUrlForVersion(desiredVersion)
+	} else if nonStable {
+		// For the latest (beta) CRDB, we use the `edge-binaries.cockroachdb.com` host.
 		u := &url.URL{
 			Scheme: "https",
 			Host:   "edge-binaries.cockroachdb.com",
@@ -97,7 +100,7 @@ func GetDownloadResponse(nonStable bool) (*http.Response, string, error) {
 		dbUrl = u.String()
 	} else {
 		// For the latest stable CRDB, we use the url provided in the CRDB release page.
-		dbUrl, latestStableVersion, err = getLatestStableVersionInfo()
+		dbUrl, desiredVersion, err = getLatestStableVersionInfo()
 		if err != nil {
 			return nil, "", err
 		}
@@ -113,22 +116,24 @@ func GetDownloadResponse(nonStable bool) (*http.Response, string, error) {
 		return nil, "", fmt.Errorf("error downloading %s: %d (%s)", dbUrl,
 			response.StatusCode, response.Status)
 	}
-	return response, latestStableVersion, nil
+	return response, desiredVersion, nil
 }
 
 // downloadBinary saves the latest version of CRDB into a local binary file,
 // and returns the path for this local binary.
+// To download a specific cockroach version, specify desiredVersion. Otherwise,
+// the latest stable or non-stable version will be chosen.
 // To download the latest STABLE version of CRDB, set `nonStable` to false.
 // To download the bleeding edge version of CRDB, set `nonStable` to true.
-func downloadBinary(tc *TestConfig, nonStable bool) (string, error) {
-	response, latestStableVersion, err := GetDownloadResponse(nonStable)
+func downloadBinary(tc *TestConfig, desiredVersion string, nonStable bool) (string, error) {
+	response, desiredVersion, err := GetDownloadResponse(desiredVersion, nonStable)
 	if err != nil {
 		return "", err
 	}
 
 	defer func() { _ = response.Body.Close() }()
 
-	filename, err := GetDownloadFilename(response, nonStable, latestStableVersion)
+	filename, err := GetDownloadFilename(response, nonStable, desiredVersion)
 	if err != nil {
 		return "", err
 	}
@@ -226,7 +231,7 @@ func downloadBinary(tc *TestConfig, nonStable bool) (string, error) {
 
 // GetDownloadFilename returns the local filename of the downloaded CRDB binary file.
 func GetDownloadFilename(
-	response *http.Response, nonStableDB bool, latestStableVersion string,
+	response *http.Response, nonStableDB bool, desiredVersion string,
 ) (string, error) {
 	if nonStableDB {
 		const contentDisposition = "Content-Disposition"
@@ -240,7 +245,7 @@ func GetDownloadFilename(
 		}
 		return filename, nil
 	}
-	filename := fmt.Sprintf("cockroach-%s", latestStableVersion)
+	filename := fmt.Sprintf("cockroach-%s", desiredVersion)
 	if runtime.GOOS == "windows" {
 		filename += ".exe"
 	}
@@ -264,17 +269,24 @@ func getLatestStableVersionInfo() (string, string, error) {
 	if !ok {
 		return "", "", fmt.Errorf("api/updates response is of wrong format")
 	}
-	var downloadUrl string
+
+	downloadUrl := getDownloadUrlForVersion(latestStableVersion)
+
+	latestStableVerFormatted := strings.ReplaceAll(latestStableVersion, ".", "-")
+	return downloadUrl, latestStableVerFormatted, nil
+}
+
+func getDownloadUrlForVersion(version string) string {
 	switch runtime.GOOS {
 	case "linux":
-		downloadUrl = fmt.Sprintf(linuxUrlpat, latestStableVersion)
+		return fmt.Sprintf(linuxUrlpat, version)
 	case "darwin":
-		downloadUrl = fmt.Sprintf(macUrlpat, latestStableVersion)
+		return fmt.Sprintf(macUrlpat, version)
 	case "windows":
-		downloadUrl = fmt.Sprintf(winUrlpat, latestStableVersion)
+		return fmt.Sprintf(winUrlpat, version)
 	}
-	latestStableVerFormatted := strings.ReplaceAll(latestStableVersion, ".", "-")
-	return downloadUrl, latestStableVerFormatted, nil
+
+	panic(errors.New("could not get supported go os version"))
 }
 
 // downloadBinaryFromResponse copies the http response's body directly into a local binary.

testserver/testserver.go

@@ -220,6 +220,7 @@ type testServerArgs struct {
 	listenAddrPorts        []int
 	testConfig             TestConfig
 	nonStableDB            bool
+	customVersion          string // custom cockroach version to use
 	cockroachBinary        string // path to cockroach executable file
 	upgradeCockroachBinary string // path to cockroach binary for upgrade
 	numNodes               int
@@ -288,6 +289,14 @@ func NonStableDbOpt() TestServerOpt {
 	}
 }
 
+// CustomVersionOpt is a TestServer option that can be passed to NewTestServer to
+// download the a specific version of CRDB.
+func CustomVersionOpt(version string) TestServerOpt {
+	return func(args *testServerArgs) {
+		args.customVersion = version
+	}
+}
+
 // ExposeConsoleOpt is a TestServer option that can be passed to NewTestServer to
 // expose the console of the server on the given port.
 // Warning: This is kept around for backwards compatibility, use AddHttpPortOpt
@@ -390,7 +399,7 @@ func NewTestServer(opts ...TestServerOpt) (TestServer, error) {
 			serverArgs.cockroachBinary = cockroachBinary
 		}
 	} else {
-		serverArgs.cockroachBinary, err = downloadBinary(&serverArgs.testConfig, serverArgs.nonStableDB)
+		serverArgs.cockroachBinary, err = downloadBinary(&serverArgs.testConfig, serverArgs.customVersion, serverArgs.nonStableDB)
 		if err != nil {
 			if errors.Is(err, errStoppedInMiddle) {
 				// If the testserver is intentionally killed in the middle of downloading,

testserver/testserver_test.go

@@ -120,6 +120,12 @@ func TestRunServer(t *testing.T) {
 			name:          "InsecureNonStable",
 			instantiation: func(t *testing.T) (*sql.DB, func()) { return testserver.NewDBForTest(t, testserver.NonStableDbOpt()) },
 		},
+		{
+			name: "InsecureCustomVersion",
+			instantiation: func(t *testing.T) (*sql.DB, func()) {
+				return testserver.NewDBForTest(t, testserver.CustomVersionOpt("21.2.15"))
+			},
+		},
 		{
 			name: "InsecureWithCustomizedMemSizeNonStable",
 			instantiation: func(t *testing.T) (*sql.DB, func()) {
@@ -646,7 +652,7 @@ func testFlockWithDownloadKilled(t *testing.T) (*sql.DB, func()) {
 
 // getLocalFile returns the to-be-downloaded CRDB binary's local path.
 func getLocalFile(nonStable bool) (string, error) {
-	response, latestStableVersion, err := testserver.GetDownloadResponse(nonStable)
+	response, latestStableVersion, err := testserver.GetDownloadResponse("", nonStable)
 	if err != nil {
 		return "", err
 	}