Finalized add function

commands/add.go

@@ -11,6 +11,8 @@ import (
 	"regexp"
 	"strconv"
 	"time"
+	"bytes"
+	"bufio"
 )
 
 type Field struct {
@@ -67,14 +69,6 @@ func Add(tableName string, model string) {
 		tableModel = promptForTableModel()
 	}
 
-
-// TODO: Do Stuff:
-	/*
-		- Generate Test Data for the new table if wanted
-		- Add Description and mermaid and DBML to the README.md file
-    - Fix prompting to use SelectMenu
-	*/
-
 	slug := utils.Slugify(tableName)
 
 	// Update datasmith.yaml config and save the new table definition
@@ -122,6 +116,15 @@ func Add(tableName string, model string) {
 		fmt.Printf("Updated CHANGELOG with version '%s'\n", newVersion)
 	}
 
+	// Update README.md with the new table information
+	if err := updateReadme(projectDir, tableName, tableModel); err != nil {
+		fmt.Printf("Error updating README.md: %v\n", err)
+	} else {
+		fmt.Printf("Updated README.md with table '%s'\n", tableName)
+	}
+	fmt.Print("\n")
+	fmt.Print("------------")
+	fmt.Print("\n")
 	fmt.Printf("Added new table '%s' to the project\n", slug)
 }
 
@@ -272,6 +275,115 @@ func generatePostgreSQLSQL(file *os.File, tableName string, tableModel TableMode
 	return err
 }
 
