From 2bde65c5a2a8a39b5617769f8ec4ffb9a4726e74 Mon Sep 17 00:00:00 2001
From: Fluffy-Bean <michal-gdula@protonmail.com>
Date: Mon, 6 May 2024 19:30:23 +0100
Subject: [PATCH] Improve migrate command

---
 cmd/database.go                               |  3 +-
 .../migrations/20240506163617-init_tables.sql | 14 ++--
 database/migrator.go                          | 67 ++++++++++++++++---
 3 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/cmd/database.go b/cmd/database.go
index 83ea9c0..d5784c5 100644
--- a/cmd/database.go
+++ b/cmd/database.go
@@ -12,6 +12,7 @@ func migrate(flags []string) {
 	cmd := flag.NewFlagSet("migrate", flag.ExitOnError)
 
 	downgrade := cmd.Bool("downgrade", false, "Downgrade Database")
+	confirm := cmd.Bool("y", false, "Skip questioning")
 
 	err := cmd.Parse(flags)
 	if err != nil {
@@ -19,7 +20,7 @@ func migrate(flags []string) {
 		os.Exit(1)
 	}
 
-	err = db.Migrate(!*downgrade)
+	err = db.Migrate(*downgrade, *confirm)
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)
diff --git a/database/migrations/20240506163617-init_tables.sql b/database/migrations/20240506163617-init_tables.sql
index 35b5d7b..c47c76e 100644
--- a/database/migrations/20240506163617-init_tables.sql
+++ b/database/migrations/20240506163617-init_tables.sql
@@ -1,10 +1,16 @@
 
 -- +migrate Up
 CREATE TABLE IF NOT EXISTS Item (
-       uuid TEXT NOT NULL PRIMARY KEY,
-       name TEXT NOT NULL,
-       price INTEGER NOT NULL,
-       description TEXT
+    item_uuid TEXT NOT NULL PRIMARY KEY,
+    name TEXT NOT NULL,
+    price INTEGER NOT NULL,
+    description TEXT
+);
+
+CREATE TABLE IF NOT EXISTS Images (
+    image_uuid TEXT NOT NULL PRIMARY KEY,
+    show BOOLEAN NOT NULL,
+    item_uuid TEXT NOT NULL
 );
 
 -- +migrate Down
diff --git a/database/migrator.go b/database/migrator.go
index 547cb6b..043a46f 100644
--- a/database/migrator.go
+++ b/database/migrator.go
@@ -1,8 +1,11 @@
 package database
 
 import (
+	"bufio"
 	"embed"
 	"fmt"
+	"os"
+	"strings"
 
 	migrate "github.com/rubenv/sql-migrate"
 )
@@ -14,12 +17,50 @@ var migrations = migrate.EmbedFileSystemMigrationSource{
 	Root:       "migrations",
 }
 
-func Migrate(Up bool) error {
-	var direction migrate.MigrationDirection
-	if Up {
+func Migrate(downgrade, confirm bool) error {
+	var direction = migrate.Down
+	if !downgrade {
 		direction = migrate.Up
+	}
+
+	var pending int
+	if !downgrade {
+		n, _, err := migrate.PlanMigration(Conn, "sqlite3", migrations, migrate.Up, 0)
+		if err != nil {
+			return err
+		}
+		pending = len(n)
 	} else {
-		direction = migrate.Down
+		n, err := migrate.GetMigrationRecords(Conn, "sqlite3")
+		if err != nil {
+			return err
+		}
+		pending = len(n)
+	}
+
+	if pending == 0 {
+		fmt.Println("Nothing to change")
+		os.Exit(0)
+	}
+
+	if !confirm {
+		if downgrade {
+			fmt.Printf("Downgrade %d migrations? [Y/n] ", pending)
+		} else {
+			fmt.Printf("Apply %d pending migrations? [Y/n] ", pending)
+		}
+
+		reader := bufio.NewReader(os.Stdin)
+		input, _ := reader.ReadString('\n')
+
+		// Format input
+		input = strings.TrimSpace(input)
+		input = strings.ToLower(input)
+
+		if !strings.HasPrefix(input, "y") && input != "" {
+			fmt.Println("Canceling migration")
+			os.Exit(0)
+		}
 	}
 
 	n, err := migrate.Exec(Conn, "sqlite3", migrations, direction)
@@ -27,7 +68,11 @@ func Migrate(Up bool) error {
 		return err
 	}
 
-	fmt.Printf("Applied %d migrations!\n", n)
+	if downgrade {
+		fmt.Printf("Downgraded %d migrations!\n", n)
+	} else {
+		fmt.Printf("Applied %d migrations!\n", n)
+	}
 	return nil
 }
 
@@ -39,17 +84,21 @@ func Status() error {
 	}
 
 	// Get the list of migrations that are yet to be applied
-	planned, _, err := migrate.PlanMigration(Conn, "sqlite3", migrations, migrate.Up, 0)
+	pending, _, err := migrate.PlanMigration(Conn, "sqlite3", migrations, migrate.Up, 0)
 	if err != nil {
 		return err
 	}
 
 	fmt.Println("Available migrations:")
 	for _, migration := range applied {
-		fmt.Printf("  Complete: %s\n", migration.Id)
+		fmt.Printf("  [x]  %s\n", migration.Id)
 	}
-	for _, migration := range planned {
-		fmt.Printf("  Pending:  %s\n", migration.Id)
+	for _, migration := range pending {
+		fmt.Printf("  [ ]  %s\n", migration.Id)
+	}
+
+	if len(pending) > 0 {
+		fmt.Printf("\nTo apply %d pending migrations, run `TastyBites migrate`\n", len(pending))
 	}
 
 	return nil