diff --git a/g.go b/g.go index a63f1ae..b91698f 100644 --- a/g.go +++ b/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] diff --git a/g_models.go b/g_models.go index f7f4a79..69b3482 100644 --- a/g_models.go +++ b/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 }