Merge #150614

150614: workload: add SQL generation and placeholder‐rewriting utilities r=nameisbhaskar a=Pradyum-Git

This patch introduces two new files to support end-to-end SQL workload extraction and rewriting:

  • sql_utils.go: low-level helpers for parsing statement-statistics TSV headers, classifying transactions as read vs write, extracting unaliased table names, and emitting BEGIN/COMMIT blocks in read/write SQL files :contentReference[oaicite:0]{index=0}
  • generate_sql.go: core workload generator that scans debug-log TSVs, filters by database, unquotes raw SQL, parses statements into an AST, and rewrites all “_”/“__more__” placeholders via a placeholderRewriter (handling LIMIT, INSERT, UPDATE, BETWEEN, IN, CASE, COALESCE, etc.), before grouping statements into `<db>_read.sql` and `<db>_write.sql` :contentReference[oaicite:1]{index=1}

Key components:
  - `replacePlaceholders`: parses raw SQL, applies `buildPlaceholderRewriter`, and formats rewritten statements
  - `placeholderRewriter` visitor: comprehensive AST pass for tagging placeholders with column metadata
  - Specialized handlers (`handleSelectLimit`, `handleInsert`, `handleUpdateSet`, `handleRangeCondition`, `handleInOperator`, etc.) for different SQL constructs
  - Updated `GenerateWorkload` entrypoint: drives TSV scanning, placeholder rewriting, and file output

Fixes: CRDB-51752

Release note (cli change): Adds SQL workload extraction and rewriting support, enabling placeholder‐driven data-generation workflows from CockroachDB debug logs.

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

Author: craig[bot]

pkg/workload/workload_generator/BUILD.bazel

