Merge #153521

153521: workload_generator: add support for external DDL r=shailendra-patel a=nameisbhaskar

Currently, the DDL generator relies on the presence of DDL CREATE statements within the debug zip. However, this dependency breaks when the debug zip is captured for a time range that does not include the creation of those objects.

To address this limitation, this PR introduces support for supplying a schema dump as an alternative input source. Users can generate this dump by running:
```
cockroach sql --url='postgresql://<url>/<db name>' --execute="SHOW CREATE ALL TABLES;" > ddl_file.sql
```
The resulting file can then be provided to the generator via the `--ddl-file` flag.

Epic: None
Release note: None

Co-authored-by: Bhaskarjyoti Bora <[email protected]>

Author: craig[bot]

pkg/workload/workload_generator/schema_generator.go

@@ -98,9 +98,22 @@ var (
 // TODO: The "anonymize" parameter is unused for now.
 func generateDDLs(
 	zipDir,
-	dbName string, anonymize bool,
+	dbName, ddlFile string, anonymize bool,
 ) (allSchemas map[string]*TableSchema, createStmts map[string]string, retErr error) {
 
+	if ddlFile != "" {
+		// DDL file location is present. We will use this instead of the debug zip.
+		f, err := os.Open(ddlFile)
+		if err != nil {
+			return nil, nil, errors.Wrap(err, "failed to open DDL file")
+		}
+		defer func() {
+			if cerr := f.Close(); cerr != nil && retErr == nil {
+				retErr = errors.Wrap(cerr, "failed to close input DDL file")
+			}
+		}()
+		return generateDDLFromDDLFile(bufio.NewReader(f), dbName, anonymize)
+	}
 	f, err := openCreateStatementsTSV(zipDir)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "failed to open TSV file")
@@ -111,14 +124,158 @@ func generateDDLs(
 		}
 	}()
 
