Merge #150926

150926: workload: add simple CHECK constraint support and bit/bytes generators r=nameisbhaskar a=Pradyum-Git

This patch enhances the workload_generator command by:

  • Parsing simple comparison CHECK constraints (>, >=, <, <=) from DDL and
    applying them to the generated schema (setting min/max or start/end args).
  • Introducing GenTypeBit and GenTypeBytes to handle BIT(n) and BYTES/BYTEA
    columns with dedicated generators.
  • Fixing minor bugs in SQL/YAML generation and type mapping.

Fixes: CRDB-51752
Epic: None

Release note (cli change):
Adds support for simple CHECK constraints and bit/bytes column generators in workload_generator.

Co-authored-by: Pradyum <[email protected]>

Author: craig[bot]

pkg/workload/workload_generator/column_generator.go

@@ -46,6 +46,10 @@ func buildGenerator(col ColumnMeta, batchIdx, batchSize int, schema Schema) Gene
 		base = buildBooleanGenerator(col, rng)
 	case GenTypeJson: //missed json type ig, will check
 		base = buildJsonGenerator(col, rng)
+	case GenTypeBit:
+		base = buildBitGenerator(col, rng)
+	case GenTypeBytes:
+		base = buildBytesGenerator(col, rng)
 
 	default:
 		panic("type not supported: " + col.Type)
@@ -226,3 +230,18 @@ func buildJsonGenerator(col ColumnMeta, rng *rand.Rand) Generator {
 	sg := &StringGen{r: rng, min: minArg, max: maxArg, nullPct: nullPct}
 	return &JsonGen{strGen: sg}
 }
+
+// buildBitGenerator produces random BIT(n) values as strings of '0'/'1'.
+func buildBitGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+	// size comes from mapBitType → args["size"]
+	size := getIntArg(col.Args, "size", 1)
+	nullPct := getFloatArg(col.Args, "null_pct", 0.0)
+	return &BitGen{r: rng, size: size, nullPct: nullPct}
+}
+
+// buildBytesGenerator produces random []byte for BYTEA/BYTES columns.
+func buildBytesGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+	size := getIntArg(col.Args, "size", 1)
+	nullPct := getFloatArg(col.Args, "null_pct", 0.0)
+	return &BytesGen{r: rng, min: size, max: size, nullPct: nullPct}
+}

pkg/workload/workload_generator/schema_generator.go

