workload_generator: Fixes and enhancements

This is a continuation of the effort for the workload generator which was worked on as an intern project.
This PR addresses multiple issues and enhancements as mentioned:

Columns were not identified properly in case there were Joins because only 1 table name was considered from a statement.
e.g. SELECT o.id, r.c2, r.c3 FROM orders o JOIN ref_data r ON o.acc_no = r.acc_no JOIN ref_data r2 ON o.acc_no = r2.acc_no WHERE r.c2 = _

Columns with table_alias.column_name is not handled as table aliases were not maintained.
e.g. SELECT o.id, r.c2 FROM orders o, ref_data r WHERE o.acc_no = r.acc_no AND o.amount > _ AND r.acc_no = _ OR o.acc_no = _;

table names prefixed with database name in it were not parsed correctly. In this case, the database_name column in the crdb_internal.node_statement_statistics.txt file, can have a different database name and not the actual database name.
e.g. select a,b,c from db_1.table_name.

If a query is run from the current database scope, the queries are marked as the current database even if the query is run for another database.
e.g. SQL> use db1;
SQL> select * from db2.table_in_db2;
The query for db is marked under the database name db1 in the crdb_internal.node_statement_statistics.txt file. This causes an issue as unsupported queries also get added for run.

Currently, the workload cannot be re-run as the last value of a sequence column is not maintained from the previous run. So, we had to fresh import the data before each run. Now, the last_value is stored in the YAML so that it can be reused in the next run.

There are some additional cleanup changes:

Use references instead of pass by value
The block for each table was a list, which was unnecessary.
checkIfAllPkAreFk had a bug where the function used to return true if there were no primary key in the query.
Ensure that we use import instead of insert during dataload. This is fixed by initializing the workload schema in tha Tables call, which ensures that the "support fixtures" check is set to true.

Fixes: CRDB-51572
Epic: None

Author: Pradyum-Git

pkg/workload/workload_generator/column_generator.go