-	return generateDDLFromReader(bufio.NewReader(f), dbName, anonymize)
+	return generateDDLFromCSV(bufio.NewReader(f), dbName, anonymize)
+}
+
+// generateDDLFromDDLFile reads DDL statements from a SQL dump file
+// and returns a map of table names to their schemas and a map of
+// short table names to their CREATE TABLE statements.
+// The file can be generated by running the following:
+//
+//	cockroach sql --url='postgresql://<url>/<db name>' --execute="SHOW CREATE ALL TABLES;" > ddl_file.sql
+func generateDDLFromDDLFile(
+	reader *bufio.Reader, dbName string, anonymize bool,
+) (map[string]*TableSchema, map[string]string, error) {
+	// the results are stored in these Maps.
+	tableStatements := make(map[string]string)
+	order := make([]string, 0)
+	seen := make(map[string]struct{})
+
+	// Buffer accumulates the SQL statements
+	var currentStmt strings.Builder
+	// inStatement helps handling multi line statements
+	inStatement := false
+
+	// The file is read line by line
+	for {
+		line, err := reader.ReadString('\n')
+		if err != nil {
+			if err.Error() == "EOF" {
+				break
+			}
+			return nil, nil, errors.Wrap(err, "failed while reading SQL file")
+		}
+
+		// Empty lines and comments are skipped
+		trimmedLine := strings.TrimSpace(line)
+		if !inStatement {
+			// The generated statement has a quote at the start of the statement. This is trimmed.
+			trimmedLine = strings.TrimLeft(trimmedLine, "\"")
+		}
+		if trimmedLine == "" || strings.HasPrefix(trimmedLine, "--") ||
+			strings.HasPrefix(trimmedLine, "create_statement") {
+			continue
+		}
+
+		// A new statement is expected to start with CREATE TABLE.
+		if strings.HasPrefix(strings.ToUpper(trimmedLine), "CREATE TABLE") {
+			// If we were already in a statement, the previous statement is processed
+			if inStatement {
+				tableStatements, order, seen = processStatement(currentStmt.String(), tableStatements, order, seen, dbName)
+			}
+
+			// A new statement is started.
+			currentStmt.Reset()
+			currentStmt.WriteString(trimmedLine)
+			inStatement = true
+		} else if strings.HasPrefix(strings.ToUpper(trimmedLine), "ALTER TABLE") {
+			// If we were in a CREATE TABLE statement, the statement is processed
+			if inStatement {
+				tableStatements, order, seen = processStatement(currentStmt.String(), tableStatements, order, seen, dbName)
+			}
+
+			// A new ALTER TABLE statement is started.
+			currentStmt.Reset()
+			currentStmt.WriteString(trimmedLine)
+			inStatement = true
+		} else if inStatement {
+			if strings.HasSuffix(trimmedLine, ";\"") {
+				// The generated statement has a quote at the end of the statement. This needs to be trimmed.
+				trimmedLine = strings.TrimRight(trimmedLine, "\"")
+			}
+			// The current statement is accumulated.
+			currentStmt.WriteString(trimmedLine)
+
+			// if the statement is complete (ends with semicolon or has closing parenthesis followed by options), it is processed.
+			if strings.HasSuffix(trimmedLine, ";") ||
+				(strings.Contains(trimmedLine, ");") && !strings.HasPrefix(trimmedLine, "--")) {
+				tableStatements, order, seen = processStatement(currentStmt.String(), tableStatements, order, seen, dbName)
+				inStatement = false
+			}
+		}
+	}
+
+	// Any remaining statement is processed.
+	if inStatement {
+		tableStatements, order, _ = processStatement(currentStmt.String(), tableStatements, order, seen, dbName)
+	}
+
+	return buildSchemas(order, tableStatements), buildCreateStmts(tableStatements), nil
+}
+
+// processStatement processes a single SQL statement and adds it to the tableStatements map if it's a CREATE TABLE statement
+// It returns the updated tableStatements, order, and seen maps
+func processStatement(
+	stmt string,
+	tableStatements map[string]string,
+	order []string,
+	seen map[string]struct{},
+	dbName string,
+) (map[string]string, []string, map[string]struct{}) {
+	// Only process CREATE TABLE statements
+	if !strings.HasPrefix(strings.ToUpper(strings.TrimSpace(stmt)), "CREATE TABLE") {
+		return tableStatements, order, seen
+	}
+
+	// Extract the table name using the tablePattern regex
+	tableMatch := tablePattern.FindStringSubmatch(stmt)
+	if tableMatch == nil {
+		return tableStatements, order, seen
+	}
+
+	// Extract and normalize the table name
+	tableName := tableMatch[1]
+	parts := strings.Split(tableName, ".")
+	for i := range parts {
+		parts[i] = strings.Trim(parts[i], `"`) // Remove quotes from parts
+	}
+
+	// If the table name doesn't have a schema, assume it's "public"
+	var schemaName string
+	var simpleTableName string
+
+	if len(parts) == 1 {
+		schemaName = "public"
+		simpleTableName = parts[0]
+	} else if len(parts) == 2 {
+		schemaName = parts[0]
+		simpleTableName = parts[1]
+	} else {
+		// Skip tables with more complex names
+		return tableStatements, order, seen
+	}
+
+	// Create a regex for the schema name
+	schemaPattern := regexp.MustCompile(`\b` + regexp.QuoteMeta(schemaName) + `\.`)
+
+	// Process the DDL record
+	fullTableName, statement := processDDLRecord(dbName, schemaName, simpleTableName, stmt, schemaPattern)
+
+	// Add to the maps if not seen before
+	if _, ok := seen[fullTableName]; !ok && fullTableName != "" {
+		tableStatements[fullTableName] = statement
+		order = append(order, fullTableName)
+		seen[fullTableName] = struct{}{}
+	}
+
+	return tableStatements, order, seen
 }
 
