Merge #149513

149513: Ddl and yaml r=nameisbhaskar a=Pradyum-Git

Changed names of ddl_generator.go and its test file to schema_generator.go and schema_generator_test.go to better represent what they do
Also added in the code which takes the output from GenerateDDLs and converts them into a schema format which has all the necessary information to be  used later by data generators
the code declares all the required structs, adds buildWorkloadSchema function to teh schema_generator.go file which is the entry point for the new features, addsa. util.go file which have all the remaining fucntiosn taht fill out teh structs wiuth enough and proper data or constraint aware data generation for initial bulk load of data into tables


Release note: None
Epic: None

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

Author: craig[bot]

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 (