Merge pull request #506 from Sanjaiy/master

Added whitelist flags for generate command

cmd/jet/main.go

@@ -50,8 +50,17 @@ var (
 	tablePkg string
 	viewPkg  string
 	enumPkg  string
+
+	tables string
+	views  string
+	enums  string
 )
 
+type templateFilter struct {
+	names []string
+	ignore bool
+}
+
 func init() {
 	flag.StringVar(&source, "source", "", "Database system name (postgres, mysql, cockroachdb, mariadb or sqlite)")
 
@@ -84,6 +93,10 @@ func init() {
 	flag.StringVar(&tablePkg, "rel-table-path", "table", "Relative path for the Table files package from the destination directory.")
 	flag.StringVar(&viewPkg, "rel-view-path", "view", "Relative path for the View files package from the destination directory.")
 	flag.StringVar(&enumPkg, "rel-enum-path", "enum", "Relative path for the Enum files package from the destination directory.")
+
+	flag.StringVar(&tables, "tables", "", `Comma-separated list of tables to generate.`)
+	flag.StringVar(&views, "views", "", `Comma-separated list of views to generate.`)
+	flag.StringVar(&enums, "enums", "", `Comma-separated list of enums to generate.`)
 }
 
 func main() {
@@ -95,15 +108,15 @@ func main() {
 	}
 
 	source := getSource()
-	ignoreTablesList := parseList(ignoreTables)
+	tablesFilter := createTemplateFilter(ignoreTables, tables, "tables")
-	ignoreViewsList := parseList(ignoreViews)
+	viewsFilter := createTemplateFilter(ignoreViews, views, "views")
-	ignoreEnumsList := parseList(ignoreEnums)
+	enumsFilter := createTemplateFilter(ignoreEnums, enums, "enums")
 
 	var err error
 
 	switch source {
 	case "postgresql", "postgres", "cockroachdb", "cockroach":
-		generatorTemplate := genTemplate(postgres2.Dialect, ignoreTablesList, ignoreViewsList, ignoreEnumsList)
+		generatorTemplate := genTemplate(postgres2.Dialect, tablesFilter, viewsFilter, enumsFilter)
 
 		if dsn != "" {
 			err = postgresgen.GenerateDSN(dsn, schemaName, destDir, generatorTemplate)
@@ -129,7 +142,7 @@ func main() {
 		)
 
 	case "mysql", "mysqlx", "mariadb":
-		generatorTemplate := genTemplate(mysql.Dialect, ignoreTablesList, ignoreViewsList, ignoreEnumsList)
+		generatorTemplate := genTemplate(mysql.Dialect, tablesFilter, viewsFilter, enumsFilter)
 
 		if dsn != "" {
 			err = mysqlgen.GenerateDSN(dsn, destDir, generatorTemplate)
@@ -158,7 +171,7 @@ func main() {
 		err = sqlitegen.GenerateDSN(
 			dsn,
 			destDir,
-			genTemplate(sqlite.Dialect, ignoreTablesList, ignoreViewsList, ignoreEnumsList),
+			genTemplate(sqlite.Dialect, tablesFilter, viewsFilter, enumsFilter),
 		)
 
 	case "":
@@ -184,7 +197,8 @@ func usage() {
 		"path",
 		"ignore-tables", "ignore-views", "ignore-enums",
 		"skip-model", "skip-sql-builder",
-		"rel-model-path", "rel-table-path", "rel-view-path", "rel-enum-path",
+		"rel-model-path", "rel-table-path", "rel-view-path", "rel-enum-path", "tables", "views",
+		"enums",
 	}
 
 	for _, name := range order {
@@ -245,38 +259,27 @@ func parseList(list string) []string {
 	return ret
 }
 
-func genTemplate(dialect jet.Dialect, ignoreTables []string, ignoreViews []string, ignoreEnums []string) template.Template {
 
-	shouldSkipTable := func(table metadata.Table) bool {
-		return strslice.Contains(ignoreTables, strings.ToLower(table.Name))
-	}
-
-	shouldSkipView := func(view metadata.Table) bool {
-		return strslice.Contains(ignoreViews, strings.ToLower(view.Name))
-	}
-
-	shouldSkipEnum := func(enum metadata.Enum) bool {
-		return strslice.Contains(ignoreEnums, strings.ToLower(enum.Name))
-	}
 
+func genTemplate(dialect jet.Dialect, tablesFilter, viewsFilter, enumsFilter templateFilter) template.Template {
 	return template.Default(dialect).
 		UseSchema(func(schemaMetaData metadata.Schema) template.Schema {
 			return template.DefaultSchema(schemaMetaData).
 				UseModel(template.DefaultModel().ShouldSkip(skipModel).UsePath(modelPkg).
 					UseTable(func(table metadata.Table) template.TableModel {
-						if shouldSkipTable(table) {
+						if shouldSkipTable(table, tablesFilter) {
 							return template.TableModel{Skip: true}
 						}
 						return template.DefaultTableModel(table)
 					}).
 					UseView(func(view metadata.Table) template.ViewModel {
-						if shouldSkipView(view) {
+						if shouldSkipTable(view, viewsFilter) {
 							return template.ViewModel{Skip: true}
 						}
 						return template.DefaultViewModel(view)
 					}).
 					UseEnum(func(enum metadata.Enum) template.EnumModel {
-						if shouldSkipEnum(enum) {
+						if shouldSkipEnum(enum, enumsFilter) {
 							return template.EnumModel{Skip: true}
 						}
 						return template.DefaultEnumModel(enum)
@@ -284,21 +287,21 @@ func genTemplate(dialect jet.Dialect, ignoreTables []string, ignoreViews []strin
 				).
 				UseSQLBuilder(template.DefaultSQLBuilder().ShouldSkip(skipSQLBuilder).
 					UseTable(func(table metadata.Table) template.TableSQLBuilder {
-						if shouldSkipTable(table) {
+						if shouldSkipTable(table, tablesFilter) {
 							return template.TableSQLBuilder{Skip: true}
 						}
 
 						return template.DefaultTableSQLBuilder(table).UsePath(tablePkg)
 					}).
 					UseView(func(table metadata.Table) template.ViewSQLBuilder {
-						if shouldSkipView(table) {
+						if shouldSkipTable(table, viewsFilter) {
 							return template.ViewSQLBuilder{Skip: true}
 						}
 
 						return template.DefaultViewSQLBuilder(table).UsePath(viewPkg)
 					}).
 					UseEnum(func(enum metadata.Enum) template.EnumSQLBuilder {
-						if shouldSkipEnum(enum) {
+						if shouldSkipEnum(enum, enumsFilter) {
 							return template.EnumSQLBuilder{Skip: true}
 						}
 
@@ -307,3 +310,37 @@ func genTemplate(dialect jet.Dialect, ignoreTables []string, ignoreViews []strin
 				)
 		})
 }
+
+func createTemplateFilter(ignoreList, allowList, filterType string) templateFilter {
+	if ignoreList != "" && allowList != "" {
+		printErrorAndExit(fmt.Sprintf("ERROR: cannot use both -%s and -ignore-%s flags simultaneously. Please specify only one option.", filterType, filterType))
+	}
+
+	if allowList != "" {
+		return templateFilter{
+			names: parseList(allowList),
+			ignore: false,
+		}
+	}
+
+	return templateFilter{
+		names: parseList(ignoreList),
+		ignore: true,
+	}
+}
+
+func shouldSkipTable(table metadata.Table, filter templateFilter) bool {
+	if filter.ignore {
+		return strslice.Contains(filter.names, strings.ToLower(table.Name))
+	}
+
+	return !strslice.Contains(filter.names, strings.ToLower(table.Name))
+}
+
+func shouldSkipEnum(enum metadata.Enum, filter templateFilter) bool {
+	if filter.ignore {
+		return strslice.Contains(filter.names, strings.ToLower(enum.Name))
+	}
+
+	return !strslice.Contains(filter.names, strings.ToLower(enum.Name))
+}

tests/mysql/generator_test.go

@@ -1,6 +1,7 @@
 package mysql
 
 import (
+	"bytes"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -1002,3 +1003,116 @@ func newAllTypesTableImpl(schemaName, tableName, alias string) allTypesTable {
 	}
 }
 `
+
+func TestAllowTablesViewsEnums(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+	}{
+		{
+			name: "with dsn",
+			args: []string{
+				"-dsn=mysql://" + dbconfig.MySQLConnectionString(sourceIsMariaDB(), "dvds"),
+				"-tables=actor,ADDRESS,Category, city ,country,staff,store,rental",
+				"-views=actor_info,CUSTomER_LIST, film_list",
+				"-enums=film_list_rating,film_rating",
+				"-path=" + genTestDir3,
+			},
+		},
+		{
+			name: "without dsn",
+			args: []string{
+				"-source=MySQL",
+				"-dbname=dvds",
+				"-host=" + dbconfig.MySqLHost,
+				"-port=" + strconv.Itoa(dbconfig.MySQLPort),
+				"-user=" + dbconfig.MySQLUser,
+				"-password=" + dbconfig.MySQLPassword,
+				"-tables=actor,ADDRESS,Category, city ,country,staff,store,rental",
+				"-views=actor_info,CUSTomER_LIST, film_list",
+				"-enums=film_list_rating,film_rating",
+				"-path=" + genTestDir3,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cmd := exec.Command("jet", tt.args...)
+
+			cmd.Stderr = os.Stderr
+			cmd.Stdout = os.Stdout
+
+			err := cmd.Run()
+			require.NoError(t, err)
+
+			testutils.AssertFileNamesEqual(t, genTestDir3+"/dvds/table", "category.go", "actor.go", "address.go",
+				"city.go", "country.go", "staff.go", "store.go", "rental.go", "table_use_schema.go")
+
+			testutils.AssertFileNamesEqual(t, genTestDir3+"/dvds/view", "actor_info.go", "customer_list.go",
+				"film_list.go", "view_use_schema.go")
+
+			testutils.AssertFileNamesEqual(t, genTestDir3+"/dvds/enum", "film_list_rating.go", "film_rating.go")
+
+			testutils.AssertFileNamesEqual(t, genTestDir3+"/dvds/model", "category.go", "actor.go", "address.go",
+				"city.go", "country.go", "staff.go", "store.go", "rental.go", "actor_info.go",
+				"customer_list.go", "film_list.go", "film_list_rating.go", "film_rating.go")
+		})
+	}
+}
+
+func TestAllowAndIgnoreTablesViewsEnums(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+	}{
+		{
+			name: "with dsn",
+			args: []string{
+				"-dsn=mysql://" + dbconfig.MySQLConnectionString(sourceIsMariaDB(), "dvds"),
+				"-tables=actor,ADDRESS,Category, city ,country,staff,store,rental",
+				"-views=actor_info,CUSTomER_LIST, film_list",
+				"-enums=film_list_rating,film_rating",
+				"-ignore-tables=actor,ADDRESS,store,rental",
+				"-ignore-views=film_list",
+				"-ignore-enums=film_rating",
+				"-path=" + genTestDir3,
+			},
+		},
+		{
+			name: "without dsn",
+			args: []string{
+				"-source=MySQL",
+				"-dbname=dvds",
+				"-host=" + dbconfig.MySqLHost,
+				"-port=" + strconv.Itoa(dbconfig.MySQLPort),
+				"-user=" + dbconfig.MySQLUser,
+				"-password=" + dbconfig.MySQLPassword,
+				"-tables=actor,ADDRESS,Category, city ,country,staff,store,rental",
+				"-views=actor_info,CUSTomER_LIST, film_list",
+				"-enums=film_list_rating,film_rating",
+				"-ignore-tables=actor,ADDRESS,store,rental",
+				"-ignore-views=film_list",
+				"-ignore-enums=film_rating",
+				"-path=" + genTestDir3,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cmd := exec.Command("jet", tt.args...)
+
+			var stdOut bytes.Buffer
+			cmd.Stderr = os.Stderr
+			cmd.Stdout = &stdOut
+
+			err := cmd.Run()
+			require.Error(t, err)
+			require.Equal(t, "exit status 1", err.Error())
+
+			stdOutput := stdOut.String()
+			require.Contains(t, stdOutput, "ERROR: cannot use both -tables and -ignore-tables flags simultaneously. Please specify only one option.")
+		})
+	}
+}

tests/postgres/generator_test.go

@@ -1,6 +1,7 @@
 package postgres
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"os/exec"
@@ -1384,3 +1385,123 @@ func newLinkTableImpl(schemaName, tableName, alias string) linkTable {
 	}
 }
 `
+
+func TestAllowTablesViewsEnums(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+	}{
+		{
+			name: "with dsn",
+			args: []string{
+				"-dsn=" + defaultDSN(),
+				"-schema=dvds",
+				"-tables=actor,ADDRESS,country, Film , cITY,",
+				"-views=Actor_info, FILM_LIST ,staff_list",
+				"-enums=mpaa_rating",
+				"-path=" + genTestDir2,
+			},
+		},
+		{
+			name: "without dsn",
+			args: []string{
+				"-source=PostgreSQL",
+				"-host=localhost",
+				"-port=" + strconv.Itoa(dbconfig.PgPort),
+				"-user=jet",
+				"-password=jet",
+				"-dbname=jetdb",
+				"-schema=dvds",
+				"-tables=actor,ADDRESS,country, Film , cITY,",
+				"-views=Actor_info, FILM_LIST ,staff_list",
+				"-enums=mpaa_rating",
+				"-path=" + genTestDir2,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := os.RemoveAll(genTestDir2)
+			require.NoError(t, err)
+
+			cmd := exec.Command("jet", tt.args...)
+
+			fmt.Println(cmd.Args)
+			cmd.Stderr = os.Stderr
+			cmd.Stdout = os.Stdout
+
+			err = cmd.Run()
+			require.NoError(t, err)
+
+			// Table SQL Builder files
+			testutils.AssertFileNamesEqual(t, "./.gentestdata2/jetdb/dvds/table", "actor.go", "address.go",
+				"country.go", "film.go", "city.go", "table_use_schema.go")
+
+			// View SQL Builder files
+			testutils.AssertFileNamesEqual(t, "./.gentestdata2/jetdb/dvds/view", "actor_info.go", "film_list.go",
+				"staff_list.go", "view_use_schema.go")
+
+			// Enums SQL Builder files
+			file.Exists(t, "./.gentestdata2/jetdb/dvds/enum", "mpaa_rating.go")
+
+			// Model files
+			testutils.AssertFileNamesEqual(t, "./.gentestdata2/jetdb/dvds/model", "actor.go", "address.go",
+				"country.go", "film.go", "city.go", "actor_info.go", "film_list.go", "staff_list.go", "mpaa_rating.go")
+		})
+	}
+}
+
+func TestAllowAndIgnoreEnums(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+	}{
+		{
+			name: "with dsn",
+			args: []string{
+				"-dsn=" + defaultDSN(),
+				"-schema=dvds",
+				"-enums=mpaa_rating",
+				"-ignore-enums=mpaa_rating",
+				"-path=" + genTestDir2,
+			},
+		},
+		{
+			name: "without dsn",
+			args: []string{
+				"-source=PostgreSQL",
+				"-host=localhost",
+				"-port=" + strconv.Itoa(dbconfig.PgPort),
+				"-user=jet",
+				"-password=jet",
+				"-dbname=jetdb",
+				"-schema=dvds",
+				"-enums=mpaa_rating",
+				"-ignore-enums=mpaa_rating",
+				"-path=" + genTestDir2,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := os.RemoveAll(genTestDir2)
+			require.NoError(t, err)
+
+			cmd := exec.Command("jet", tt.args...)
+
+			fmt.Println(cmd.Args)
+			var stdOut bytes.Buffer
+			cmd.Stderr = os.Stderr
+			cmd.Stdout = &stdOut
+
+			err = cmd.Run()
+			require.Error(t, err)
+			require.Equal(t, "exit status 1", err.Error())
+
+			stdOutput := stdOut.String()
+			require.Contains(t, stdOutput, "ERROR: cannot use both -enums and -ignore-enums flags simultaneously. Please specify only one option.")
+		})
+	}
+}

tests/sqlite/generator_test.go

@@ -1,6 +1,7 @@
 package sqlite
 
 import (
+	"bytes"
 	"os"
 	"os/exec"
 	"reflect"
@@ -427,3 +428,47 @@ type Address struct {
 	LastUpdate time.Time
 }
 `
+
+func TestAllowTablesEnums(t *testing.T) {
+	cmd := exec.Command("jet",
+		"-source=SQLite",
+		"-dsn=file://"+testDatabaseFilePath,
+		"-tables=actor,Address,CATEGORY , city ,film,rental,store",
+		"-views=customer_list, film_list,STAFF_LIst",
+		"-path="+genDestDir)
+
+	cmd.Stderr = os.Stderr
+	cmd.Stdout = os.Stdout
+
+	err := cmd.Run()
+	require.NoError(t, err)
+
+	testutils.AssertFileNamesEqual(t, genDestDir+"/table", "actor.go", "address.go", "category.go", "city.go",
+		"film.go", "rental.go", "store.go", "table_use_schema.go")
+
+	testutils.AssertFileNamesEqual(t, genDestDir+"/view", "customer_list.go", "film_list.go", "staff_list.go",
+		"view_use_schema.go")
+
+	testutils.AssertFileNamesEqual(t, genDestDir+"/model", "actor.go", "address.go", "category.go", "city.go",
+		"film.go", "rental.go", "store.go", "customer_list.go", "film_list.go", "staff_list.go")
+}
+
+func TestAllowAndIgnoreViews(t *testing.T) {
+	cmd := exec.Command("jet",
+		"-source=SQLite",
+		"-dsn=file://"+testDatabaseFilePath,
+		"-views=customer_list, film_list,STAFF_LIst",
+		"-ignore-views=customer_list",
+		"-path="+genDestDir)
+
+	var stdOut bytes.Buffer
+	cmd.Stderr = os.Stderr
+	cmd.Stdout = &stdOut
+
+	err := cmd.Run()
+	require.Error(t, err)
+	require.Equal(t, "exit status 1", err.Error())
+
+	stdOutput := stdOut.String()
+	require.Contains(t, stdOutput, "ERROR: cannot use both -views and -ignore-views flags simultaneously. Please specify only one option.")
+}