-// generateDDLFromReader takes a reader for a TSV file containing DDL statements,
+// generateDDLFromCSV 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 generateDDLFromReader(
+func generateDDLFromCSV(
 	r io.Reader, dbName string, anonymize bool,
 ) (map[string]*TableSchema, map[string]string, error) {
 	reader := csv.NewReader(r)

pkg/workload/workload_generator/schema_generator_test.go

@@ -8,6 +8,7 @@
 package workload_generator
 
 import (
+	"bufio"
 	_ "embed"
 	"strings"
 	"testing"
@@ -115,13 +116,57 @@ var data string
 
 func TestGenerateDDLsIntegration(t *testing.T) {
 	t.Run("expect success", func(t *testing.T) {
-		schemas, stmts, err := generateDDLFromReader(strings.NewReader(data), testDBName, false)
+		schemas, stmts, err := generateDDLFromCSV(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)
 	})
 }
+
+func TestGenerateDDLFromDDLFile(t *testing.T) {
+	// Sample DDL statements for testing
+	sampleDDL := `"CREATE TABLE public.users (
+		id INT8 NOT NULL,
+		name VARCHAR(50) NOT NULL,
+		email VARCHAR(100) UNIQUE,
+		CONSTRAINT users_pkey PRIMARY KEY (id)
+	);"
+
+	CREATE TABLE public.orders (
+		order_id INT8 NOT NULL,
+		user_id INT8 NOT NULL,
+		amount DECIMAL(10,2) NOT NULL,
+		CONSTRAINT orders_pkey PRIMARY KEY (order_id)
+	);
+
+	ALTER TABLE public.orders ADD CONSTRAINT orders_user_id_fkey FOREIGN KEY (user_id) REFERENCES public.users(id);
+	`
+
+	t.Run("expect success with DDL file", func(t *testing.T) {
+		reader := strings.NewReader(sampleDDL)
+		schemas, stmts, err := generateDDLFromDDLFile(bufio.NewReader(reader), testDBName, false)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, schemas)
+		assert.NotEmpty(t, stmts)
+
+		// Verify that both tables were parsed correctly
+		assert.Contains(t, schemas, "users")
+		assert.Contains(t, schemas, "orders")
+
+		// Verify that the users table has the correct columns
+		usersSchema := schemas["users"]
+		assert.Contains(t, usersSchema.Columns, "id")
+		assert.Contains(t, usersSchema.Columns, "name")
+		assert.Contains(t, usersSchema.Columns, "email")
+
+		// Verify that the orders table has the correct columns
+		ordersSchema := schemas["orders"]
+		assert.Contains(t, ordersSchema.Columns, "order_id")
+		assert.Contains(t, ordersSchema.Columns, "user_id")
+		assert.Contains(t, ordersSchema.Columns, "amount")
+	})
+}

pkg/workload/workload_generator/workload.go

@@ -104,6 +104,10 @@ var workloadGeneratorMeta = workload.Meta{
 		g.flags.FlagSet = pflag.NewFlagSet("db_workload", pflag.ContinueOnError)
 		g.flags.StringVar(&g.debugLogsLocation, "debug-logs", "",
 			"Path to unzipped debug logs directory.")
+		g.flags.StringVar(&g.ddlFile, "ddl-file", "",
+			"The File containing the DDL. If this is provided, the DDL is read from this file instead of the debug logs."+
+				" To generate the file, execute: \n"+
+				"cockroach sql --url='postgresql://<url>/<db name>' --execute=\"SHOW CREATE ALL TABLES;\" > ddl_file.sql")
 		g.flags.IntVar(&g.rowCount, "rows", 1000,
 			"Base row count for tables without foreign keys; other tables scale by foreign key depth/fanout.")
 		g.flags.StringVar(&g.inputYAML, "input-yaml", "",
@@ -126,6 +130,7 @@ type workloadGenerator struct {
 	flags             workload.Flags
 	connFlags         *workload.ConnFlags
 	debugLogsLocation string // path to the unzipped debug zip file
+	ddlFile           string // path to the DDL file.
 	dbName            string // database name to use for the workload
 	rowCount          int    // base number of rows per table before FK‐depth scaling
 
@@ -200,7 +205,7 @@ func (w *workloadGenerator) initializeGenerator() error {
 	w.dbName = dbName
 
 	// 2) Parsing DDLs out of the debug logs.
-	schemas, stmts, err := generateDDLs(w.debugLogsLocation, w.dbName, false)
+	schemas, stmts, err := generateDDLs(w.debugLogsLocation, w.dbName, w.ddlFile, false)
 	if err != nil {
 		return errors.Wrap(err, "failed to generate DDLs from debug logs")
 	}
@@ -592,7 +597,6 @@ func (w *workloadGenerator) initGenerators(db *gosql.DB) error {
 			maxRows = tblBlock.Count
 		}
 	}
-	// TODO: make the globalBatchNumber dynamic
 	globalNumBatches := (maxRows + baseBatchSize - 1) / baseBatchSize
 
 	// 1) The generator + empty cache for every table.col is built.