From 2556579ad0781f4d4f1be9478a90883c1b931e95 Mon Sep 17 00:00:00 2001
From: ZhengYang <zhengyang4k@gmail.com>
Date: Thu, 7 Aug 2014 13:52:28 +0800
Subject: [PATCH] generate migration file

---
 g.go           |  6 ++--
 g_migration.go | 83 +++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/g.go b/g.go
index 6fdb3ae..a63f1ae 100644
--- a/g.go
+++ b/g.go
@@ -90,9 +90,9 @@ func generateCode(cmd *Command, args []string) {
 		generateModel(string(driver), string(conn), string(level), curpath)
 	case "migration":
 		if len(args) == 2 {
-			filename := args[1]
-			ColorLog("[INFO] Using '%s' as migration file name\n", filename)
-			generateMigration(filename, curpath)
+			mname := args[1]
+			ColorLog("[INFO] Using '%s' as migration name\n", mname)
+			generateMigration(mname, curpath)
 		} else {
 			ColorLog("[ERRO] Wrong number of arguments\n")
 			ColorLog("[HINT] Usage: bee generate migration [filename]\n")
diff --git a/g_migration.go b/g_migration.go
index 0a301f0..8c3622f 100644
--- a/g_migration.go
+++ b/g_migration.go
@@ -14,9 +14,84 @@
 
 package main
 
-import "fmt"
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path"
+	"strings"
+	"time"
+)
 
-func generateMigration(filename string, curpath string) {
-	fmt.Println("filename:", filename)
-	fmt.Println("curpath:", curpath)
+const (
+	M_PATH        = "migrations"
+	M_DATE_FORMAT = "2006-01-02"
+)
+
+// generateMigration generates migration file template for database schema update.
+// The generated file template consists of an up() method for updating schema and
+// a down() method for reverting the update.
+func generateMigration(mname string, curpath string) {
+	migrationFilePath := path.Join(curpath, M_PATH)
+	if _, err := os.Stat(migrationFilePath); os.IsNotExist(err) {
+		// create migrations directory
+		if err := os.Mkdir(migrationFilePath, 0777); err != nil {
+			ColorLog("[ERRO] Could not create migration directory: %s\n", err)
+			os.Exit(2)
+		}
+	}
+	// create file
+	today := time.Now().Format(M_DATE_FORMAT)
+	fpath := path.Join(migrationFilePath, fmt.Sprintf("%s_%s.go", today, mname))
+	if f, err := os.OpenFile(fpath, os.O_CREATE|os.O_EXCL|os.O_RDWR, 0666); err == nil {
+		defer f.Close()
+		content := strings.Replace(MIGRATION_TPL, "{{StructName}}", camelCase(mname), -1)
+		content = strings.Replace(content, "{{DateFormat}}", M_DATE_FORMAT, -1)
+		content = strings.Replace(content, "{{CurrTime}}", today, -1)
+		f.WriteString(content)
+		// gofmt generated source code
+		formatSourceCode(fpath)
+		ColorLog("[INFO] Migration file generated: %s\n", fpath)
+	} else {
+		// error creating file
+		ColorLog("[ERRO] Could not create migration file: %s\n", err)
+		os.Exit(2)
+	}
 }
+
+// formatSourceCode formats the source code using gofmt
+func formatSourceCode(fpath string) {
+	cmd := exec.Command("gofmt", "-w", fpath)
+	cmd.Run()
+}
+
+const MIGRATION_TPL = `
+package main
+
+import (
+	"time"
+
+	"github.com/astaxie/beego"
+	"github.com/astaxie/beego/migration"
+)
+
+func init() {
+	m := &{{StructName}}{}
+	m.Created = time.Parse("{{DateFormat}}", "{{CurrTime}}")
+	migration.Register(m)
+}
+
+type {{StructName}} struct {
+	migration.Migration
+}
+
+// Run the migrations
+func (m *{{StructName}}) up() {
+	// use m.Sql("create table ...") to make schema update
+}
+
+// Reverse the migrations
+func (m *{{StructName}}) down() {
+	// use m.Sql("drop table ...") to reverse schema update
+}
+`