add full export

Author: githubesson

.gitignore

@@ -1 +1,2 @@
-/*.exe
+/*.exe
+/*.sql

README.md

@@ -38,6 +38,7 @@ sqlparser [-format=txt|csv|json|jsonl] [-output=filename] [-workers=N] <sqlfile>
   - `jsonl`: JSON lines format with table structure
 - `-output`: Output file path (default: stdout)
 - `-workers`: Number of worker threads (default: 1)
+- `-all`: Export all tables (default: false)
 - `<sqlfile>`: Input SQL file containing INSERT statements
 
 ### Environment Variables

cmd/sqlparser/main.go

@@ -12,66 +12,99 @@ import (
 )
 
 func main() {
-	format := flag.String("format", "txt", "Output format (txt, csv, json, jsonl)")
-	output := flag.String("output", "", "Output file (if not specified, prints to stdout)")
+	format := flag.String("format", "", "Output format (txt, csv, json, jsonl)")
+	output := flag.String("output", "", "Output file (for single table export)")
 	workers := flag.Int("workers", getWorkerCount(), "Number of worker threads")
+	exportAll := flag.Bool("all", false, "Export all tables (creates a directory named after the input file)")
 	flag.Parse()
 
 	args := flag.Args()
 	if len(args) < 1 {
-		fmt.Printf("Usage: sqlparser [-format=txt|csv|json] [-output=filename] [-workers=N] <sqlfile>\n")
-		fmt.Printf("  -format: Output format (default: txt)\n")
-		fmt.Printf("  -output: Output file (default: stdout)\n")
+		fmt.Printf("Usage: sqlparser [-format=txt|csv|json] [-output=filename] [-workers=N] [-all] <sqlfile>\n")
+		fmt.Printf("  -format: Output format (txt, csv, json, jsonl). If specified without -output, creates files in a directory\n")
+		fmt.Printf("  -output: Output file (optional, defaults to directory output if format is specified)\n")
 		fmt.Printf("  -workers: Number of worker threads (default: %d)\n", getWorkerCount())
+		fmt.Printf("  -all: Export all tables into separate files (default: false)\n")
 		os.Exit(1)
 	}
 
 	filename := args[0]
 
-	// Scan for tables and prompt for selection
+	// Scan for tables
 	tables, err := parser.ScanTables(filename)
 	if err != nil {
 		fmt.Printf("Error scanning tables: %v\n", err)
 		os.Exit(1)
 	}
 
-	selectedTable, err := parser.PromptTableSelection(tables)
-	if err != nil {
-		fmt.Printf("Error selecting table: %v\n", err)
-		os.Exit(1)
+	var selectedTables []*parser.TableInfo
+	if *exportAll {
+		selectedTables = make([]*parser.TableInfo, len(tables))
+		for i := range tables {
+			selectedTables[i] = &tables[i]
+		}
+		fmt.Printf("\nExporting all %d tables\n", len(tables))
+	} else {
+		selectedTable, err := parser.PromptTableSelection(tables)
+		if err != nil {
+			// Check if all tables were selected from the menu
+			if allSelected, ok := err.(*parser.AllTablesSelected); ok {
+				selectedTables = make([]*parser.TableInfo, len(allSelected.Tables))
+				for i := range allSelected.Tables {
+					selectedTables[i] = &allSelected.Tables[i]
+				}
+				fmt.Printf("\nExporting all %d tables\n", len(tables))
+			} else {
+				fmt.Printf("Error selecting table: %v\n", err)
+				os.Exit(1)
+			}
+		} else {
+			selectedTables = []*parser.TableInfo{selectedTable}
+			fmt.Printf("\nSelected table: %s\n", selectedTable.Name)
+		}
 	}
 
-	fmt.Printf("\nSelected table: %s\n", selectedTable.Name)
-
+	// If format is specified but no output, use directory output by default
+	useDirectoryOutput := *exportAll || len(selectedTables) > 1 || (*format != "" && *output == "")
 	outputFormat := models.OutputFormat(*format)
+	if *format == "" {
+		outputFormat = models.FormatText // default to text if no format specified
+	}
 
-	// Create output file or use stdout
-	var out *os.File
-	if *output != "" {
-		out, err = os.Create(*output)
-		if err != nil {
-			fmt.Printf("Error creating output file: %v\n", err)
-			os.Exit(1)
-		}
-		defer out.Close()
+	var w writer.Writer
+	if useDirectoryOutput {
+		w, err = writer.CreateMultiWriter(outputFormat, filename)
 	} else {
-		out = os.Stdout
+		if *output == "" {
+			w, err = writer.CreateWriter(outputFormat, os.Stdout)
+		} else {
+			out, err := os.Create(*output)
+			if err != nil {
+				fmt.Printf("Error creating output file: %v\n", err)
+				os.Exit(1)
+			}
+			defer out.Close()
+			w, err = writer.CreateWriter(outputFormat, out)
+			if err != nil {
+				fmt.Printf("Error creating writer: %v\n", err)
+				os.Exit(1)
+			}
+		}
 	}
 
-	// Initialize the output writer based on format
-	w, err := writer.CreateWriter(outputFormat, out)
 	if err != nil {
 		fmt.Printf("Error creating writer: %v\n", err)
 		os.Exit(1)
 	}
 	defer w.Close()
 
-	// Process the file
+	// Process each selected table
 	fmt.Printf("Processing with %d workers...\n", *workers)
-	err = parser.ProcessSQLFileInBatches(filename, w, *workers, selectedTable)
-	if err != nil {
-		fmt.Printf("Error processing SQL file: %v\n", err)
-		os.Exit(1)
+	for _, table := range selectedTables {
+		if err := parser.ProcessSQLFileInBatches(filename, w, *workers, table); err != nil {
+			fmt.Printf("Error processing table %s: %v\n", table.Name, err)
+			os.Exit(1)
+		}
 	}
 }
 

pkg/models/types.go

@@ -9,6 +9,13 @@ var (
 	BatchSize = getBatchSize()
 )
 
+const (
+	FormatText  OutputFormat = "txt"
+	FormatCSV   OutputFormat = "csv"
+	FormatJSON  OutputFormat = "json"
+	FormatJSONL OutputFormat = "jsonl"
+)
+
 func getBatchSize() int {
 	if val := os.Getenv("BATCH_SIZE"); val != "" {
 		if size, err := strconv.Atoi(val); err == nil && size > 0 {
@@ -20,12 +27,9 @@ func getBatchSize() int {
 
 type OutputFormat string
 
-const (
-	FormatText  OutputFormat = "txt"
-	FormatCSV   OutputFormat = "csv"
-	FormatJSON  OutputFormat = "json"
-	FormatJSONL OutputFormat = "jsonl"
-)
+func (f OutputFormat) Extension() string {
+	return string(f)
+}
 
 type Row struct {
 	TableName string                 `json:"table_name,omitempty"`

pkg/parser/database.go

@@ -78,16 +78,29 @@ func PromptTableSelection(tables []TableInfo) (*TableInfo, error) {
 	}
 
 	fmt.Println("\nFound the following tables with INSERT statements:")
+	fmt.Printf("0. Export all tables\n")
 	for i, table := range tables {
 		fmt.Printf("%d. %s\n", i+1, table.Name)
 	}
 
 	var choice int
-	fmt.Print("\nEnter the number of the table you want to parse (1-" + fmt.Sprint(len(tables)) + "): ")
+	fmt.Print("\nEnter the number of the table you want to parse (0-" + fmt.Sprint(len(tables)) + "): ")
 	_, err := fmt.Scanf("%d", &choice)
-	if err != nil || choice < 1 || choice > len(tables) {
+	if err != nil || choice < 0 || choice > len(tables) {
 		return nil, fmt.Errorf("invalid selection")
 	}
 
+	if choice == 0 {
+		return nil, &AllTablesSelected{Tables: tables}
+	}
+
 	return &tables[choice-1], nil
 }
+
+type AllTablesSelected struct {
+	Tables []TableInfo
+}
+
+func (e *AllTablesSelected) Error() string {
+	return "all tables selected"
+}

pkg/parser/parser.go

@@ -88,10 +88,6 @@ func ProcessSQLFileInBatches(filename string, writer writer.Writer, numWorkers i
 
 				// Process rows immediately
 				for _, row := range result.rows {
-					if writer.Type() == models.FormatJSONL {
-						row.TableName = currentTableName
-					}
-					
 					rowCount++
 					row.RowNumber = rowCount
 					currentBatch = append(currentBatch, row)
@@ -195,7 +191,15 @@ func processStatement(statement string, numWorkers int) (string, []models.Row, e
 	statement = strings.TrimSpace(strings.TrimSuffix(statement, ";"))
 
 	if strings.HasPrefix(strings.ToUpper(statement), "INSERT INTO") {
-		return parseInsert(statement, numWorkers)
+		tableName, rows, err := parseInsert(statement, numWorkers)
+		if err != nil {
+			return "", nil, err
+		}
+		// Set the table name for each row
+		for i := range rows {
+			rows[i].TableName = tableName
+		}
+		return tableName, rows, nil
 	}
 	return "", nil, nil
 }

pkg/parser/rows.go

@@ -29,7 +29,8 @@ func parseRowsSequential(tableName string, columns []string, values [][]string)
 			}
 		}
 		rows[i] = models.Row{
-			Data: rowData,
+			TableName: tableName,
+			Data:      rowData,
 		}
 	}
 	return tableName, rows, nil
@@ -70,7 +71,8 @@ func parseRowsParallel(tableName string, columns []string, values [][]string, nu
 					}
 				}
 				rows[idx] = models.Row{
-					Data: rowData,
+					TableName: tableName,
+					Data:      rowData,
 				}
 			}
 		}(start, end)

pkg/writer/writer.go

@@ -4,6 +4,8 @@ import (
 	"bufio"
 	"fmt"
 	"io"
+	"os"
+	"path/filepath"
 
 	"sqlparser/pkg/models"
 )
@@ -18,6 +20,12 @@ type Writer interface {
 	Type() models.OutputFormat
 }
 
+type MultiWriter struct {
+	format  models.OutputFormat
+	writers map[string]Writer
+	baseDir string
+}
+
 func CreateWriter(format models.OutputFormat, output io.Writer) (Writer, error) {
 	bufferedWriter := bufio.NewWriterSize(output, bufferSize)
 
@@ -34,3 +42,80 @@ func CreateWriter(format models.OutputFormat, output io.Writer) (Writer, error)
 		return nil, fmt.Errorf("unsupported format: %s", format)
 	}
 }
+
+func CreateMultiWriter(format models.OutputFormat, inputPath string) (*MultiWriter, error) {
+	// Use the input file name (without extension) as the base directory
+	baseDir := filepath.Base(inputPath)
+	if filepath.Ext(baseDir) != "" {
+		baseDir = baseDir[:len(baseDir)-len(filepath.Ext(baseDir))]
+	}
+
+	// Create the directory
+	if err := os.MkdirAll(baseDir, 0755); err != nil {
+		return nil, fmt.Errorf("failed to create output directory: %v", err)
+	}
+
+	return &MultiWriter{
+		format:  format,
+		baseDir: baseDir,
+		writers: make(map[string]Writer),
+	}, nil
+}
+
+func (mw *MultiWriter) WriteTableStart(tableName string) error {
+	// Create a new file for this table
+	filename := filepath.Join(mw.baseDir, tableName+"."+mw.format.Extension())
+	file, err := os.Create(filename)
+	if err != nil {
+		return fmt.Errorf("failed to create file for table %s: %v", tableName, err)
+	}
+
+	// Create a writer for this table
+	writer, err := CreateWriter(mw.format, file)
+	if err != nil {
+		file.Close()
+		return err
+	}
+
+	mw.writers[tableName] = writer
+	return writer.WriteTableStart(tableName)
+}
+
+func (mw *MultiWriter) WriteRows(rows []models.Row) error {
+	if len(rows) == 0 {
+		return nil
+	}
+
+	// Get the writer for this table's rows
+	tableName := rows[0].TableName
+	writer, exists := mw.writers[tableName]
+	if !exists {
+		return fmt.Errorf("no writer found for table %s", tableName)
+	}
+
+	return writer.WriteRows(rows)
+}
+
+func (mw *MultiWriter) WriteTableEnd() error {
+	var lastErr error
+	for tableName, writer := range mw.writers {
+		if err := writer.WriteTableEnd(); err != nil {
+			lastErr = fmt.Errorf("failed to end table %s: %v", tableName, err)
+		}
+	}
+	return lastErr
+}
+
+func (mw *MultiWriter) Close() error {
+	var lastErr error
+	for tableName, writer := range mw.writers {
+		if err := writer.Close(); err != nil {
+			lastErr = fmt.Errorf("failed to close writer for table %s: %v", tableName, err)
+		}
+	}
+	return lastErr
+}
+
+func (mw *MultiWriter) Type() models.OutputFormat {
+	return mw.format
+}