From bfb02acb43e472835b3838ed188a76e35836e9cc Mon Sep 17 00:00:00 2001 From: SanjaiyKumar Date: Wed, 18 Jun 2025 00:13:18 +0530 Subject: [PATCH] Added new flags '-allow-tables', '-allow-views', and '-allow-enums' for whitelisting the tables, views, and enums while generate command. --- cmd/jet/main.go | 86 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 61 insertions(+), 25 deletions(-) diff --git a/cmd/jet/main.go b/cmd/jet/main.go index 8512e43..384653c 100644 --- a/cmd/jet/main.go +++ b/cmd/jet/main.go @@ -50,8 +50,17 @@ var ( tablePkg string viewPkg string enumPkg string + + allowTables string + allowViews string + allowEnums string ) +type templateFilter struct { + names []string + allow 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(&allowTables, "allow-tables", "", `Comma-separated list of tables to allow.`) + flag.StringVar(&allowViews, "allow-views", "", `Comma-separated list of views to allow. Will override ignore-views flag.`) + flag.StringVar(&allowEnums, "allow-enums", "", `Comma-separated list of enums to allow.`) } func main() { @@ -95,15 +108,15 @@ func main() { } source := getSource() - ignoreTablesList := parseList(ignoreTables) - ignoreViewsList := parseList(ignoreViews) - ignoreEnumsList := parseList(ignoreEnums) + tablesFilter := createTemplateFilter(ignoreTables, allowTables) + viewsFilter := createTemplateFilter(ignoreViews, allowViews) + enumsFilter := createTemplateFilter(ignoreEnums, allowEnums) 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", "allow-tables", "allow-views", + "allow-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,36 @@ func genTemplate(dialect jet.Dialect, ignoreTables []string, ignoreViews []strin ) }) } + +func createTemplateFilter(ignoreList, allowList string) templateFilter { + ignoreListParsed := parseList(ignoreList) + allowListParsed := parseList(allowList) + + if len(allowListParsed) > 0 { + return templateFilter{ + names: allowListParsed, + allow: true, + } + } + + return templateFilter{ + names: ignoreListParsed, + allow: false, + } +} + +func shouldSkipTable(table metadata.Table, config templateFilter) bool { + if config.allow { + return !strslice.Contains(config.names, strings.ToLower(table.Name)) + } + + return strslice.Contains(config.names, strings.ToLower(table.Name)) +} + +func shouldSkipEnum(enum metadata.Enum, config templateFilter) bool { + if config.allow { + return !strslice.Contains(config.names, strings.ToLower(enum.Name)) + } + + return strslice.Contains(config.names, strings.ToLower(enum.Name)) +}