Merge pull request #118 from FoGhost/add-target-table-arg

Add target-tables argument

cmd/root.go

@@ -125,8 +125,9 @@ func setRootOpts(cmd *cobra.Command, opts *internal.ArgType) {
 	cmd.Flags().BoolVar(&opts.SingleFile, "single-file", false, "toggle single file output")
 	cmd.Flags().StringVarP(&opts.Package, "package", "p", "", "package name used in generated Go code")
 	cmd.Flags().StringVar(&opts.CustomTypePackage, "custom-type-package", "", "Go package name to use for custom or unknown types")
-	cmd.Flags().StringArrayVar(&opts.IgnoreFields, "ignore-fields", nil, "fields to exclude from the generated Go code types")
-	cmd.Flags().StringArrayVar(&opts.IgnoreTables, "ignore-tables", nil, "tables to exclude from the generated Go code types")
+	cmd.Flags().StringArrayVar(&opts.TargetTables, "target-tables", nil, "tables to include from the generated Go code")
+	cmd.Flags().StringArrayVar(&opts.IgnoreFields, "ignore-fields", nil, "fields to exclude from the generated Go code")
+	cmd.Flags().StringArrayVar(&opts.IgnoreTables, "ignore-tables", nil, "tables to exclude from the generated Go code")
 	cmd.Flags().StringVar(&opts.TemplatePath, "template-path", "", "user supplied template path")
 	cmd.Flags().StringVar(&opts.Tags, "tags", "", "build tags to add to package header")
 	cmd.Flags().StringVar(&opts.InflectionRuleFile, "inflection-rule-file", "", "custom inflection rule file")

internal/argtype.go

@@ -51,6 +51,10 @@ type ArgType struct {
 	// CustomTypePackage is the Go package name to use for unknown types.
 	CustomTypePackage string
 
+	// TargetTables allows the user to specify table names which should be
+	// handled by yo in the generated code.
+	TargetTables []string
+
 	// IgnoreFields allows the user to specify field names which should not be
 	// handled by yo in the generated code.
 	IgnoreFields []string

internal/loader.go

@@ -100,6 +100,7 @@ func (tl *TypeLoader) LoadTable(args *ArgType) (map[string]*Type, error) {
 	for _, ti := range tableList {
 		ignore := false
 
+		// Ignore tables specified by IgnoreTables argument.
 		for _, ignoreTable := range args.IgnoreTables {
 			if ignoreTable == ti.TableName {
 				// Skip adding this table if user has specified they are not
@@ -112,6 +113,16 @@ func (tl *TypeLoader) LoadTable(args *ArgType) (map[string]*Type, error) {
 			}
 		}
 
+		// If the 'TargetTables' argument is passed, ignore any tables that are not specified in the array.
+		if len(args.TargetTables) != 0 {
+			ignore = true
+			for _, t := range args.TargetTables {
+				if t == ti.TableName {
+					ignore = false
+				}
+			}
+		}
+
 		if ignore {
 			continue
 		}