workload_generator: convert parsed DDL schemas into generator-ready YAML format

Previously, the workload_generator only parsed CREATE TABLE DDLs to extract schema metadata and built anonymization mappings for future use.

This was inadequate because there was no structured representation of generator parameters per column or table, nor any built-in support for marshaling that information into YAML for downstream workload generation.

To address this, this patch extends the parsed TableSchema into TableBlock and ColumnMeta types containing all necessary generator arguments (fanout, nullability, defaults, FK seeds, etc.), centralizes the buildWorkloadSchema orchestration, and adds easy YAML marshalling for future generator pipelines.

Fixes: CRDB-51752

Release note (cli change): The workload_generator now supports parsing DDL schemas into a structured YAML format, enabling more flexible and detailed workload generation configurations.

Author: Pradyum-Git

pkg/workload/workload_generator/BUILD.bazel

@@ -4,17 +4,24 @@ go_library(
     name = "workload_generator",
     srcs = [
         "constants.go",
-        "ddl_generator.go",
         "schema_designs.go",
+        "schema_generator.go",
+        "utils.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/workload/workload_generator",
     visibility = ["//visibility:public"],
-    deps = ["@com_github_cockroachdb_errors//:errors"],
+    deps = [
+        "//pkg/util/timeutil",
+        "@com_github_cockroachdb_errors//:errors",
+    ],
 )
 
 go_test(
     name = "workload_generator_test",
-    srcs = ["ddl_generator_test.go"],
+    srcs = [
+        "schema_generator_test.go",
+        "utils_test.go",
+    ],
     embed = [":workload_generator"],
     embedsrcs = ["test_data/debug/crdb_internal.create_statements.txt"],
     deps = ["@com_github_stretchr_testify//assert"],

pkg/workload/workload_generator/schema_designs.go

@@ -147,7 +147,7 @@ func (ts *TableSchema) SetPrimaryKeys(pks []string) {
 // ColumnMeta is the per-column metadata (type, args, FK info, etc.) that
 // drives our per batch generators.
 type ColumnMeta struct {
-	Type          string                 `yaml:"type"`
+	Type          GeneratorType          `yaml:"type"`
 	Args          map[string]interface{} `yaml:"args"`
 	IsPrimaryKey  bool                   `yaml:"isPrimaryKey"`
 	IsUnique      bool                   `yaml:"isUnique"`

pkg/workload/workload_generator/ddl_generator.go renamed to pkg/workload/workload_generator/schema_generator.go

@@ -10,11 +10,13 @@ import (
 	"encoding/csv"
 	"fmt"
 	"io"
+	"math/rand"
 	"os"
 	"path/filepath"
 	"regexp"
 	"strings"
 
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 )
 
@@ -524,3 +526,52 @@ func processDDLRecord(
 	}
 	tableStatements[fullTable] = stmt
 }
+
+// buildWorkloadSchema constructs the complete workload schema used for data generation.
+// It takes each parsed TableSchema and produces one or more TableBlock entries per table,
+// wiring up foreign‐key relationships and scaling row counts appropriately.
+//
+// The steps are as follows:
+//  1. buildInitialBlocks:
+//     • Creates a TableBlock for each table with a baseline row count (baseRowCount).
+//     • Converts each Column into ColumnMeta (type, null probability, default, etc.).
+//     • Collects “seeds” for foreign‐key columns to enable parent→child linkage.
+//  2. wireForeignKeys:
+//     • Scans each TableSchema’s ForeignKeys and populates the corresponding ColumnMeta. FK,
+//     FKMode, Fanout, CompositeID, and ParentSeed entries.
+//  3. adjustFanoutForPureFKPKs:
+//     • If a table’s every primary‐key column is also a foreign key, drops its fan-out to 1,
+//     ensuring exactly one child per parent in that “pure FK–PK” scenario.
+//  4. computeRowCounts:
+//     • For each table, computes the total row count by multiplying baseRowCount
+//     by the smallest product of FK fan-outs, ensuring referential integrity.
+//
+// Parameters:
+//   - allSchemas:    map of simple table name → *TableSchema, parsed from the DDL.
+//   - dbName:        the database name (used for qualifying schema references).
+//   - baseRowCount:  the initial number of rows per table before applying FK scaling.
+//
+// Returns:
+//   - Schema: the finalized workload schema, mapping each table name to a slice of one
+//     or more TableBlock, each of which drives per-batch generators with the
+//     correct column metadata and row counts.
+func buildWorkloadSchema(
+	allSchemas map[string]*TableSchema, dbName string, baseRowCount int,
+) Schema {
+	// Initialize RNG for seeding and composite IDs
+	rng := rand.New(rand.NewSource(timeutil.Now().UnixNano()))
+
+	// 1) Build initial blocks and capture FK seeds
+	blocks, fkSeed := buildInitialBlocks(allSchemas, dbName, rng, baseRowCount)
+
+	// 2) Wire up foreign-key relationships in the blocks
+	wireForeignKeys(blocks, allSchemas, fkSeed, rng)
+
+	// 3) If a table's PK cols are all FKs, drop its fanout to 1
+	adjustFanoutForPureFKPKs(blocks)
+
+	// 4) Recompute each block's row count based on FK fanouts
+	computeRowCounts(blocks, baseRowCount)
+
+	return blocks
+}

pkg/workload/workload_generator/ddl_generator_test.go renamed to pkg/workload/workload_generator/schema_generator_test.go

@@ -3,8 +3,8 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
-// ddl_generator_test.go
-// Unit tests for ddl_generator.go
+// schema_generator_test.go
+// Unit tests for schema_generator.go
 package workload_generator
 
 import (