Merge #149394

149394: pkg/cli: Enhance TSV parsing in zip_upload_table_dumps r=Abhinav1299 a=Abhinav1299

This commit introduces a new function, makeQuotedTSVIterator, to handle TSV files with quoted fields containing newlines. This is used to correctly parse complex fields such as SQL statements that spans multiple lines. It also updates the existing makeTableIterator to correctly manage empty files.

The processTableDump function is modified to utilize the new parser for specific files, improving robustness in handling complex TSV formats.

Part of: CRDB-51738
Epic: CRDB-52093
Release note: None

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

Author: craig[bot]

pkg/cli/zip_upload_table_dumps.go

@@ -27,6 +27,10 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
+const (
+	createStatementsFileName = "crdb_internal.create_statements.txt"
+)
+
 type tsvColumnParserFn func(string) (any, error)
 
 type columnParserMap map[string]tsvColumnParserFn
@@ -55,7 +59,7 @@ var clusterWideTableDumps = map[string]columnParserMap{
 	"system.rangelog.txt":                           {},
 	"crdb_internal.table_indexes.txt":               {},
 	"crdb_internal.index_usage_statistics.txt":      {},
-	"crdb_internal.create_statements.txt":           {},
+	createStatementsFileName:                        {},
 	"system.job_info.txt":                           {},
 	"crdb_internal.create_schema_statements.txt":    {},
 	"crdb_internal.default_privileges.txt":          {},
@@ -253,11 +257,13 @@ func processTableDump(
 		tags = append(tags, makeDDTag(nodeIDTag, strings.Split(fileName, "/")[1]))
 	}
 
-	header, iter := makeTableIterator(f)
-	if err := iter(func(row string) error {
-		cols := strings.Split(row, "\t")
+	// Common processing function for all the parsers
+	processFields := func(header []string, cols []string) error {
 		if len(header) != len(cols) {
-			return errors.Newf("the number of headers is not matching the number of columns in the row")
+			return errors.Newf(
+				"the number of headers (%d) is not matching the number of columns (%d) in the row",
+				len(header), len(cols),
+			)
 		}
 
 		headerColumnMapping := map[string]any{
@@ -302,8 +308,23 @@ func processTableDump(
 
 		lines = append(lines, jsonRow)
 		return nil
-	}); err != nil {
-		return err
+	}
+
+	var processErr error
+
+	iterMaker := makeTableIterator // default parser
+	if strings.HasSuffix(fileName, createStatementsFileName) {
+		// Using makeQuotedTSVIterator parser for create_statements.txt because it has
+		// SQL CREATE statements with embedded newlines and tabs
+		iterMaker = makeQuotedTSVIterator
+	}
+	header, iter := iterMaker(f)
+	processErr = iter(func(cols []string) error {
+		return processFields(header, cols)
+	})
+
+	if processErr != nil {
+		return processErr
 	}
 
 	// flush the remaining lines if any
@@ -316,28 +337,37 @@ func processTableDump(
 	return nil
 }
 
-// makeTableIterator returns the headers slice and an iterator
-func makeTableIterator(f io.Reader) ([]string, func(func(string) error) error) {
+// makeTableIterator returns the headers slice and an iterator (original implementation for most files)
+func makeTableIterator(f io.Reader) ([]string, func(func([]string) error) error) {
 	reader := bufio.NewReader(f)
 
 	// Read first line for headers
 	headerLine, err := reader.ReadString('\n')
 	if err != nil && err != io.EOF {
-		return nil, func(func(string) error) error { return err }
+		return nil, func(func([]string) error) error { return err }
 	}
 
 	// Trim the newline character if present
 	headerLine = strings.TrimSuffix(headerLine, "\n")
-	headers := strings.Split(headerLine, "\t")
 
-	return headers, func(fn func(string) error) error {
+	// Handle empty files correctly
+	var headers []string
+	if headerLine == "" {
+		headers = []string{}
+	} else {
+		headers = strings.Split(headerLine, "\t")
+	}
+
+	return headers, func(fn func([]string) error) error {
 		for {
 			line, err := reader.ReadString('\n')
 			if err != nil {
 				if err == io.EOF {
 					// Process any remaining content before EOF
 					if line != "" {
-						if err := fn(strings.TrimSuffix(line, "\n")); err != nil {
+						line = strings.TrimSuffix(line, "\n")
+						cols := strings.Split(line, "\t")
+						if err := fn(cols); err != nil {
 							return err
 						}
 					}
@@ -346,16 +376,123 @@ func makeTableIterator(f io.Reader) ([]string, func(func(string) error) error) {
 				return err
 			}
 
-			// Trim the newline character
+			// Trim the newline character and split into fields
 			line = strings.TrimSuffix(line, "\n")
-			if err := fn(line); err != nil {
+			cols := strings.Split(line, "\t")
+			if err := fn(cols); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+}
+
+// makeQuotedTSVIterator returns the headers slice and an iterator that properly handles
+// TSV files with quoted fields containing newlines. This function is ONLY used for
+// special cases where TSV fields contain complex content that spans multiple lines,
+// such as SQL CREATE statements with embedded newlines and tabs.
+//
+// When to use:
+//   - Use this function when TSV fields are quoted and may contain newlines (e.g., crdb_internal.create_statements.txt)
+//   - DO NOT use for regular TSV files - use makeTableIterator instead
+//
+// Example problem this solves:
+//
+//	A TSV field containing: "CREATE TABLE foo (\n  id INT,\n  name STRING\n)"
+//	Simple line-by-line parsing would incorrectly split this into multiple records,
+//	but this function correctly treats it as a single field value.
+//
+// Unlike makeTableIterator, this function uses readTSVRecord to correctly parse
+// complex fields and returns parsed field slices for each record.
+func makeQuotedTSVIterator(f io.Reader) ([]string, func(func([]string) error) error) {
+	reader := bufio.NewReader(f)
+
+	_, headers, err := readTSVRecord(reader)
+	if err != nil {
+		if err == io.EOF {
+			return []string{}, func(func([]string) error) error { return nil }
+		}
+		return nil, func(func([]string) error) error { return err }
+	}
+
+	if len(headers) == 0 {
+		return headers, func(func([]string) error) error { return nil }
+	}
+
+	return headers, func(fn func([]string) error) error {
+		for {
+			_, fields, err := readTSVRecord(reader)
+			if err != nil {
+				if err == io.EOF {
+					break
+				}
+				return err
+			}
+
+			if len(fields) == 0 {
+				break
+			}
+
+			if err := fn(fields); err != nil {
 				return err
 			}
 		}
 		return nil
 	}
 }
 
+// readTSVRecord reads a complete TSV record from a buffered reader, properly handling
+// quoted fields that may contain embedded newlines and tab characters. Unlike simple
+// line-by-line parsing, this function tracks quote state to correctly parse fields
+// that span multiple lines. Returns the complete raw line, parsed field values as a
+// slice, and any error encountered during reading.
+func readTSVRecord(reader *bufio.Reader) (string, []string, error) {
+	var line strings.Builder
+	var fields []string
+	var currentField strings.Builder
+	inQuotes := false
+	hasContent := false
+
+	for {
+		b, err := reader.ReadByte()
+		if err != nil {
+			if err == io.EOF {
+				if hasContent && (line.Len() > 0 || currentField.Len() > 0) {
+					fields = append(fields, currentField.String())
+					return line.String(), fields, nil
+				}
+				return "", fields, io.EOF
+			}
+			return "", nil, err
+		}
+
+		hasContent = true
+		line.WriteByte(b)
+
+		switch b {
+		case '"':
+			inQuotes = !inQuotes
+			currentField.WriteByte(b)
+		case '\t':
+			if inQuotes {
+				currentField.WriteByte(b)
+			} else {
+				fields = append(fields, currentField.String())
+				currentField.Reset()
+			}
+		case '\n':
+			if inQuotes {
+				currentField.WriteByte(b)
+			} else {
+				fields = append(fields, currentField.String())
+				return line.String(), fields, nil
+			}
+		default:
+			currentField.WriteByte(b)
+		}
+	}
+}
+
 func getNodeSpecificTableDumps(debugDirPath string) ([]string, error) {
 	allTxtFiles, err := expandPatterns([]string{path.Join(debugDirPath, zippedNodeTableDumpsPattern)})
 	if err != nil {

pkg/cli/zip_upload_table_dumps_test.go

@@ -100,14 +100,65 @@ func TestMakeTableIterator(t *testing.T) {
 			)),
 			headers: []string{"h1"},
 		},
+		{
+			name:    "empty file",
+			input:   bytes.NewBufferString(""),
+			headers: []string{},
+		},
+		{
+			name:    "header only",
+			input:   bytes.NewBufferString("h1\th2\th3\n"),
+			headers: []string{"h1", "h2", "h3"},
+		},
 	}
 
 	for _, tc := range tt {
-		headers, iter := makeTableIterator(tc.input)
-		assert.Equal(t, tc.headers, headers)
-		require.NoError(t, iter(func(row string) error {
-			assert.Len(t, strings.Split(row, "\t"), len(headers))
-			return nil
-		}))
+		t.Run("Old_"+tc.name, func(t *testing.T) {
+			headers, iter := makeTableIterator(tc.input)
+			assert.Equal(t, tc.headers, headers)
+			require.NoError(t, iter(func(fields []string) error {
+				assert.Len(t, fields, len(headers))
+				return nil
+			}))
+		})
+	}
+
+	// Test cases specifically for the quoted TSV parser with quoted fields
+	quotedTSVTestCases := []struct {
+		name    string
+		input   io.Reader
+		headers []string
+	}{
+		{
+			name:    "simple",
+			input:   bytes.NewBufferString("h1\th2\th3\nr1c1\tr1c2\tr1c3\nr2c1\tr2c2\tr2c3\n"),
+			headers: []string{"h1", "h2", "h3"},
+		},
+		{
+			name:    "quoted fields with newlines",
+			input:   bytes.NewBufferString("col1\tcol2\tcol3\nval1\t\"CREATE TABLE test (\n\tid INT,\n\tname STRING\n)\"\tval3\n"),
+			headers: []string{"col1", "col2", "col3"},
+		},
+		{
+			name:    "empty file",
+			input:   bytes.NewBufferString(""),
+			headers: []string{},
+		},
+		{
+			name:    "header only",
+			input:   bytes.NewBufferString("h1\th2\th3\n"),
+			headers: []string{"h1", "h2", "h3"},
+		},
+	}
+
+	for _, tc := range quotedTSVTestCases {
+		t.Run("QuotedTSV_"+tc.name, func(t *testing.T) {
+			headers, iter := makeQuotedTSVIterator(tc.input)
+			assert.Equal(t, tc.headers, headers)
+			require.NoError(t, iter(func(fields []string) error {
+				assert.Len(t, fields, len(headers))
+				return nil
+			}))
+		})
 	}
 }