workload_generator: add support for external DDL

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

Author: nameisbhaskar

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.