+func updateReadme(projectDir, tableName string, tableModel TableModel) error {
+	readmeFilePath := filepath.Join(projectDir, "README.md")
+
+	// Read the existing README.md content
+	file, err := os.Open(readmeFilePath)
+	if err != nil {
+		return fmt.Errorf("error reading README.md file: %v", err)
+	}
+	defer file.Close()
+
+	var buffer bytes.Buffer
+	scanner := bufio.NewScanner(file)
+	inTablesSection := false
+	inMermaidSection := false
+	inDbmlSection := false
+	var tablesSection []string
+	var mermaidSection []string
+	var dbmlSection []string
+
+	for scanner.Scan() {
+		line := scanner.Text()
+		trimmedLine := strings.TrimSpace(line)
+
+		if strings.HasPrefix(trimmedLine, "```mermaid") {
+			inMermaidSection = true
+			buffer.WriteString(line + "\n")
+			continue
+		}
+
+		if strings.HasPrefix(trimmedLine, "```dbml") {
+			inDbmlSection = true
+			buffer.WriteString(line + "\n")
+			continue
+		}
+
+		if inMermaidSection && strings.HasPrefix(trimmedLine, "```") {
+			inMermaidSection = false
+			mermaidSection = append(mermaidSection, generateMermaidDiagram(tableName, tableModel))
+			buffer.WriteString(strings.Join(mermaidSection, "\n") + "\n")
+			buffer.WriteString(line + "\n")
+			continue
+		}
+
+		if inDbmlSection && strings.HasPrefix(trimmedLine, "```") {
+			inDbmlSection = false
+			dbmlSection = append(dbmlSection, generateDBML(tableName, tableModel))
+			buffer.WriteString(strings.Join(dbmlSection, "\n") + "\n")
+			buffer.WriteString(line + "\n")
+			continue
+		}
+
+		if strings.Contains(trimmedLine, "NO TABLES GIVEN") {
+			inTablesSection = true
+			buffer.WriteString(strings.Replace(line, "NO TABLES GIVEN", "Tables:\n", 1))
+			continue
+		}
+
+		if inTablesSection && strings.HasPrefix(trimmedLine, "```") {
+			inTablesSection = false
+			buffer.WriteString(strings.Join(tablesSection, "\n") + "\n")
+			buffer.WriteString(line + "\n")
+			continue
+		}
+
+		if inTablesSection {
+			tablesSection = append(tablesSection, "- "+tableName)
+		}
+
+		buffer.WriteString(line + "\n")
+	}
+
+	// Append table if it was not in any sections
+	if inTablesSection {
+		buffer.WriteString(strings.Join(tablesSection, "\n") + "\n")
+	}
+	if inMermaidSection {
+		mermaidSection = append(mermaidSection, generateMermaidDiagram(tableName, tableModel))
+		buffer.WriteString(strings.Join(mermaidSection, "\n") + "\n")
+	}
+	if inDbmlSection {
+		dbmlSection = append(dbmlSection, generateDBML(tableName, tableModel))
+		buffer.WriteString(strings.Join(dbmlSection, "\n") + "\n")
+	}
+
+	if err := scanner.Err(); err != nil {
+		return fmt.Errorf("error scanning README.md file: %v", err)
+	}
+
+	// Write the updated content back to the README.md file
+	if err := os.WriteFile(readmeFilePath, buffer.Bytes(), 0644); err != nil {
+		return fmt.Errorf("error writing updated README.md file: %v", err)
+	}
+
+	return nil
+}
+
+func generateMermaidDiagram(tableName string, tableModel TableModel) string {
+	var mermaid strings.Builder
+
+	mermaid.WriteString(fmt.Sprintf("    %s {\n", tableName))
+
+	for _, field := range tableModel.Fields {
+		mermaid.WriteString(fmt.Sprintf("        %s %s\n", field.Name, field.Type))
+	}
+
+	mermaid.WriteString("    }\n")
+
+	return mermaid.String()
+}
 func generateDBML(tableName string, tableModel TableModel) string {
 	var dbml strings.Builder
 
@@ -439,18 +551,18 @@ func promptForTableModel() TableModel {
 		utils.PromptIfEmpty(&field.Name, "Enter field name: ")
 		field.Type = promptForFieldType()
 
-		field.PrimaryKey = utils.PromptForBool("Is this a primary key? (y/n): ") // TODO: use SelectMenu
+		field.PrimaryKey = utils.PromptForBool("Is this a primary key? (y/n): ") 
 		if !field.PrimaryKey {
 			utils.PromptIfEmpty(&field.ForeignKey, "Enter foreign key (leave empty if not applicable): ")
 		}
-		field.Unique = utils.PromptForBool("Is this field unique? (y/n): ") // TODO: use SelectMenu
-		field.NotNull = utils.PromptForBool("Is this field not null? (y/n): ") // TODO: use SelectMenu
-		field.AutoIncrement = utils.PromptForBool("Is this field auto-increment? (y/n): ") // TODO: use SelectMenu
+		field.Unique = utils.PromptForBool("Is this field unique? (y/n): ") 
+		field.NotNull = utils.PromptForBool("Is this field not null? (y/n): ") 
+		field.AutoIncrement = utils.PromptForBool("Is this field auto-increment? (y/n): ")
 		field.DefaultValue = utils.PromptForString("Enter default value (leave empty if not applicable): ")
 
 		fields = append(fields, field)
 
-		if !utils.PromptForBool("Do you want to add another field? (y/n): ") { // TODO: use SelectMenu
+		if !utils.PromptForBool("Do you want to add another field? (y/n): ") {
 			break
 		}
 	}

utils/utils.go

@@ -46,20 +46,18 @@ func PromptIfEmpty(value *string, prompt string) {
 
 // PromptForBool prompts the user for a yes/no response and returns a boolean
 func PromptForBool(prompt string) bool {
-	for {
-		fmt.Print(prompt)
-		reader := bufio.NewReader(os.Stdin)
-		input, _ := reader.ReadString('\n')
-		input = strings.ToLower(strings.TrimSpace(input))
-
-		if input == "y" || input == "yes" {
-			return true
-		} else if input == "n" || input == "no" {
-			return false
-		} else {
-			fmt.Println("Invalid input, please enter 'y' or 'n'")
-		}
+	options := []MenuOption{
+		{Display: "Yes", Value: "yes"},
+		{Display: "No", Value: "no"},
 	}
+
+	selected, err := SelectMenu(options, prompt)
+	if err != nil {
+		fmt.Printf("Error selecting option: %v\n", err)
+		return false
+	}
+
+	return selected == "yes"
 }
 
 // PromptForString prompts the user for a string response and returns it