workload: enable workload_generator to run SQL workloads with init artifacts

This patch wires up the `workload_generator` command to consume the outputs from `workload init`—the schema YAML and generated SQL files—and drive end‑to‑end data‑driven query execution:

  • Load and unmarshal the schema YAML produced during initialization
  • Discover `<db>_read.sql` and `<db>_write.sql` files emitted by the SQL generator
  • Apply the existing placeholderRewriter at runtime to inject generated data into each statement
  • Execute the rewritten SQL workload against a CockroachDB cluster, generating rows and collecting metrics

Fixes: CRDB‑51752

Release note (cli change):
Enables `workload_generator` to read init‑time schema and SQL artifacts and run SQL workloads with placeholder‑driven data generation.

Author: Pradyum-Git

pkg/workload/workload_generator/BUILD.bazel

@@ -5,6 +5,7 @@ go_library(
     srcs = [
         "column_generator.go",
         "constants.go",
+        "run_utils.go",
         "schema_designs.go",
         "schema_generator.go",
         "schema_utils.go",
@@ -23,10 +24,13 @@ go_library(
         "//pkg/sql/sem/tree/treecmp",
         "//pkg/sql/types",
         "//pkg/util/bufalloc",
+        "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/workload",
         "//pkg/workload/histogram",
+        "@com_github_cockroachdb_cockroach_go_v2//crdb",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_lib_pq//:pq",
         "@com_github_spf13_pflag//:pflag",
         "@in_gopkg_yaml_v2//:yaml_v2",
     ],

pkg/workload/workload_generator/constants.go

@@ -18,4 +18,6 @@ const (
 	sqlForeignKey = "FOREIGN KEY" // sqlForeignKey is a constant string used to represent foreign key constraints in column definitions
 	sqlFamily     = "FAMILY"      // sqlFamily is a constant string used to represent family definitions in column definitions
 	sqlConstraint = "CONSTRAINT"  // sqlConstraint is a constant string used to represent constraint definitions in column definitions
+	insert        = "INSERT"      // insert is a constant string used to represent INSERT statements in SQL
+	update        = "UPDATE"      // update is a constant string used to represent UPDATE statements in SQL
 )

pkg/workload/workload_generator/run_utils.go

@@ -0,0 +1,278 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package workload_generator
+
+import (
+	"bufio"
+	gosql "database/sql"
+	"fmt"
+	"io"
+	"os"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+var (
+	// txnRe defines how to separate the transactions from the <dbName.sql> file
+	txnRe = regexp.MustCompile(`(?m)^-------Begin Transaction------\s*([\s\S]*?)-------End Transaction-------`)
+	// placeholderRe defines the structure of the placeholders that need to be converted into $1, $2...
+	placeholderRe = regexp.MustCompile(`"?:-:\|(.+?)\|:-:"?`)
+)
+
+// setColumnValue sets the value for a placeholder in the args slice at index i.
+func setColumnValue(raw string, placeholder Placeholder, args []interface{}, i int) error {
+	var arg interface{}
+	// If we got an empty string and this column is nullable, emitting a SQL NULL.
+	if raw == "" && placeholder.IsNullable {
+		arg = setNullType(placeholder)
+	} else {
+		// Otherwise the raw string is parsed into the right Go/sql type.
+		typedValue, err := setNotNullType(placeholder, raw)
+		if err != nil {
+			return err
+		}
+		arg = typedValue
+	}
+
+	args[i] = arg
+	return nil
+}
+
+// setNotNullType converts the raw string value to the appropriate SQL type based on the placeholder's column type.
+func setNotNullType(placeholder Placeholder, raw string) (interface{}, error) {
+	var arg interface{}
+	switch t := strings.ToUpper(placeholder.ColType); {
+	case strings.HasPrefix(t, "INT"): // integer types
+		iv, err := strconv.ParseInt(raw, 10, 64)
+		if err != nil {
+			return nil, err
+		}
+		arg = gosql.NullInt64{Int64: iv, Valid: true}
+	case strings.HasPrefix(t, "FLOAT"), strings.HasPrefix(t, "DECIMAL"), strings.HasPrefix(t, "NUMERIC"), strings.HasPrefix(t, "DOUBLE"): // floating point types
+		fv, err := strconv.ParseFloat(raw, 64)
+		if err != nil {
+			return nil, err
+		}
+		arg = gosql.NullFloat64{Float64: fv, Valid: true}
+	case t == "BOOL", t == "BOOLEAN": // boolean types
+		bv, err := strconv.ParseBool(raw)
+		if err != nil {
+			return nil, err
+		}
+		arg = gosql.NullBool{Bool: bv, Valid: true}
+	// The remaining types are parsed as raw strings.
+	default:
+		// Everything else is treated as text/varchar/etc.
+		arg = gosql.NullString{String: raw, Valid: raw != ""}
+	}
+	return arg, nil
+}
+
+// setNullType sets the argument to a SQL NULL value based on the column type of the placeholder.
+func setNullType(placeholder Placeholder) interface{} {
+	var arg interface{}
+	switch t := strings.ToUpper(placeholder.ColType); {
+	case strings.HasPrefix(t, "INT"):
+		arg = gosql.NullInt64{Valid: false}
+	case strings.HasPrefix(t, "FLOAT"), strings.HasPrefix(t, "DECIMAL"), strings.HasPrefix(t, "NUMERIC"), strings.HasPrefix(t, "DOUBLE"):
+		arg = gosql.NullFloat64{Valid: false}
+	case t == "BOOL", t == "BOOLEAN":
+		arg = gosql.NullBool{Valid: false}
+	default:
+		arg = gosql.NullString{Valid: false}
+	}
+	return arg
+}
+
+// getTableName checks if the name field of placeholder is a column in the TableName column in the allSchema map inside d.
+// If yes, then returns that table name. otherwise looks for a table with that column and returns that.
+func getTableName(p Placeholder, d *workloadGenerator) string {
+	for _, block := range d.workloadSchema[p.TableName] {
+		for colName := range block.Columns {
+			if colName == p.Name {
+				return p.TableName
+			}
+		}
+	}
+	for tableName, blocks := range d.workloadSchema {
+		block := blocks[0]
+		for colName := range block.Columns {
+			if colName == p.Name {
+				return tableName
+			}
+		}
+	}
+	return p.TableName
+}
+
+// getColumnValue retrieves the value for a placeholder based on its clause and whether it has a foreign key dependency.
+func getColumnValue(
+	allPksAreFK bool,
+	p Placeholder,
+	d *workloadGenerator,
+	inserted map[string][]interface{},
+	indexes []int,
+	i int,
+) string {
+	var raw string
+	if allPksAreFK && (p.Clause == insert || p.Clause == update) {
+		tableName := getTableName(p, d)
+		key := fmt.Sprintf("%s.%s", tableName, p.Name)
+		fk := d.columnGens[key].columnMeta.FK
+		parts := strings.Split(fk, ".")
+		parentCol := parts[len(parts)-1] // The last part is the column name.
+		if vals, ok := inserted[parentCol]; ok && len(vals) > 0 {
+			raw = vals[0].(string)         // Using the first value from the inserted map.
+			inserted[parentCol] = vals[1:] // Remove the first value from the map.
+		} else {
+			//Fallback that shouldn't really happen.
+			raw = d.getRegularColumnValue(p, indexes[i])
+		}
+	} else {
+		raw = d.getRegularColumnValue(p, indexes[i])
+	}
+	return raw
+}
+
+// checkIfAllPkAreFk checks if all primary keys in the SQL query are foreign keys.
+func checkIfAllPkAreFk(sqlQuery SQLQuery, d *workloadGenerator) bool {
+	allPksAreFK := true
+	for _, p := range sqlQuery.Placeholders {
+		tableName := getTableName(p, d)
+		key := fmt.Sprintf("%s.%s", tableName, p.Name)
+		if p.IsPrimaryKey && !d.columnGens[key].columnMeta.HasForeignKey {
+			allPksAreFK = false
+			break
+		}
+	}
+	return allPksAreFK
+}
+
+// readSQL reads <dbName><read/write>.sql and returns a slice of Transactions.
+// It will number placeholders $1…$N separately in each SQL statement.
+func readSQL(path, typ string) ([]Transaction, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("open file: %w", err)
+	}
+	defer func() {
+		_ = f.Close()
+	}()
+
+	data, err := io.ReadAll(bufio.NewReader(f))
+	if err != nil {
+		return nil, fmt.Errorf("read file: %w", err)
+	}
+	text := string(data)
+
+	// Each transaction block
+	blocks := txnRe.FindAllStringSubmatch(text, -1)
+
+	// Currently we are defining two types of transactions - read and write.
+	var txns []Transaction
+	// For every transaction block.
+	for _, blk := range blocks {
+		body := blk[1]
+		lines := strings.Split(body, "\n")
+		var txn Transaction
+		var curr SQLQuery
+		// For every sql query line.
+		for _, line := range lines {
+			line = strings.TrimSpace(line)
+			if line == "" || line == "BEGIN;" || line == "COMMIT;" {
+				continue
+			}
+			// build up the SQL text
+			curr.SQL += line + " "
+			if strings.HasSuffix(line, ";") {
+				// once we hit the end of a statement, re-number from $1
+				stmtPos := 1
+				var placeholders []Placeholder
+				// sqlOut is the rewritten SQL where the placeholders have been replaced with $x.
+				sqlOut := placeholderRe.ReplaceAllStringFunc(curr.SQL, getPlaceholderReplacer(&placeholders, &stmtPos))
+
+				curr.SQL = strings.TrimSpace(sqlOut)
+				curr.Placeholders = placeholders
+				txn.Queries = append(txn.Queries, curr)
+
+				// reset for next statement
+				curr = SQLQuery{}
+			}
+		}
+		// Decide whether the transaction is a read type or write type.
+		txn.typ = typ
+		txns = append(txns, txn)
+	}
+
+	return txns, nil
+}
+
+// getPlaceholderReplacer returns a ReplaceAllStringFunc that
+// appends to placeholders and increments stmtPos.
+func getPlaceholderReplacer(placeholders *[]Placeholder, stmtPos *int) func(string) string {
+	return func(match string) string {
+		inner := placeholderRe.FindStringSubmatch(match)[1]
+		parts := splitQuoted(inner)
+
+		var p Placeholder
+		//Set all the fields in the placeholder struct based on the information from the sql.
+		p.Name = trimQuotes(parts[0])
+		p.ColType = trimQuotes(parts[1])
+		p.IsNullable = trimQuotes(parts[2]) == "NULL"
+		p.IsPrimaryKey = strings.Contains(trimQuotes(parts[3]), "PRIMARY KEY")
+		if d := trimQuotes(parts[4]); d != "" {
+			p.Default = &d
+		}
+		p.IsUnique = trimQuotes(parts[5]) == "UNIQUE"
+		if fk := trimQuotes(parts[6]); fk != "" {
+			fkParts := strings.Split(strings.TrimPrefix(fk, "FK→"), ".")
+			p.FKReference = &FKRef{Table: fkParts[0], Column: fkParts[1]}
+		}
+		if chk := trimQuotes(parts[7]); chk != "" {
+			p.InlineCheck = &chk
+		}
+		p.Clause = trimQuotes(parts[8])
+		p.Position = *stmtPos
+		p.TableName = trimQuotes(parts[9])
+
+		*stmtPos++
+		*placeholders = append(*placeholders, p)
+		return fmt.Sprintf("$%d::%s", p.Position, p.ColType)
+	}
+}
+
+// splitQuoted splits a string like "'a','b','c'" into ["'a'", "'b'", "'c'"].
+func splitQuoted(s string) []string {
+	var out []string
+	buf := ""
+	inQuote := false
+	for _, r := range s {
+		switch r {
+		case '\'':
+			inQuote = !inQuote
+			buf += string(r)
+		case ',':
+			if inQuote {
+				buf += string(r)
+			} else {
+				out = append(out, strings.TrimSpace(buf))
+				buf = ""
+			}
+		default:
+			buf += string(r)
+		}
+	}
+	if buf != "" {
+		out = append(out, strings.TrimSpace(buf))
+	}
+	return out
+}
+
+// trimQuotes removes leading/trailing single-quotes.
+func trimQuotes(s string) string {
+	return strings.Trim(s, "'")
+}

pkg/workload/workload_generator/sql_generator.go

@@ -120,7 +120,13 @@ func generateWorkload(
 			// 5d) The sql query is processed to replace _ and __more__ with new placeholders which contain information about the column they refer to.
 			rewritten, err := replacePlaceholders(rawSQL, allSchemas)
 			if err != nil {
-				f.Close()
+				if errClose := f.Close(); errClose != nil {
+					// Wrap the original placeholder-rewrite error, then attach the close error.
+					return errors.WithSecondaryError(
+						errors.Wrapf(err, "rewriting SQL %q", rawSQL),
+						errClose,
+					)
+				}
 				return errors.Wrapf(err, "rewriting SQL %q", rawSQL)
 			}
 

pkg/workload/workload_generator/sql_utils.go

@@ -803,19 +803,25 @@ func writeTransaction(
 	if err != nil {
 		return errors.Wrapf(err, "creating %s", outPathRead)
 	}
-	defer outReadFile.Close()
+	defer func() {
+		_ = outReadFile.Close()
+	}()
 
 	outWriteFile, err := os.Create(outPathWrite)
 	if err != nil {
 		return errors.Wrapf(err, "creating %s", outPathWrite)
 	}
-	defer outWriteFile.Close()
+	defer func() {
+		_ = outWriteFile.Close()
+	}()
 
 	outUnhandledFile, err := os.Create(outPathUnhandled)
 	if err != nil {
 		return errors.Wrapf(err, "creating %s", outPathUnhandled)
 	}
-	defer outUnhandledFile.Close()
+	defer func() {
+		_ = outUnhandledFile.Close()
+	}()
 
 	for _, txnID := range txnOrder {
 		stmts := txnMap[txnID]

pkg/workload/workload_generator/types.go

@@ -214,7 +214,7 @@ func (u *UniqueWrapper) Next() string {
 			// accept it
 			u.seen[v] = struct{}{}
 			u.order = append(u.order, v)
-			// evict oldest if over capacity
+			// evict oldest value if over capacity
 			if len(u.order) > u.capacity {
 				oldest := u.order[0]
 				u.order = u.order[1:]