From f664bef1eb1bf7eb9e2f1aabef74b630a45e59e2 Mon Sep 17 00:00:00 2001
From: foghost <foghostor@gmail.com>
Date: Mon, 3 Apr 2023 20:55:07 +0900
Subject: [PATCH] Add target-tables argument

---
 cmd/root.go         |  5 +++--
 internal/argtype.go |  4 ++++
 internal/loader.go  | 11 +++++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 1a05ce4..cf1a8d2 100644
--- a/cmd/root.go
+++ b/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")
diff --git a/internal/argtype.go b/internal/argtype.go
index a9f05a0..28025e5 100644
--- a/internal/argtype.go
+++ b/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
diff --git a/internal/loader.go b/internal/loader.go
index 78a017d..f7357bb 100644
--- a/internal/loader.go
+++ b/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
 		}