Merge pull request #37 from ZhengYang/master

Generate Models for selected tables + Generate Migrations

g.go

@@ -20,12 +20,16 @@ var cmdGenerate = &Command{
 	UsageLine: "generate [Command]",
 	Short:     "generate code based on application",
 	Long: `
-bee generate model [-driver=mysql] [-conn=root:@tcp(127.0.0.1:3306)/test] [-level=1]
+bee generate model [-tables=""] [-driver=mysql] [-conn=root:@tcp(127.0.0.1:3306)/test] [-level=1]
     generate model based on an existing database
+    -tables: a list of table names separated by ',', default is empty, indicating all tables
     -driver: [mysql | postgresql | sqlite], the default is mysql
     -conn:   the connection string used by the driver, the default is root:@tcp(127.0.0.1:3306)/test
     -level:  [1 | 2 | 3], 1 = model; 2 = models,controller; 3 = models,controllers,router
 
+bee generate migration [filename]
+    generate migration file for making database schema update
+
 bee generate controller [modelfile]
     generate RESTFul controllers based on modelfile             
 
@@ -43,9 +47,11 @@ bee generate test [routerfile]
 var driver docValue
 var conn docValue
 var level docValue
+var tables docValue
 
 func init() {
 	cmdGenerate.Run = generateCode
+	cmdGenerate.Flag.Var(&tables, "tables", "specify tables to generate model")
 	cmdGenerate.Flag.Var(&driver, "driver", "database driver: mysql, postgresql, etc.")
 	cmdGenerate.Flag.Var(&conn, "conn", "connection string used by the driver to connect to a database instance")
 	cmdGenerate.Flag.Var(&level, "level", "1 = models only; 2 = models and controllers; 3 = models, controllers and routers")
@@ -83,8 +89,19 @@ func generateCode(cmd *Command, args []string) {
 		}
 		ColorLog("[INFO] Using '%s' as 'driver'\n", driver)
 		ColorLog("[INFO] Using '%s' as 'conn'\n", conn)
+		ColorLog("[INFO] Using '%s' as 'tables'", tables)
 		ColorLog("[INFO] Using '%s' as 'level'\n", level)
-		generateModel(string(driver), string(conn), string(level), curpath)
+		generateModel(string(driver), string(conn), string(level), string(tables), curpath)
+	case "migration":
+		if len(args) == 2 {
+			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")
+			os.Exit(2)
+		}
 	default:
 		ColorLog("[ERRO] command is missing\n")
 	}

g_migration.go

@@ -0,0 +1,97 @@
+// Copyright 2013 bee authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"): you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path"
+	"strings"
+	"time"
+)
+
+const (
+	M_PATH        = "migrations"
+	M_DATE_FORMAT = "20060102_150405"
+)
+
+// 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)+"_"+today, -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/migration"
+)
+
+// 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
+}
+
+// Reverse the migrations
+func (m *{{StructName}}) Down() {
+	// use m.Sql("DROP TABLE ...") to reverse schema update
+}
+`

g_models.go

@@ -191,7 +191,7 @@ func (tag *OrmTag) String() string {
 	return fmt.Sprintf("`orm:\"%s\"`", strings.Join(ormOptions, ";"))
 }
 
-func generateModel(driver string, connStr string, level string, currpath string) {
+func generateModel(driver string, connStr string, level string, tables string, currpath string) {
 	var mode byte
 	if level == "1" {
 		mode = O_MODEL
@@ -204,12 +204,19 @@ func generateModel(driver string, connStr string, level string, currpath string)
 		ColorLog("[HINT] Level must be either 1, 2 or 3\n")
 		os.Exit(2)
 	}
-	gen(driver, connStr, mode, currpath)
+	var selectedTables map[string]bool
+	if tables != "" {
+		selectedTables = make(map[string]bool)
+		for _, v := range strings.Split(tables, ",") {
+			selectedTables[v] = true
+		}
+	}
+	gen(driver, connStr, mode, selectedTables, currpath)
 }
 
 // Generate takes table, column and foreign key information from database connection
 // and generate corresponding golang source files
-func gen(dbms string, connStr string, mode byte, currpath string) {
+func gen(dbms string, connStr string, mode byte, selectedTableNames map[string]bool, currpath string) {
 	db, err := sql.Open(dbms, connStr)
 	if err != nil {
 		ColorLog("[ERRO] Could not connect to %s: %s\n", dbms, connStr)
@@ -224,7 +231,7 @@ func gen(dbms string, connStr string, mode byte, currpath string) {
 	mvcPath.ControllerPath = path.Join(currpath, "controllers")
 	mvcPath.RouterPath = path.Join(currpath, "routers")
 	createPaths(mode, mvcPath)
-	writeSourceFiles(tables, mode, mvcPath)
+	writeSourceFiles(tables, mode, mvcPath, selectedTableNames)
 }
 
 // getTables gets a list table names in current database
@@ -418,24 +425,30 @@ func createPaths(mode byte, paths *MvcPath) {
 // writeSourceFiles generates source files for model/controller/router
 // It will wipe the following directories and recreate them:./models, ./controllers, ./routers
 // Newly geneated files will be inside these folders.
-func writeSourceFiles(tables []*Table, mode byte, paths *MvcPath) {
+func writeSourceFiles(tables []*Table, mode byte, paths *MvcPath, selectedTables map[string]bool) {
 	if (O_MODEL & mode) == O_MODEL {
 		ColorLog("[INFO] Creating model files...\n")
-		writeModelFiles(tables, paths.ModelPath)
+		writeModelFiles(tables, paths.ModelPath, selectedTables)
 	}
 	if (O_CONTROLLER & mode) == O_CONTROLLER {
 		ColorLog("[INFO] Creating controller files...\n")
-		writeControllerFiles(tables, paths.ControllerPath)
+		writeControllerFiles(tables, paths.ControllerPath, selectedTables)
 	}
 	if (O_ROUTER & mode) == O_ROUTER {
 		ColorLog("[INFO] Creating router files...\n")
-		writeRouterFile(tables, paths.RouterPath)
+		writeRouterFile(tables, paths.RouterPath, selectedTables)
 	}
 }
 
 // writeModelFiles generates model files
-func writeModelFiles(tables []*Table, mPath string) {
+func writeModelFiles(tables []*Table, mPath string, selectedTables map[string]bool) {
 	for _, tb := range tables {
+		// if selectedTables map is not nil and this table is not selected, ignore it
+		if selectedTables != nil {
+			if _, selected := selectedTables[tb.Name]; !selected {
+				continue
+			}
+		}
 		filename := getFileName(tb.Name)
 		fpath := path.Join(mPath, filename+".go")
 		f, err := os.OpenFile(fpath, os.O_CREATE|os.O_EXCL|os.O_RDWR, 0666)
@@ -462,8 +475,14 @@ func writeModelFiles(tables []*Table, mPath string) {
 }
 
 // writeControllerFiles generates controller files
-func writeControllerFiles(tables []*Table, cPath string) {
+func writeControllerFiles(tables []*Table, cPath string, selectedTables map[string]bool) {
 	for _, tb := range tables {
+		// if selectedTables map is not nil and this table is not selected, ignore it
+		if selectedTables != nil {
+			if _, selected := selectedTables[tb.Name]; !selected {
+				continue
+			}
+		}
 		if tb.Pk == "" {
 			continue
 		}
@@ -486,9 +505,15 @@ func writeControllerFiles(tables []*Table, cPath string) {
 }
 
 // writeRouterFile generates router file
-func writeRouterFile(tables []*Table, rPath string) {
+func writeRouterFile(tables []*Table, rPath string, selectedTables map[string]bool) {
 	var nameSpaces []string
 	for _, tb := range tables {
+		// if selectedTables map is not nil and this table is not selected, ignore it
+		if selectedTables != nil {
+			if _, selected := selectedTables[tb.Name]; !selected {
+				continue
+			}
+		}
 		if tb.Pk == "" {
 			continue
 		}