feat[db]: allow exclude tables (#94)

* feat[db]: allow exclude tables Signed-off-by: rogerogers <[email protected]> * chore: remove print Signed-off-by: rogerogers <[email protected]> * chore: merge variable declarations and condition Signed-off-by: rogerogers <[email protected]> --------- Signed-off-by: rogerogers <[email protected]>
cloudwego · Nov 3, 2023 · 462b628 · 462b628
1 parent 3480534
commit 462b628
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 15 deletions.
diff --git a/.gitignore b/.gitignore
@@ -28,3 +28,5 @@ output/*
 # Vscode files
 .vscode
 
+# mac
+.DS_Store
diff --git a/cmd/dynamic/db.go b/cmd/dynamic/db.go
@@ -73,11 +73,16 @@ func parsePass(da *config.ModelArgument, pass string) error {
 	f.BoolVar(&da.FieldSignable, consts.Signable, false, "")
 	f.BoolVar(&da.FieldWithTypeTag, consts.TypeTag, false, "")
 	f.BoolVar(&da.FieldWithIndexTag, consts.IndexTag, false, "")
-	var tables utils.FlagStringSlice
+	var (
+		tables        utils.FlagStringSlice
+		excludeTables utils.FlagStringSlice
+	)
 	f.Var(&tables, consts.Tables, "")
+	f.Var(&excludeTables, consts.ExcludeTables, "")
 	if err := f.Parse(utils.StringSliceSpilt([]string{pass})); err != nil {
 		return err
 	}
 	da.Tables = tables
+	da.ExcludeTables = excludeTables
 	return nil
 }
diff --git a/cmd/static/model_flags.go b/cmd/static/model_flags.go
@@ -42,6 +42,7 @@ func modelFlags() []cli.Flag {
 		&cli.StringFlag{Name: consts.OutDir, Usage: "Specify output directory", Value: consts.DefaultDbOutDir, DefaultText: consts.DefaultDbOutDir},
 		&cli.StringFlag{Name: consts.OutFile, Usage: "Specify output filename", Value: consts.DefaultDbOutFile, DefaultText: consts.DefaultDbOutFile},
 		&cli.StringSliceFlag{Name: consts.Tables, Usage: "Specify databases tables"},
+		&cli.StringSliceFlag{Name: consts.ExcludeTables, Usage: "Specify exclude tables"},
 		&cli.BoolFlag{Name: consts.UnitTest, Usage: "Specify generate unit test", Value: false, DefaultText: "false"},
 		&cli.BoolFlag{Name: consts.OnlyModel, Usage: "Specify only generate model code", Value: false, DefaultText: "false"},
 		&cli.StringFlag{Name: consts.ModelPkgName, Usage: "Specify model package name", Value: "", DefaultText: ""},

diff --git a/config/model.go b/config/model.go
@@ -27,6 +27,7 @@ type ModelArgument struct {
 	DSN               string
 	Type              string
 	Tables            []string
+	ExcludeTables     []string
 	OnlyModel         bool
 	OutPath           string
 	OutFile           string
@@ -49,6 +50,7 @@ func (c *ModelArgument) ParseCli(ctx *cli.Context) error {
 	c.DSN = ctx.String(consts.DSN)
 	c.Type = strings.ToLower(ctx.String(consts.DBType))
 	c.Tables = ctx.StringSlice(consts.Tables)
+	c.ExcludeTables = ctx.StringSlice(consts.ExcludeTables)
 	c.OnlyModel = ctx.Bool(consts.OnlyModel)
 	c.OutPath = ctx.String(consts.OutDir)
 	c.OutFile = ctx.String(consts.OutFile)

diff --git a/pkg/consts/const.go b/pkg/consts/const.go
@@ -121,16 +121,17 @@ const (
 	Pass            = "pass"
 	ProtoSearchPath = "proto_search_path"
 
-	DSN          = "dsn"
-	DBType       = "db_type"
-	Tables       = "tables"
-	OnlyModel    = "only_model"
-	OutFile      = "out_file"
-	UnitTest     = "unittest"
-	ModelPkgName = "model_pkg"
-	Nullable     = "nullable"
-	Signable     = "signable"
-	IndexTag     = "index_tag"
-	TypeTag      = "type_tag"
-	HexTag       = "hex"
+	DSN           = "dsn"
+	DBType        = "db_type"
+	Tables        = "tables"
+	ExcludeTables = "exclude_tables"
+	OnlyModel     = "only_model"
+	OutFile       = "out_file"
+	UnitTest      = "unittest"
+	ModelPkgName  = "model_pkg"
+	Nullable      = "nullable"
+	Signable      = "signable"
+	IndexTag      = "index_tag"
+	TypeTag       = "type_tag"
+	HexTag        = "hex"
 )
diff --git a/pkg/model/model.go b/pkg/model/model.go
@@ -18,6 +18,7 @@ package model
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/cloudwego/cwgo/config"
 	"github.com/cloudwego/cwgo/pkg/consts"
@@ -33,15 +34,33 @@ func Model(c *config.ModelArgument) error {
 		return err
 	}
 
-	g := gen.NewGenerator(gen.Config{
+	genConfig := gen.Config{
 		OutPath:           c.OutPath,
 		OutFile:           c.OutFile,
 		ModelPkgPath:      c.ModelPkgName,
 		WithUnitTest:      c.WithUnitTest,
 		FieldNullable:     c.FieldNullable,
 		FieldSignable:     c.FieldSignable,
 		FieldWithIndexTag: c.FieldWithIndexTag,
-	})
+	}
+
+	if len(c.ExcludeTables) > 0 || c.Type == string(consts.Sqlite) {
+		genConfig.WithTableNameStrategy(func(tableName string) (targetTableName string) {
+			if c.Type == string(consts.Sqlite) && strings.HasPrefix(tableName, "sqlite") {
+				return ""
+			}
+			if len(c.ExcludeTables) > 0 {
+				for _, table := range c.ExcludeTables {
+					if tableName == table {
+						return ""
+					}
+				}
+			}
+			return tableName
+		})
+	}
+
+	g := gen.NewGenerator(genConfig)
 
 	g.UseDB(db)
-Original file line number
+Diff line change
@@ Expand Up / @@ -28,3 +28,5 @@ output/* @@
     # Vscode files
     .vscode
+    # mac
+    .DS_Store