fix: set tidb tunning variables only when connected to valid tidb

cmd/go-tpc/tpch.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"os"
 	"runtime"
@@ -26,19 +27,26 @@ var queryTuningVars = []struct {
 	{"tidb_prefer_broadcast_join_by_exchange_data_size", "ON"},
 }
 
-func appendQueryTuningVarsToConnParams() {
-	for _, v := range queryTuningVars {
-		if !strings.Contains(connParams, v.name) {
-			connParams = fmt.Sprintf("%s&%s=%s", connParams, v.name, v.value)
-		}
+// isSysVarSupported determines if a system variable is supported in given TiDB version
+// INFO: Should be optimized if some vars are changed in the future.
+func isSysVarSupported(ver util.SemVersion, sysVar string) bool {
+	if ver.Major > 7 {
+		return true
+	}
+	if ver.Major < 7 {
+		return false
+	}
+	if ver.Minor > 1 {
+		return true
+	}
+	if ver.Minor < 1 {
+		return false
 	}
+
+	return ver.Patch >= 0
 }
 
 func executeTpch(action string) {
-	if action == "run" && driver == "mysql" && tpchConfig.EnableQueryTuning {
-		appendQueryTuningVarsToConnParams()
-	}
-
 	openDB()
 	defer closeDB()
 
@@ -50,6 +58,18 @@ func executeTpch(action string) {
 		runtime.GOMAXPROCS(maxProcs)
 	}
 
+	if action == "run" && tpchConfig.EnableQueryTuning {
+		serverVer, err := getServerVersion(globalDB)
+		if err != nil {
+			panic(fmt.Errorf("get server version failed: %v", err))
+		}
+		if semVer, ok := util.NewTiDBSemVersion(serverVer); ok {
+			if err := setQueryTuningVars(globalDB, semVer); err != nil {
+				panic(fmt.Errorf("set session variables failed: %v", err))
+			}
+		}
+	}
+
 	tpchConfig.PlanReplayerConfig.Host = hosts[0]
 	tpchConfig.PlanReplayerConfig.StatusPort = statusPort
 
@@ -67,6 +87,23 @@ func executeTpch(action string) {
 	w.OutputStats(true)
 }
 
+func getServerVersion(db *sql.DB) (string, error) {
+	var version string
+	err := db.QueryRow("SELECT VERSION()").Scan(&version)
+	return version, err
+}
+
+func setQueryTuningVars(db *sql.DB, ver util.SemVersion) error {
+	for _, v := range queryTuningVars {
+		if isSysVarSupported(ver, v.name) {
+			if _, err := db.Exec(fmt.Sprintf("SET SESSION %s = %s", v.name, v.value)); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 func registerTpch(root *cobra.Command) {
 	cmd := &cobra.Command{
 		Use: "tpch",

pkg/util/version.go

@@ -0,0 +1,52 @@
+package util
+
+import (
+	"strconv"
+	"strings"
+)
+
+type SemVersion struct {
+	Major int
+	Minor int
+	Patch int
+}
+
+// @version is the `SELECT VERSION()` output of TiDB
+func NewTiDBSemVersion(version string) (SemVersion, bool) {
+	isTiDB := strings.Contains(strings.ToLower(version), "tidb")
+	if !isTiDB {
+		return SemVersion{}, false
+	}
+
+	verItems := strings.Split(version, "-v")
+	if len(verItems) < 2 {
+		return SemVersion{}, false
+	}
+	verStr := strings.Split(verItems[1], "-")[0]
+
+	parts := strings.Split(verStr, ".")
+	if len(parts) < 3 {
+		return SemVersion{}, false
+	}
+
+	major, err := strconv.Atoi(parts[0])
+	if err != nil {
+		return SemVersion{}, false
+	}
+	minor, err := strconv.Atoi(parts[1])
+	if err != nil {
+		return SemVersion{}, false
+	}
+
+	patch, err := strconv.Atoi(parts[2])
+	if err != nil {
+		return SemVersion{}, false
+	}
+
+	return SemVersion{
+		Major: major,
+		Minor: minor,
+		Patch: patch,
+	}, true
+
+}

pkg/util/version_test.go

@@ -0,0 +1,75 @@
+package util
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewTiDBSemVersion(t *testing.T) {
+	testCases := []struct {
+		name     string
+		input    string
+		expected SemVersion
+		ok       bool
+	}{
+		{
+			name:     "normal case with addition",
+			input:    "5.7.25-TiDB-v7.1.0-alpha",
+			expected: SemVersion{Major: 7, Minor: 1, Patch: 0},
+			ok:       true,
+		},
+		{
+			name:     "version without addition",
+			input:    "5.7.25-TiDB-v7.4.1",
+			expected: SemVersion{Major: 7, Minor: 4, Patch: 1},
+			ok:       true,
+		},
+		{
+			name:     "multi-part addition",
+			input:    "5.7.25-TiDB-v6.5.3-beta.2",
+			expected: SemVersion{Major: 6, Minor: 5, Patch: 3},
+			ok:       true,
+		},
+		{
+			name:     "empty addition due to trailing hyphen",
+			input:    "5.7.25-TiDB-v7.1.0-",
+			expected: SemVersion{Major: 7, Minor: 1, Patch: 0},
+			ok:       true,
+		},
+		{
+			name:     "non-tidb database",
+			input:    "MySQL 8.0.35",
+			expected: SemVersion{},
+			ok:       false,
+		},
+		{
+			name:     "missing version prefix",
+			input:    "TiDB-7.2.0",
+			expected: SemVersion{},
+			ok:       false,
+		},
+		{
+			name:     "invalid patch version",
+			input:    "5.7.25-TiDB-v7.1.x",
+			expected: SemVersion{},
+			ok:       false,
+		},
+		{
+			name:     "insufficient version parts",
+			input:    "5.7.25-TiDB-v7.1",
+			expected: SemVersion{},
+			ok:       false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			actual, ok := NewTiDBSemVersion(tc.input)
+			assert.Equal(t, tc.ok, ok, "ok mismatch")
+			if tc.ok {
+				assert.Equal(t, tc.expected, actual, "version mismatch")
+			}
+		})
+	}
+}