Merge pull request #20 from sboy99/feat/storage-selection

feat: storage selection and update config

config/config.go

@@ -40,7 +40,7 @@ const (
 )
 
 const (
-	FILE  StorageEnum = "file"
+	LOCAL StorageEnum = "local"
 	CLOUD StorageEnum = "cloud"
 )
 
@@ -72,6 +72,9 @@ func Save(config *Config) error {
 	viper.Set("db.port", config.DB.Port)
 	viper.Set("db.username", config.DB.Username)
 	viper.Set("db.password", config.DB.Password)
+	// Storage //
+	viper.Set("storage.type", config.Storage.Type)
+	viper.Set("storage.dest", config.Storage.Dest)
 
 	if err := viper.WriteConfig(); err != nil {
 		return err

config/validator.go

@@ -11,7 +11,7 @@ func setDefaults() {
 	viper.SetDefault("app.name", "go-vault")
 	viper.SetDefault("app.version", "0.0.1")
 
-	viper.SetDefault("storage.type", "file")
+	viper.SetDefault("storage.type", "local")
 	viper.SetDefault("storage.dest", "./generated")
 }
 

internal/config/service.go

@@ -45,6 +45,12 @@ func (c *ConfigService) SetupConfig() {
 		return
 	}
 
+	storageType, err := ui.DisplaySelectStorageTypePrompt()
+	if err != nil {
+		logger.Error("Failed to display select storage prompt\nDetails: %v", err)
+		return
+	}
+
 	// Save config //
 	cfg := config.GetConfig()
 	cfg.DB.Type = dbType
@@ -53,6 +59,9 @@ func (c *ConfigService) SetupConfig() {
 	cfg.DB.Name = dbName
 	cfg.DB.Username = dbUser
 	cfg.DB.Password = dbPass
+
+	cfg.Storage.Type = storageType
+
 	if err = config.Save(cfg); err != nil {
 		logger.Error("Failed to save config\nDetails: %v", err)
 	}

internal/storage/local_storage.go

@@ -0,0 +1,52 @@
+package storage
+
+import (
+	"os"
+)
+
+type LocalStorage struct {
+	BasePath string
+}
+
+func NewLocalStorage(basePath string) *LocalStorage {
+	return &LocalStorage{BasePath: basePath}
+}
+
+func (ls *LocalStorage) Save(filename string, data []byte) error {
+	file, err := os.Create(ls.filePath(filename))
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	_, err = file.Write(data)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (ls *LocalStorage) Load(filename string) ([]byte, error) {
+	file, err := os.Open(ls.filePath(filename))
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	data := make([]byte, 0)
+	_, err = file.Read(data)
+	if err != nil {
+		return nil, err
+	}
+
+	return data, nil
+}
+
+func (ls *LocalStorage) Delete(filename string) error {
+	return os.Remove(ls.filePath(filename))
+}
+
+func (ls *LocalStorage) filePath(filename string) string {
+	return ls.BasePath + "/" + filename
+}

internal/storage/storage.go

@@ -16,7 +16,7 @@ type Storage struct {
 func NewStorage() *Storage {
 	return &Storage{
 		storageMap: map[config.StorageEnum]IStorage{
-			config.FILE:  NewFileStorage(config.GetConfig().Storage.Dest),
+			config.LOCAL: NewLocalStorage(config.GetConfig().Storage.Dest),
 			config.CLOUD: nil,
 		},
 	}

internal/ui/ui.go

@@ -93,3 +93,17 @@ func DisplayInputDatabasePasswordPrompt() (string, error) {
 	}
 	return result, nil
 }
+
+func DisplaySelectStorageTypePrompt() (config.StorageEnum, error) {
+	// Prompt for Storage selection //
+	prompt := promptui.Select{
+		Label: "Select Storage",
+		Items: []config.StorageEnum{config.LOCAL, config.CLOUD},
+	}
+	// Run the prompt //
+	_, result, err := prompt.Run()
+	if err != nil {
+		return "", err
+	}
+	return config.StorageEnum(result), nil
+}

main.go

@@ -12,8 +12,6 @@ func main() {
 	// Logger
 	logger.Init(logger.DEBUG)
 
-	logger.Info("args %v", os.Args)
-
 	// User Config //
 	if len(os.Args) > 1 && os.Args[1] != "setup" {
 		if err := config.Load(); err != nil {

pkg/pg_dump/pg_dump.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"fmt"
 	"strings"
+	"time"
 
 	_ "github.com/lib/pq"
 )
@@ -36,8 +37,16 @@ func (p *PgDump) Dump() ([]byte, error) {
 	}
 
 	var dumpSql bytes.Buffer
-	dumpSql.WriteString(makeSqlComment("PostgreSQL database dump"))
-	dumpSql.WriteString(makeSqlComment(fmt.Sprintf("Database version: %s", dbVersion)))
+	dumpSql.WriteString(makeSqlComment(
+		"PostgreSQL database dump" +
+			"\n" +
+			fmt.Sprintf("Database version: %s", dbVersion) +
+			"\n" +
+			"Generated by go-vault" +
+			"\n" +
+			fmt.Sprintf("Generated at: %s", time.Now().Format(time.RFC3339)),
+	))
+	dumpSql.WriteString("\n")
 	dumpSql.WriteString("SET statement_timeout = 0;\n")
 	dumpSql.WriteString("SET lock_timeout = 0;\n")
 	dumpSql.WriteString("SET client_encoding = 'UTF8';\n")
@@ -181,8 +190,9 @@ func (p *PgDump) generateCreateTableStatementsForTables(tables []string) (string
 			continue
 		}
 		dumpSql.WriteString(createTableStatement)
-		dumpSql.WriteString("\n\n")
+		dumpSql.WriteString("\n")
 	}
+	dumpSql.WriteString("\n")
 	return dumpSql.String(), nil
 }
 

pkg/pg_dump/sql_util.go

@@ -1,9 +1,12 @@
 package pgdump
 
-import "fmt"
+import (
+	"fmt"
+	"strings"
+)
 
 func makeSqlComment(comment string) string {
-	return fmt.Sprintf("--\n-- %s\n--\n", comment)
+	return fmt.Sprintf("--\n-- %s\n--\n", strings.Replace(comment, "\n", "\n-- ", -1))
 }
 
 func getListSchemasQuery() string {