@@ -20,7 +20,7 @@ import (
 //   - batchIdx:  which batch we’re in (0,1,2…)
 //   - baseBatchSize: how many rows per batch (for sequences)
 //   - schema:    full YAML schema, so we can recurse on FKs
-func buildGenerator(col ColumnMeta, batchIdx, batchSize int, schema Schema) Generator {
+func buildGenerator(col *ColumnMeta, batchIdx, batchSize int, schema Schema) Generator {
 	// seed the per-batch RNG
 	origSeed := getIntArg(col.Args, "seed", 0)
 	seed64 := buildBatchSeed(origSeed, batchIdx)
@@ -80,7 +80,7 @@ func buildGenerator(col ColumnMeta, batchIdx, batchSize int, schema Schema) Gene
 		parentTable := pathSegments[len(pathSegments)-1] // "district"
 
 		// now look up in your schema map:
-		parentMeta := schema[parentTable][0].Columns[childCol]
+		parentMeta := schema[parentTable].Columns[childCol]
 		parentGen := buildGenerator(parentMeta, batchIdx, batchSize, schema)
 		base = NewFkWrapper(parentGen, col.Fanout)
 	}
@@ -150,21 +150,28 @@ func buildBatchSeed(origSeed, batchIdx int) int64 {
 }
 
 // buildSequenceGenerator creates a SequenceGen that generates sequential integers
-func buildSequenceGenerator(col ColumnMeta, batchIdx int, batchSize int) Generator {
-	baseStart := getIntArg(col.Args, "start", 0)
-	return &SequenceGen{cur: baseStart + batchIdx*batchSize}
+func buildSequenceGenerator(col *ColumnMeta, batchIdx int, batchSize int) Generator {
+	current := 0
+	// LastValue should always be present in the YAML. Just in case this value is not set, the default behavior
+	// is to start from the "start" argument, which defaults to 0 + batchIdx*batchSize.
+	if col.LastValue != "" {
+		current, _ = strconv.Atoi(col.LastValue)
+	} else {
+		current = getIntArg(col.Args, "start", 0) + batchIdx*batchSize
+	}
+	return &SequenceGen{cur: current}
 }
 
 // buildIntegerGenerator creates an IntegerGen that generates random integers
-func buildIntegerGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+func buildIntegerGenerator(col *ColumnMeta, rng *rand.Rand) Generator {
 	minArg := getIntArg(col.Args, "min", 0)
 	maxArg := getIntArg(col.Args, "max", 0)
 	nullPct := getFloatArg(col.Args, "null_pct", 0.0)
 	return &IntegerGen{r: rng, min: minArg, max: maxArg, nullPct: nullPct}
 }
 
 // buildFloatGenerator creates a FloatGen that generates random floats
-func buildFloatGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+func buildFloatGenerator(col *ColumnMeta, rng *rand.Rand) Generator {
 	minArg := getFloatArg(col.Args, "min", 0.0)
 	maxArg := getFloatArg(col.Args, "max", 0.0)
 	round := getIntArg(col.Args, "round", 2)
@@ -173,15 +180,15 @@ func buildFloatGenerator(col ColumnMeta, rng *rand.Rand) Generator {
 }
 
 // buildStringGenerator creates a StringGen that generates random strings
-func buildStringGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+func buildStringGenerator(col *ColumnMeta, rng *rand.Rand) Generator {
 	minArg := getIntArg(col.Args, "min", 0)
 	maxArg := getIntArg(col.Args, "max", 0)
 	nullPct := getFloatArg(col.Args, "null_pct", 0.0)
 	return &StringGen{r: rng, min: minArg, max: maxArg, nullPct: nullPct}
 }
 
 // buildTimestampGenerator creates a TimestampGen that generates random timestamps
-func buildTimestampGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+func buildTimestampGenerator(col *ColumnMeta, rng *rand.Rand) Generator {
 	// parse Python-style format → Go layout
 	startStr := getStringArg(col.Args, "start", "2000-01-01")
 	endStr := getStringArg(col.Args, "end", timeutil.Now().Format("2006-01-02"))
@@ -211,18 +218,18 @@ func buildTimestampGenerator(col ColumnMeta, rng *rand.Rand) Generator {
 }
 
 // buildUuidGenerator creates a UUIDGen that generates random UUIDs
-func buildUuidGenerator(_ ColumnMeta, rng *rand.Rand) Generator {
+func buildUuidGenerator(_ *ColumnMeta, rng *rand.Rand) Generator {
 	return &UUIDGen{r: rng}
 }
 
 // buildBooleanGenerator creates a BoolGen that generates random booleans
-func buildBooleanGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+func buildBooleanGenerator(col *ColumnMeta, rng *rand.Rand) Generator {
 	nullPct := getFloatArg(col.Args, "null_pct", 0.0)
 	return &BoolGen{r: rng, nullPct: nullPct}
 }
 
 // buildJsonGenerator creates a JsonGen that generates random JSON strings
-func buildJsonGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+func buildJsonGenerator(col *ColumnMeta, rng *rand.Rand) Generator {
 	// JSON is just a StringGen plus a wrapper
 	minArg := getIntArg(col.Args, "min", defaultJSONMinLen)
 	maxArg := getIntArg(col.Args, "max", defaultJSONMaxLen)
@@ -232,15 +239,15 @@ func buildJsonGenerator(col ColumnMeta, rng *rand.Rand) Generator {
 }
 
 // buildBitGenerator produces random BIT(n) values as strings of '0'/'1'.
-func buildBitGenerator(col ColumnMeta, rng *rand.Rand) Generator {
+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 {
+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/column_generator_test.go

@@ -62,15 +62,15 @@ func TestSplitmix64AndBatchSeed(t *testing.T) {
 }
 
 func TestSequenceGenerator(t *testing.T) {
-	col := ColumnMeta{Args: map[string]interface{}{"start": 5}}
+	col := &ColumnMeta{Args: map[string]interface{}{"start": 5}}
 	gen := buildSequenceGenerator(col, 2, 10)
 	assert.Equal(t, "25", gen.Next(), "first value")
 	assert.Equal(t, "26", gen.Next(), "second value")
 }
 
 func TestIntegerAndFloatGenerators(t *testing.T) {
 	rng := rand.New(rand.NewSource(0))
-	icol := ColumnMeta{Args: map[string]interface{}{"min": 1, "max": 3, "null_pct": 0.0}}
+	icol := &ColumnMeta{Args: map[string]interface{}{"min": 1, "max": 3, "null_pct": 0.0}}
 	igen := buildIntegerGenerator(icol, rng)
 	for i := 0; i < 10; i++ {
 		v := igen.Next()
@@ -79,7 +79,7 @@ func TestIntegerAndFloatGenerators(t *testing.T) {
 		assert.GreaterOrEqual(t, n, 1)
 		assert.LessOrEqual(t, n, 3)
 	}
-	fcol := ColumnMeta{Args: map[string]interface{}{"min": 1.0, "max": 2.0, "round": 1, "null_pct": 0.0}}
+	fcol := &ColumnMeta{Args: map[string]interface{}{"min": 1.0, "max": 2.0, "round": 1, "null_pct": 0.0}}
 	fgen := buildFloatGenerator(fcol, rng)
 	for i := 0; i < 10; i++ {
 		v := fgen.Next()
@@ -88,7 +88,7 @@ func TestIntegerAndFloatGenerators(t *testing.T) {
 }
 
 func TestTimestampGenerator(t *testing.T) {
-	col := ColumnMeta{Args: map[string]interface{}{"start": "2000-01-02", "end": "2000-01-03", "format": "%Y-%m-%d"}}
+	col := &ColumnMeta{Args: map[string]interface{}{"start": "2000-01-02", "end": "2000-01-03", "format": "%Y-%m-%d"}}
 	rng := rand.New(rand.NewSource(0))
 	gen := buildTimestampGenerator(col, rng)
 	v := gen.Next()
@@ -98,7 +98,7 @@ func TestTimestampGenerator(t *testing.T) {
 
 func TestUuidGenerator(t *testing.T) {
 	rng := rand.New(rand.NewSource(0))
-	col := ColumnMeta{}
+	col := &ColumnMeta{}
 	gen := buildUuidGenerator(col, rng)
 	v1 := gen.Next()
 	v2 := gen.Next()
@@ -110,20 +110,20 @@ func TestUuidGenerator(t *testing.T) {
 
 func TestStringBoolJsonGenerators(t *testing.T) {
 	rng := rand.New(rand.NewSource(1))
-	scol := ColumnMeta{Args: map[string]interface{}{"min": 2, "max": 4, "null_pct": 0.0}}
+	scol := &ColumnMeta{Args: map[string]interface{}{"min": 2, "max": 4, "null_pct": 0.0}}
 	sgen := buildStringGenerator(scol, rng)
 	for i := 0; i < 5; i++ {
 		v := sgen.Next()
 		assert.GreaterOrEqual(t, len(v), 2)
 		assert.LessOrEqual(t, len(v), 4)
 	}
-	bcol := ColumnMeta{Args: map[string]interface{}{"null_pct": 0.0}}
+	bcol := &ColumnMeta{Args: map[string]interface{}{"null_pct": 0.0}}
 	bgen := buildBooleanGenerator(bcol, rng)
 	for i := 0; i < 5; i++ {
 		v := bgen.Next()
 		assert.Contains(t, []string{"0", "1"}, v)
 	}
-	jcol := ColumnMeta{Args: map[string]interface{}{"min": 1, "max": 1, "null_pct": 0.0}}
+	jcol := &ColumnMeta{Args: map[string]interface{}{"min": 1, "max": 1, "null_pct": 0.0}}
 	jgen := buildJsonGenerator(jcol, rng)
 	v := jgen.Next()
 	assert.True(t, strings.HasPrefix(v, "{\"k\":\"") && strings.HasSuffix(v, "\"}"), "JSON object")
@@ -177,7 +177,7 @@ func TestMakeGeneratorAllTypes(t *testing.T) {
 		}},
 	}
 	for _, tt := range types {
-		gen := buildGenerator(ColumnMeta{Type: tt.typ, Args: tt.args}, 0, 1, schema)
+		gen := buildGenerator(&ColumnMeta{Type: tt.typ, Args: tt.args}, 0, 1, schema)
 		v := gen.Next()
 		assert.True(t, tt.checker(v), "%s generator produced %q", tt.typ, v)
 	}