workload/ycsb: add --isolation-level flag

Informs cockroachdb#107112.

This commit adds an `--isolation-level` flag to ycsb, which controls the
isolation level to run the workload transactions under. If unset, the
workload will run with the default isolation level of the database.

Release note: None

Author: nvb

pkg/workload/connection.go

@@ -107,3 +107,33 @@ func SanitizeUrls(gen Generator, dbOverride string, urls []string) (string, erro
 	}
 	return dbName, nil
 }
+
+// SetDefaultIsolationLevel configures the provided URLs with the specified
+// default transaction isolation level, if any.
+func SetDefaultIsolationLevel(urls []string, isoLevel string) error {
+	if isoLevel == "" {
+		return nil
+	}
+	// As a convenience, replace underscores with spaces. This allows users of the
+	// workload tool to pass --isolation-level=read_committed instead of needing
+	// to pass --isolation-level="read committed".
+	isoLevel = strings.ReplaceAll(isoLevel, "_", " ")
+	// NOTE: validation of the isolation level value is done by the server during
+	// connection establishment.
+	return setUrlParam(urls, "default_transaction_isolation", isoLevel)
+}
+
+// setUrlParam sets the given parameter to the given value in the provided URLs.
+func setUrlParam(urls []string, param, value string) error {
+	for i := range urls {
+		parsed, err := url.Parse(urls[i])
+		if err != nil {
+			return err
+		}
+		q := parsed.Query()
+		q.Set(param, value)
+		parsed.RawQuery = q.Encode()
+		urls[i] = parsed.String()
+	}
+	return nil
+}

pkg/workload/ycsb/ycsb.go

@@ -93,6 +93,7 @@ type ycsb struct {
 	flags     workload.Flags
 	connFlags *workload.ConnFlags
 
+	isoLevel    string
 	timeString  bool
 	insertHash  bool
 	zeroPadding int
@@ -125,8 +126,10 @@ var ycsbMeta = workload.Meta{
 		g := &ycsb{}
 		g.flags.FlagSet = pflag.NewFlagSet(`ycsb`, pflag.ContinueOnError)
 		g.flags.Meta = map[string]workload.FlagMeta{
-			`workload`: {RuntimeOnly: true},
+			`isolation-level`: {RuntimeOnly: true},
+			`workload`:        {RuntimeOnly: true},
 		}
+		g.flags.StringVar(&g.isoLevel, `isolation-level`, ``, `Isolation level to run workload transactions under [serializable, snapshot, read_committed]. If unset, the workload will run with the default isolation level of the database.`)
 		g.flags.BoolVar(&g.timeString, `time-string`, false, `Prepend field[0-9] data with current time in microsecond precision.`)
 		g.flags.BoolVar(&g.insertHash, `insert-hash`, true, `Key to be hashed or ordered.`)
 		g.flags.IntVar(&g.zeroPadding, `zero-padding`, 1, `Key using "insert-hash=false" has zeros padded to left to make this length of digits.`)
@@ -378,6 +381,9 @@ func (g *ycsb) Ops(
 	if err != nil {
 		return workload.QueryLoad{}, err
 	}
+	if err := workload.SetDefaultIsolationLevel(urls, g.isoLevel); err != nil {
+		return workload.QueryLoad{}, err
+	}
 
 	const readStmtStr = `SELECT * FROM usertable WHERE ycsb_key = $1`