generate models for specified tables

g.go

@@ -20,8 +20,9 @@ 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
@@ -46,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")
@@ -86,8 +89,9 @@ 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]

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
 		}