From ea63bf253e8006c54beaf276cab10eccb6754332 Mon Sep 17 00:00:00 2001
From: Gnanakeethan Balasubramaniam <gnanakeethan@gmail.com>
Date: Thu, 13 Jul 2017 20:30:06 +0530
Subject: [PATCH] UPDATE: creating ddl migration spec is now possible

SUMMARY: The DDL migration can now be generated by adding a `-ddl` and a
proper "alter" or "create" as argument value.

Migration generation had been modified to accommodate the complementary
code for ddl generation

Signed-off-by: Gnanakeethan Balasubramaniam <gnanakeethan@gmail.com>
---
 cmd/commands/generate/generate.go |   3 +-
 generate/g.go                     |   1 +
 generate/g_migration.go           | 111 +++++++++++++++++++++---------
 3 files changed, 80 insertions(+), 35 deletions(-)

diff --git a/cmd/commands/generate/generate.go b/cmd/commands/generate/generate.go
index ce874ff..2fa8cbe 100644
--- a/cmd/commands/generate/generate.go
+++ b/cmd/commands/generate/generate.go
@@ -22,7 +22,7 @@ import (
 	"github.com/beego/bee/config"
 	"github.com/beego/bee/generate"
 	"github.com/beego/bee/generate/swaggergen"
-	beeLogger "github.com/beego/bee/logger"
+	"github.com/beego/bee/logger"
 	"github.com/beego/bee/utils"
 )
 
@@ -71,6 +71,7 @@ func init() {
 	CmdGenerate.Flag.Var(&generate.SQLConn, "conn", "Connection string used by the SQLDriver to connect to a database instance.")
 	CmdGenerate.Flag.Var(&generate.Level, "level", "Either 1, 2 or 3. i.e. 1=models; 2=models and controllers; 3=models, controllers and routers.")
 	CmdGenerate.Flag.Var(&generate.Fields, "fields", "List of table Fields.")
+	CmdGenerate.Flag.Var(&generate.DDL, "ddl", "Generate DDL Migration")
 	commands.AvailableCommands = append(commands.AvailableCommands, CmdGenerate)
 }
 
diff --git a/generate/g.go b/generate/g.go
index e4260b2..691acc7 100644
--- a/generate/g.go
+++ b/generate/g.go
@@ -21,3 +21,4 @@ var SQLConn utils.DocValue
 var Level utils.DocValue
 var Tables utils.DocValue
 var Fields utils.DocValue
+var DDL utils.DocValue
diff --git a/generate/g_migration.go b/generate/g_migration.go
index 56ec063..3759f57 100644
--- a/generate/g_migration.go
+++ b/generate/g_migration.go
@@ -21,7 +21,7 @@ import (
 	"strings"
 	"time"
 
-	beeLogger "github.com/beego/bee/logger"
+	"github.com/beego/bee/logger"
 	"github.com/beego/bee/logger/colors"
 	"github.com/beego/bee/utils"
 )
@@ -203,11 +203,31 @@ func GenerateMigration(mname, upsql, downsql, curpath string) {
 	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 utils.CloseFile(f)
-		content := strings.Replace(MigrationTPL, "{{StructName}}", utils.CamelCase(mname)+"_"+today, -1)
-		content = strings.Replace(content, "{{CurrTime}}", today, -1)
-		content = strings.Replace(content, "{{UpSQL}}", upsql, -1)
-		content = strings.Replace(content, "{{DownSQL}}", downsql, -1)
-		f.WriteString(content)
+		ddlSpec := ""
+		spec := ""
+		up := ""
+		down := ""
+		if DDL != "" {
+			ddlSpec = "m.ddlSpec()"
+			switch strings.Title(DDL.String()) {
+			case "Create":
+				spec = strings.Replace(DDLSpecCreate, "{{StructName}}", utils.CamelCase(mname)+"_"+today, -1)
+				break
+			case "Alter":
+				spec = strings.Replace(DDLSpecAlter, "{{StructName}}", utils.CamelCase(mname)+"_"+today, -1)
+			}
+			spec = strings.Replace(spec, "{{tableName}}", mname, -1)
+		} else {
+			up = strings.Replace(MigrationUp, "{{UpSQL}}", upsql, -1)
+			up = strings.Replace(up, "{{StructName}}", utils.CamelCase(mname)+"_"+today, -1)
+			down = strings.Replace(MigrationDown, "{{DownSQL}}", downsql, -1)
+			down = strings.Replace(down, "{{StructName}}", utils.CamelCase(mname)+"_"+today, -1)
+		}
+
+		header := strings.Replace(MigrationHeader, "{{StructName}}", utils.CamelCase(mname)+"_"+today, -1)
+		header = strings.Replace(header, "{{ddlSpec}}", ddlSpec, -1)
+		header = strings.Replace(header, "{{CurrTime}}", today, -1)
+		f.WriteString(header + spec + up + down)
 		// Run 'gofmt' on the generated source code
 		utils.FormatSourceCode(fpath)
 		fmt.Fprintf(w, "\t%s%screate%s\t %s%s\n", "\x1b[32m", "\x1b[1m", "\x1b[21m", fpath, "\x1b[0m")
@@ -216,33 +236,56 @@ func GenerateMigration(mname, upsql, downsql, curpath string) {
 	}
 }
 
-const MigrationTPL = `package main
+const (
+	MigrationHeader = `package main
+						import (
+							"github.com/astaxie/beego/migration"
+						)
 
-import (
-	"github.com/astaxie/beego/migration"
+						// DO NOT MODIFY
+						type {{StructName}} struct {
+							migration.Migration
+						}
+
+						// DO NOT MODIFY
+						func init() {
+							m := &{{StructName}}{}
+							m.Created = "{{CurrTime}}"
+							{{ddlSpec}}
+							migration.Register("{{StructName}}", m)
+						}
+					   `
+
+	DDLSpecCreate = `
+				/*
+				refer beego/migration/doc.go
+				*/
+				func(m *{{StructName}}) ddlSpec(){
+				m.CreateTable("{{tableName}}", "InnoDB", "utf8")
+				m.PriCol("id").SetAuto(true).SetNullable(false).SetDataType("INT(10)").SetUnsigned(true)
+
+				}
+				`
+	DDLSpecAlter = `
+				/*
+				refer beego/migration/doc.go
+				*/
+				func(m *{{StructName}}) ddlSpec(){
+				m.AlterTable("{{tableName}}")
+
+				}
+				`
+	MigrationUp = `
+				// Run the migrations
+				func (m *{{StructName}}) Up() {
+					// use m.SQL("CREATE TABLE ...") to make schema update
+					{{UpSQL}}
+				}`
+	MigrationDown = `
+				// Reverse the migrations
+				func (m *{{StructName}}) Down() {
+					// use m.SQL("DROP TABLE ...") to reverse schema update
+					{{DownSQL}}
+				}
+				`
 )
-
-// DO NOT MODIFY
-type {{StructName}} struct {
-	migration.Migration
-}
-
-// DO NOT MODIFY
-func init() {
-	m := &{{StructName}}{}
-	m.Created = "{{CurrTime}}"
-	migration.Register("{{StructName}}", m)
-}
-
-// Run the migrations
-func (m *{{StructName}}) Up() {
-	// use m.SQL("CREATE TABLE ...") to make schema update
-	{{UpSQL}}
-}
-
-// Reverse the migrations
-func (m *{{StructName}}) Down() {
-	// use m.SQL("DROP TABLE ...") to reverse schema update
-	{{DownSQL}}
-}
-`