@@ -7,14 +7,20 @@ go_library(
         "constants.go",
         "schema_designs.go",
         "schema_generator.go",
+        "schema_utils.go",
+        "sql_generator.go",
+        "sql_utils.go",
         "types.go",
-        "utils.go",
         "workload.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/workload/workload_generator",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/col/coldata",
+        "//pkg/sql/parser",
+        "//pkg/sql/parser/statements",
+        "//pkg/sql/sem/tree",
+        "//pkg/sql/sem/tree/treecmp",
         "//pkg/sql/types",
         "//pkg/util/bufalloc",
         "//pkg/util/timeutil",
@@ -31,9 +37,13 @@ go_test(
     srcs = [
         "column_generator_test.go",
         "schema_generator_test.go",
-        "utils_test.go",
+        "schema_utils_test.go",
+        "sql_test.go",
     ],
     embed = [":workload_generator"],
     embedsrcs = ["test_data/debug/crdb_internal.create_statements.txt"],
-    deps = ["@com_github_stretchr_testify//assert"],
+    deps = [
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
+    ],
 )

pkg/workload/workload_generator/constants.go

@@ -16,4 +16,6 @@ const (
 	sqlDefault    = "DEFAULT"     // sqlDefault is a constant string used to represent default value expressions in column definitions
 	sqlCheck      = "CHECK"       // sqlCheck is a constant string used to represent CHECK constraints in column definitions
 	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
 )

pkg/workload/workload_generator/schema_designs.go

@@ -31,28 +31,55 @@ type Column struct {
 
 // String function converts the Column schema details into a parsable placeholder.
 func (c *Column) String() string {
-	parts := []string{c.Name, c.ColType}
+	// 1. An 8-element slice, containing teh information stored in the column meta is built.
+	parts := make([]string, 8)
+	parts[0] = c.Name
+	parts[1] = c.ColType
+
 	if c.IsNullable {
-		parts = append(parts, sqlNull)
+		parts[2] = "NULL"
 	} else {
-		parts = append(parts, sqlNotNull)
+		parts[2] = "NOT NULL"
 	}
+
 	if c.IsPrimaryKey {
-		parts = append(parts, sqlPrimaryKey)
+		parts[3] = "PRIMARY KEY"
+	} else {
+		parts[3] = ""
 	}
+
 	if c.Default != "" {
-		parts = append(parts, fmt.Sprintf("%s %s", sqlDefault, c.Default))
+		parts[4] = "DEFAULT " + c.Default
+	} else {
+		parts[4] = ""
 	}
+
 	if c.IsUnique {
-		parts = append(parts, sqlUnique)
+		parts[5] = "UNIQUE"
+	} else {
+		parts[5] = ""
 	}
+
 	if c.FKTable != "" && c.FKColumn != "" {
-		parts = append(parts, fmt.Sprintf("%s→%s.%s", sqlForeignKey, c.FKTable, c.FKColumn))
+		parts[6] = fmt.Sprintf("FK→%s.%s", c.FKTable, c.FKColumn)
+	} else {
+		parts[6] = ""
 	}
+
 	if c.InlineCheck != "" {
-		parts = append(parts, fmt.Sprintf("%s(%s)", sqlCheck, c.InlineCheck))
+		parts[7] = fmt.Sprintf("CHECK(%s)", c.InlineCheck)
+	} else {
+		parts[7] = ""
 	}
-	return strings.Join(parts, " ")
+
+	// 2. Each part (empty → "''") is quoted, escaping any internal apostrophes.
+	for i, p := range parts {
+		escaped := strings.ReplaceAll(p, "'", "\\'")
+		parts[i] = fmt.Sprintf("'%s'", escaped)
+	}
+
+	// 3. The parts are joined with commas.
+	return strings.Join(parts, ",")
 }
 
 // TableSchema stores table level schema information based on input ddl.

pkg/workload/workload_generator/schema_generator.go

@@ -80,7 +80,7 @@ var (
 	createTableRe = regexp.MustCompile(`(?i)^(CREATE\s+TABLE\s+)`)
 )
 
-// GenerateDDLs extracts and processes DDL statements from a CockroachDB debug zip file.
+// generateDDLs extracts and processes DDL statements from a CockroachDB debug zip file.
 // It reads the create_statements.txt file from the zip directory, filters statements
 // for the specified database, and writes them to an output file. It also parses each
 // DDL statement into a TableSchema object and returns a map of table names to their schemas.
@@ -96,7 +96,7 @@ var (
 //   - error: Any error encountered during processing
 //
 // TODO: The "anonymize" parameter is unused for now.
-func GenerateDDLs(
+func generateDDLs(
 	zipDir,
 	dbName string, anonymize bool,
 ) (allSchemas map[string]*TableSchema, createStmts map[string]string, retErr error) {
@@ -111,14 +111,14 @@ func GenerateDDLs(
 		}
 	}()
 
-	return generateDDLs(bufio.NewReader(f), dbName, anonymize)
+	return generateDDLFromReader(bufio.NewReader(f), dbName, anonymize)
 }
 
-// generateDDLs takes a reader for a TSV file containing DDL statements,
+// generateDDLFromReader takes a reader for a TSV file containing DDL statements,
 // parses the statements, and returns a map of table names to their schemas
 // and a map of short table names to their CREATE TABLE statements.
 // It has been deigned this way to maek it unit-testable
-func generateDDLs(
+func generateDDLFromReader(
 	r io.Reader, dbName string, anonymize bool,
 ) (map[string]*TableSchema, map[string]string, error) {
 	reader := csv.NewReader(r)
@@ -255,12 +255,13 @@ func splitColumnDefsAndTableConstraints(body string) (colDefs, tableConstraints
 func hasConstrainingPrefix(up string) bool {
 	// You could make this a global var if you like, to avoid reallocating the slice.
 	prefixes := []string{
-		"CONSTRAINT",
-		"PRIMARY KEY",
-		"UNIQUE",
-		"FOREIGN KEY",
-		"CHECK",
-		"INDEX",
+		sqlConstraint,
+		sqlPrimaryKey,
+		sqlUnique,
+		sqlForeignKey,
+		sqlCheck,
+		sqlIndex,
+		sqlFamily,
 	}
 	for _, p := range prefixes {
 		if strings.HasPrefix(up, p) {

pkg/workload/workload_generator/schema_generator_test.go

@@ -115,13 +115,13 @@ var data string
 
 func TestGenerateDDLsIntegration(t *testing.T) {
 	t.Run("expect success", func(t *testing.T) {
-		schemas, stmts, err := generateDDLs(strings.NewReader(data), testDBName, false)
+		schemas, stmts, err := generateDDLFromReader(strings.NewReader(data), testDBName, false)
 		assert.NoError(t, err)
 		assert.NotEmpty(t, schemas)
 		assert.NotEmpty(t, stmts)
 	})
 	t.Run("expect failure due to invalid file location", func(t *testing.T) {
-		_, _, err := GenerateDDLs("wrong_file_location", testDBName, false)
+		_, _, err := generateDDLs("wrong_file_location", testDBName, false)
 		assert.NotNil(t, err)
 	})
 }

pkg/workload/workload_generator/utils.go renamed to pkg/workload/workload_generator/schema_utils.go

@@ -3,7 +3,7 @@
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
-// utils_test.go
+// schema_utils_test.go
 package workload_generator
 
 import (