@@ -288,8 +288,9 @@ func processColumnDefs(table *TableSchema, columnDefs []string) {
 		defaultVal := colMatch[4]       // DEFAULT value
 		foreignKeyTable := colMatch[5]  // Referenced table for foreign keys
 		foreignKeyColumn := colMatch[6] // Referenced column for foreign keys
-
-		table.ColumnOrder = append(table.ColumnOrder, name)
+		// While adding columns to the order, surrounding quotes are stripped if any.
+		// This is to ensure that column names here match with the column names used as keys in schema maps.
+		table.ColumnOrder = append(table.ColumnOrder, strings.Trim(name, `"`))
 		// Extract CHECK constraint if present (requires special handling for nested parentheses)
 		inlineCheck := ""
 		checkIdx := checkInlineRe.FindStringIndex(columnDef)
@@ -565,6 +566,8 @@ func buildWorkloadSchema(
 	// 1) Build initial blocks and capture FK seeds
 	blocks, fkSeed := buildInitialBlocks(allSchemas, dbName, rng, baseRowCount)
 
+	applyCheckConstraints(blocks, allSchemas)
+
 	// 2) Wire up foreign-key relationships in the blocks
 	wireForeignKeys(blocks, allSchemas, fkSeed, rng)
 

pkg/workload/workload_generator/schema_utils.go

@@ -12,6 +12,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"time"
 )
 
 const (
@@ -23,6 +24,9 @@ const (
 	seedKeyDelimiter = "__"
 	// nullPct is the key for nullability percentage in args maps
 	nullPct = "null_pct"
+	// maxArg and minArg are keys for range limits in args maps
+	maxArg = "max"
+	minArg = "min"
 )
 
 // GeneratorType is an enum for all the data generator types.
@@ -55,6 +59,12 @@ var (
 	simpleNumberRe   = regexp.MustCompile(`^[+-]?\d+(?:\.\d+)?$`)
 	quotedStrRe      = regexp.MustCompile(`^'.*'$`)
 	booleanLiteralRe = regexp.MustCompile(`^(?i:true|false)$`)
+
+	// Regexes for simple comparison constraints.
+	gtRe  = regexp.MustCompile(`(?i)^([A-Za-z_]\w*)\s*>\s*([^\s]+)$`)
+	gteRe = regexp.MustCompile(`(?i)^([A-Za-z_]\w*)\s*>=\s*([^\s]+)$`)
+	ltRe  = regexp.MustCompile(`(?i)^([A-Za-z_]\w*)\s*<\s*([^\s]+)$`)
+	lteRe = regexp.MustCompile(`(?i)^([A-Za-z_]\w*)\s*<=\s*([^\s]+)$`)
 )
 
 // Schema is the map of TableBlocks, one per table, which is used by all data generators
@@ -285,13 +295,154 @@ func atoi(s string) int {
 
 // setArgsRange sets the "min" and "max" keys in the args map to the specified range.
 func setArgsRange(args map[string]any, min, max int) {
-	args["min"] = min
-	args["max"] = max
+	args[minArg] = min
+	args[maxArg] = max
 }
 
 // canonical replaces "." with "__" to match the legacy YAML format.
 func canonical(name string) string { return strings.ReplaceAll(name, ".", "__") }
 
+// bumpTimestampISO returns the given RFC3339Nano timestamp string
+// advanced by one nanosecond, or the original string if parsing fails.
+func bumpTimestampISO(s string) string {
+	if t, err := time.Parse(time.RFC3339Nano, s); err == nil {
+		return t.Add(time.Millisecond).Format(time.RFC3339Nano)
+	}
+	// fallback: return original
+	return s
+}
+
+// applyCheckConstraints updates each ColumnMeta in the Schema
+// to reflect simple comparison CHECK constraints (>, >=, <, <=)
+// for integer, float, and timestamp column types. It reads the
+// raw CheckConstraints from allSchemas and sets min/max/start/end
+// arguments in-place before data generation
+func applyCheckConstraints(blocks Schema, allSchemas map[string]*TableSchema) {
+	for tbl, blks := range blocks {
+		schema := allSchemas[tbl]
+		for i := range blks {
+			block := &blks[i]
+			// Iterating over each column in the current TableBlock.
+			for colName, cm := range block.Columns {
+				// Columns that are not integer, float, or timestamp are skipped.
+				switch cm.Type {
+				case GenTypeInteger, GenTypeFloat, GenTypeTimestamp:
+				default:
+					continue
+				}
+				// Examining each raw CHECK expression for supported patterns.
+				for _, chk := range schema.CheckConstraints {
+					chk = strings.TrimSpace(chk)
+					// col > val
+					applyGreaterThanCheck(chk, colName, cm)
+					// col >= val
+					applyGreaterThanEqualsCheck(chk, colName, cm)
+					// col < val
+					applyLessThanCheck(chk, colName, cm)
+					// col <= val
+					applyLessThanEqualsCheck(chk, colName, cm)
+				}
+				// The updated ColumnMeta is written back.
+				block.Columns[colName] = cm
+			}
+		}
+	}
+}
+
+// applyLessThanEqualsCheck checks if the given check constraint
+// is a "col <= val" expression and updates the ColumnMeta accordingly.
+func applyLessThanEqualsCheck(chk string, colName string, cm ColumnMeta) {
+	if m := lteRe.FindStringSubmatch(chk); m != nil && m[1] == colName {
+		lit := stripCast(m[2])
+		switch cm.Type {
+		case GenTypeInteger:
+			if v, err := strconv.Atoi(lit); err == nil {
+				cm.Args[maxArg] = v
+			}
+		case GenTypeFloat:
+			if f, err := strconv.ParseFloat(lit, 64); err == nil {
+				cm.Args[maxArg] = f
+			}
+		case GenTypeTimestamp:
+			cm.Args["end"] = lit
+		}
+	}
+}
+
+// applyLessThanCheck checks if the given check constraint
+// is a "col < val" expression and updates the ColumnMeta accordingly.
+func applyLessThanCheck(chk string, colName string, cm ColumnMeta) {
+	if m := ltRe.FindStringSubmatch(chk); m != nil && m[1] == colName {
+		lit := stripCast(m[2])
+		switch cm.Type {
+		case GenTypeInteger:
+			if v, err := strconv.Atoi(lit); err == nil {
+				cm.Args[maxArg] = v - 1
+			}
+		case GenTypeFloat:
+			if f, err := strconv.ParseFloat(lit, 64); err == nil {
+				cm.Args[maxArg] = math.Nextafter(f, math.Inf(-1))
+			}
+		case GenTypeTimestamp:
+			// Subtract one millisecond
+			if t, err := time.Parse(time.RFC3339Nano, lit); err == nil {
+				cm.Args["end"] = t.Add(-time.Millisecond).Format(time.RFC3339Nano)
+			}
+		}
+	}
+}
+
+// applyLessThanEqualsCheck checks if the given check constraint
+// is a "col <= val" expression and updates the ColumnMeta accordingly.
+func applyGreaterThanEqualsCheck(chk string, colName string, cm ColumnMeta) {
+	if m := gteRe.FindStringSubmatch(chk); m != nil && m[1] == colName {
+		lit := stripCast(m[2])
+		switch cm.Type {
+		case GenTypeInteger:
+			if v, err := strconv.Atoi(lit); err == nil {
+				cm.Args[minArg] = v
+			}
+		case GenTypeFloat:
+			if f, err := strconv.ParseFloat(lit, 64); err == nil {
+				cm.Args[minArg] = f
+			}
+		case GenTypeTimestamp:
+			cm.Args["start"] = lit
+		}
+	}
+}
+
+// applyGreaterThanCheck checks if the given check constraint
+// is a "col > val" expression and updates the ColumnMeta accordingly.
+func applyGreaterThanCheck(chk string, colName string, cm ColumnMeta) {
+	if m := gtRe.FindStringSubmatch(chk); m != nil && m[1] == colName {
+		lit := stripCast(m[2])
+		switch cm.Type {
+		case GenTypeInteger:
+			if v, err := strconv.Atoi(lit); err == nil {
+				cm.Args[minArg] = v + 1
+			}
+		case GenTypeFloat:
+			if f, err := strconv.ParseFloat(lit, 64); err == nil {
+				cm.Args[minArg] = math.Nextafter(f, math.Inf(1))
+			}
+		case GenTypeTimestamp:
+			cm.Args["start"] = bumpTimestampISO(lit)
+		}
+	}
+}
+
+// stripCast removes any Cockroach/Postgres cast suffix:
+//
+//	“123:::INT8” → “123”
+//	“‘2021-01-01’::DATE” → “'2021-01-01'”
+func stripCast(lit string) string {
+	if i := strings.Index(lit, "::"); i >= 0 {
+		return lit[:i]
+	}
+	return lit
+}
+
 // mapSQLType maps a SQL column type to the workload generator type and
 // argument set expected by cockroach workloads. The returned map may
 // include bounds, formatting information or other hints used by the
@@ -331,15 +482,16 @@ func mapSQLType(sql string, col *Column, rng *rand.Rand) (GeneratorType, map[str
 
 	case sql == "date":
 		return mapDateType(sql, col, args)
-
-	case sql == "timestamp" || sql == "timestamptz":
+	// We use hasPrefix(timestamp) to match both "timestamp" and "timestamptz"
+	// as well as any other variations like "timestamp(6)".
+	case strings.HasPrefix(sql, "timestamp"):
 		return mapTimestampType(sql, col, args)
 
 	case sql == "bool" || sql == "boolean":
 		return GenTypeBool, args
 
 	case sql == "json" || sql == "jsonb":
-		mapJsonType(sql, col, args)
+		return mapJsonType(sql, col, args)
 	}
 	setArgsRange(args, 5, 30)
 	return GenTypeString, args
@@ -404,13 +556,13 @@ func mapDecimalType(sql string, _ *Column, args map[string]any) (GeneratorType,
 			maxVal = 1.0 - fracUnit
 			minVal = -maxVal
 		}
-		args["min"] = minVal
-		args["max"] = maxVal
+		args[minArg] = minVal
+		args[maxArg] = maxVal
 		args["round"] = scale
 	} else {
 		// fallback for DECIMAL without precision
-		args["min"] = 0.0
-		args["max"] = 1.0
+		args[minArg] = 0.0
+		args[maxArg] = 1.0
 		args["round"] = 2
 	}
 	return GenTypeFloat, args

pkg/workload/workload_generator/sql_utils.go

@@ -516,10 +516,10 @@ func getFieldCol(
 	numParts := 1
 	parts := make([]string, numParts)
 	var placeholder string
-	for _, schema := range allSchemas {
+	for tn, schema := range allSchemas {
 		matched := false
 		for _, column := range schema.Columns {
-			if column.Name == col {
+			if column.Name == col && tn == tableName {
 				placeholder = column.String()
 				placeholder = placeholder[1 : len(placeholder)-1]
 				matched = true

pkg/workload/workload_generator/types.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"math/rand"
 	"strconv"
+	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 )
@@ -157,6 +158,45 @@ func (g *JsonGen) Next() string {
 	return fmt.Sprintf(`{"k":"%s"}`, v)
 }
 
+// BitGen emits either nil (per nullPct) or a string like "101001".
+type BitGen struct {
+	r       *rand.Rand
+	size    int
+	nullPct float64
+}
+
+func (b *BitGen) Next() string {
+	if b.r.Float64() < b.nullPct {
+		return ""
+	}
+	var sb strings.Builder
+	for i := 0; i < b.size; i++ {
+		if b.r.Intn(2) == 0 {
+			sb.WriteByte('0')
+		} else {
+			sb.WriteByte('1')
+		}
+	}
+	return sb.String()
+}
+
+type BytesGen struct {
+	r        *rand.Rand
+	min, max int
+	nullPct  float64
+}
+
+func (b *BytesGen) Next() string {
+	if b.r.Float64() < b.nullPct {
+		return ""
+	}
+
+	length := b.min
+	buf := make([]byte, length)
+	b.r.Read(buf) // fill with random bytes
+	return string(buf)
+}
+
 // ─── Wrappers ──────────────────────────────────────────────────────────
 
 // DefaultWrapper wraps a base Generator and emits a fixed literal