Repository: k1LoW/tbls
Pull request #409 (2023-01-19 15:08:52)

Support `--label` option for filtering tables


Patch:
diff --git a/cmd/doc.go b/cmd/doc.go
index cf0daab..0e5d56a 100644
--- a/cmd/doc.go
+++ b/cmd/doc.go
@@ -192,6 +192,7 @@ func loadDocArgs(args []string) ([]config.Option, error) {
 	options = append(options, config.BaseUrl(baseUrl))
 	options = append(options, config.Include(append(tables, includes...)))
 	options = append(options, config.Exclude(excludes))
+	options = append(options, config.IncludeLabels(labels))
 	if len(args) == 2 {
 		options = append(options, config.DSNURL(args[0]))
 		options = append(options, config.DocPath(args[1]))
@@ -232,6 +233,7 @@ func init() {
 	docCmd.Flags().StringSliceVarP(&tables, "table", "", []string{}, "target table (tables to include)")
 	docCmd.Flags().StringSliceVarP(&includes, "include", "", []string{}, "tables to include")
 	docCmd.Flags().StringSliceVarP(&excludes, "exclude", "", []string{}, "tables to exclude")
+	docCmd.Flags().StringSliceVarP(&labels, "label", "", []string{}, "table labels to be included")
 
 	if err := docCmd.MarkZshCompPositionalArgumentFile(2); err != nil {
 		_, _ = fmt.Fprintf(os.Stderr, "%s\n", err)
diff --git a/cmd/out.go b/cmd/out.go
index 08bfd34..f8ee81f 100644
--- a/cmd/out.go
+++ b/cmd/out.go
@@ -144,6 +144,7 @@ func loadOutArgs(args []string) ([]config.Option, error) {
 
 	options = append(options, config.Include(append(tables, includes...)))
 	options = append(options, config.Exclude(excludes))
+	options = append(options, config.IncludeLabels(labels))
 
 	if len(args) == 1 {
 		options = append(options, config.DSNURL(args[0]))
@@ -160,6 +161,7 @@ func init() {
 	outCmd.Flags().StringSliceVarP(&tables, "table", "", []string{}, "target table (tables to include)")
 	outCmd.Flags().StringSliceVarP(&includes, "include", "", []string{}, "tables to include")
 	outCmd.Flags().StringSliceVarP(&excludes, "exclude", "", []string{}, "tables to exclude")
+	outCmd.Flags().StringSliceVarP(&labels, "label", "", []string{}, "table labels to be included")
 	outCmd.Flags().IntVarP(&distance, "distance", "", 0, "distance between related tables to be displayed")
 	outCmd.Flags().StringVarP(&when, "when", "", "", "command execute condition")
 }
diff --git a/cmd/root.go b/cmd/root.go
index 2cb0c54..4a71229 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -65,6 +65,9 @@ var tables []string
 // tables to excludes
 var excludes []string
 
+// table labels to be included
+var labels []string
+
 const rootUsageTemplate = `Usage:{{if .Runnable}}{{if ne .UseLine "tbls [flags]" }}
   {{.UseLine}}{{end}}{{end}}{{if .HasAvailableSubCommands}}
   {{.CommandPath}} [command]{{end}}{{if gt (len .Aliases) 0}}
diff --git a/config/config.go b/config/config.go
index 075d842..a639c73 100644
--- a/config/config.go
+++ b/config/config.go
@@ -53,6 +53,10 @@ type Config struct {
 	RequiredVersion        string                 `yaml:"requiredVersion,omitempty"`
 	DisableOutputSchema    bool                   `yaml:"disableOutputSchema,omitempty"`
 	MergedDict             dict.Dict              `yaml:"-"`
+
+	// Table labels to be included
+	includeLabels []string
+
 	// Path of config file
 	Path string `yaml:"-"`
 	root string `yaml:"-"`
@@ -202,6 +206,16 @@ func Exclude(e []string) Option {
 	}
 }
 
+// IncludeLabels return Option set Config.IncludeLabels
+func IncludeLabels(l []string) Option {
+	return func(c *Config) error {
+		if len(l) > 0 {
+			c.includeLabels = l
+		}
+		return nil
+	}
+}
+
 // New return Config
 func New() (*Config, error) {
 	c := Config{}
@@ -400,21 +414,27 @@ func (c *Config) FilterTables(s *schema.Schema) error {
 	for _, t := range s.Tables {
 		li, mi := matchLength(i, t.Name)
 		le, me := matchLength(e, t.Name)
+		ml := matchLabels(c.includeLabels, t.Labels)
 		switch {
-		case len(c.Include) == 0:
+		case mi:
+			if me && li < le {
+				excludes = append(excludes, t)
+				continue
+			}
+			includes = append(includes, t)
+		case ml:
 			if me {
 				excludes = append(excludes, t)
 				continue
 			}
 			includes = append(includes, t)
-		case mi:
-			if me && li < le {
+		case len(c.Include) == 0 && len(c.includeLabels) == 0:
+			if me {
 				excludes = append(excludes, t)
 				continue
 			}
 			includes = append(includes, t)
 		default:
-			// c.Include > 0 && !mi
 			excludes = append(excludes, t)
 		}
 	}
@@ -652,6 +672,17 @@ func matchLength(s []string, e string) (int, bool) {
 	return 0, false
 }
 
+func matchLabels(il []string, l schema.Labels) bool {
+	for _, ll := range l {
+		for _, ill := range il {
+			if ll.Name == ill {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func match(s []string, e string) bool {
 	_, m := matchLength(s, e)
 	return m


Test patch:
diff --git a/config/config_test.go b/config/config_test.go
index 300f7af..e049dd2 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -186,38 +186,47 @@ func TestFilterTables(t *testing.T) {
 	tests := []struct {
 		include       []string
 		exclude       []string
+		labels        []string
 		distance      int
 		wantTables    int
 		wantRelations int
 	}{
-		{[]string{}, []string{}, 0, 5, 3},
-		{[]string{}, []string{"schema_migrations"}, 0, 4, 3},
-		{[]string{}, []string{"users"}, 0, 4, 1},
-		{[]string{"users"}, []string{}, 0, 1, 0},
-		{[]string{"user*"}, []string{}, 0, 2, 1},
-		{[]string{"*options"}, []string{}, 0, 1, 0},
-		{[]string{"*"}, []string{"user_options"}, 0, 4, 2},
-		{[]string{"not_exist"}, []string{}, 0, 0, 0},
-		{[]string{"not_exist", "*"}, []string{}, 0, 5, 3},
-		{[]string{"users"}, []string{"*"}, 0, 1, 0},
-		{[]string{"use*"}, []string{"use*"}, 0, 2, 1},
-		{[]string{"use*"}, []string{"user*"}, 0, 0, 0},
-		{[]string{"user*"}, []string{"user_*"}, 0, 1, 0},
-		{[]string{"*", "user*"}, []string{"user_*"}, 0, 4, 2},
+		{[]string{}, []string{}, []string{}, 0, 5, 3},
+		{[]string{}, []string{"schema_migrations"}, []string{}, 0, 4, 3},
+		{[]string{}, []string{"users"}, []string{}, 0, 4, 1},
+		{[]string{"users"}, []string{}, []string{}, 0, 1, 0},
+		{[]string{"user*"}, []string{}, []string{}, 0, 2, 1},
+		{[]string{"*options"}, []string{}, []string{}, 0, 1, 0},
+		{[]string{"*"}, []string{"user_options"}, []string{}, 0, 4, 2},
+		{[]string{"not_exist"}, []string{}, []string{}, 0, 0, 0},
+		{[]string{"not_exist", "*"}, []string{}, []string{}, 0, 5, 3},
+		{[]string{"users"}, []string{"*"}, []string{}, 0, 1, 0},
+		{[]string{"use*"}, []string{"use*"}, []string{}, 0, 2, 1},
+		{[]string{"use*"}, []string{"user*"}, []string{}, 0, 0, 0},
+		{[]string{"user*"}, []string{"user_*"}, []string{}, 0, 1, 0},
+		{[]string{"*", "user*"}, []string{"user_*"}, []string{}, 0, 4, 2},
 
-		{[]string{"users"}, []string{}, 1, 3, 2},
-		{[]string{"user_options"}, []string{}, 1, 2, 1},
-		{[]string{"user_options"}, []string{}, 2, 3, 2},
-		{[]string{"user_options"}, []string{}, 3, 4, 3},
-		{[]string{}, []string{}, 9, 5, 3},
-		{[]string{"posts"}, []string{}, 9, 4, 3},
-		{[]string{""}, []string{"*"}, 9, 0, 0},
+		{[]string{"users"}, []string{}, []string{}, 1, 3, 2},
+		{[]string{"user_options"}, []string{}, []string{}, 1, 2, 1},
+		{[]string{"user_options"}, []string{}, []string{}, 2, 3, 2},
+		{[]string{"user_options"}, []string{}, []string{}, 3, 4, 3},
+		{[]string{}, []string{}, []string{}, 9, 5, 3},
+		{[]string{"posts"}, []string{}, []string{}, 9, 4, 3},
+		{[]string{""}, []string{"*"}, []string{}, 9, 0, 0},
+
+		{[]string{}, []string{}, []string{"private"}, 0, 2, 1},
+		{[]string{}, []string{}, []string{"option"}, 0, 2, 0},
+		{[]string{}, []string{}, []string{"public", "private"}, 0, 4, 3},
+		{[]string{}, []string{"users"}, []string{"private"}, 0, 1, 0},
+		{[]string{}, []string{"user*"}, []string{"option"}, 0, 1, 0},
+		{[]string{"users"}, []string{}, []string{"private"}, 0, 2, 1},
 	}
 	for i, tt := range tests {
 		t.Run(fmt.Sprintf("%d.%v%v", i, tt.include, tt.exclude), func(t *testing.T) {
 			s := newSchemaForTestFilterTables(t)
 			c.Include = tt.include
 			c.Exclude = tt.exclude
+			c.includeLabels = tt.labels
 			c.Distance = tt.distance
 			err = c.FilterTables(s)
 			if err != nil {
diff --git a/testdata/filter_tables.json b/testdata/filter_tables.json
index 35b191a..92f54a5 100644
--- a/testdata/filter_tables.json
+++ b/testdata/filter_tables.json
@@ -72,7 +72,8 @@
         }
       ],
       "triggers": [],
-      "def": "CREATE TABLE users (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  username TEXT UNIQUE NOT NULL CHECK(length(username) \u003e 4)\n)"
+      "def": "CREATE TABLE users (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  username TEXT UNIQUE NOT NULL CHECK(length(username) \u003e 4)\n)",
+      "labels": [{"name":"private", "virtual": true}]
     },
     {
       "name": "categories",
@@ -132,7 +133,8 @@
         }
       ],
       "triggers": [],
-      "def": "CREATE TABLE categories (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  name TEXT UNIQUE NOT NULL\n)"
+      "def": "CREATE TABLE categories (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  name TEXT UNIQUE NOT NULL\n)",
+      "labels": [{"name":"public", "virtual": true}, {"name":"option", "virtual": true}]
     },
     {
       "name": "posts",
@@ -212,7 +214,8 @@
         }
       ],
       "triggers": [],
-      "def": "CREATE TABLE posts (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  user_id INTEGER NOT NULL,\n  category_id INTEGER NOT NULL,\n  title TEXT NOT NULL,\n  CONSTRAINT posts_user_id_fk FOREIGN KEY(user_id) REFERENCES users(id) MATCH NONE ON UPDATE NO ACTION ON DELETE CASCADE,\n  CONSTRAINT posts_category_id_fk FOREIGN KEY(category_id) REFERENCES categories(id) MATCH NONE ON UPDATE NO ACTION ON DELETE CASCADE\n)"
+      "def": "CREATE TABLE posts (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  user_id INTEGER NOT NULL,\n  category_id INTEGER NOT NULL,\n  title TEXT NOT NULL,\n  CONSTRAINT posts_user_id_fk FOREIGN KEY(user_id) REFERENCES users(id) MATCH NONE ON UPDATE NO ACTION ON DELETE CASCADE,\n  CONSTRAINT posts_category_id_fk FOREIGN KEY(category_id) REFERENCES categories(id) MATCH NONE ON UPDATE NO ACTION ON DELETE CASCADE\n)",
+      "labels": [{"name":"public", "virtual": true}]
     },
     {
       "name": "user_options",
@@ -271,7 +274,8 @@
         }
       ],
       "triggers": [],
-      "def": "CREATE TABLE user_options (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  user_id INTEGER NOT NULL,\n  email TEXT NOT NULL,\n  CONSTRAINT user_options_user_id_fk FOREIGN KEY(user_id) REFERENCES users(id) MATCH NONE ON UPDATE NO ACTION ON DELETE CASCADE\n)"
+      "def": "CREATE TABLE user_options (\n  id INTEGER PRIMARY KEY AUTOINCREMENT,\n  user_id INTEGER NOT NULL,\n  email TEXT NOT NULL,\n  CONSTRAINT user_options_user_id_fk FOREIGN KEY(user_id) REFERENCES users(id) MATCH NONE ON UPDATE NO ACTION ON DELETE CASCADE\n)",
+      "labels": [{"name":"private", "virtual": true}, {"name":"option", "virtual": true}]
     },
     {
       "name": "schema_migrations",