feat: add gen.go

Author: Skyenought

cmd/static/model_flags.go

@@ -51,5 +51,15 @@ func modelFlags() []cli.Flag {
 		&cli.BoolFlag{Name: consts.TypeTag, Usage: "Specify generate field with gorm column type tag", Value: false, DefaultText: "false"},
 		&cli.BoolFlag{Name: consts.IndexTag, Usage: "Specify generate field with gorm index tag", Value: false, DefaultText: "false"},
 		&cli.StringFlag{Name: consts.SQLDir, Usage: "Specify a sql file or directory", Value: "", DefaultText: ""},
+		&cli.StringFlag{Name: consts.Mode, Usage: modeUsage, Value: "noctx,defaultquery,queryinterface", DefaultText: "noctx,defaultquery,queryinterface"},
 	}
 }
+
+const modeUsage = `Specify gorm/gen generator mode 
+(https://gorm.io/gen/dao.html#Generator-Modes).
+There is no need to follow strict case, we offer some abbreviations
+
+- WithoutContext (noctx)
+- WithDefaultQuery (defaultquery)
+- WithQueryInterface (queryinterface)
+`

config/model.go

@@ -38,6 +38,7 @@ type ModelArgument struct {
 	FieldWithIndexTag bool
 	FieldWithTypeTag  bool
 	SQLDir            string
+	Mode              string
 }
 
 func NewModelArgument() *ModelArgument {
@@ -62,5 +63,6 @@ func (c *ModelArgument) ParseCli(ctx *cli.Context) error {
 	c.FieldWithIndexTag = ctx.Bool(consts.IndexTag)
 	c.FieldWithTypeTag = ctx.Bool(consts.TypeTag)
 	c.SQLDir = ctx.String(consts.SQLDir)
+	c.Mode = ctx.String(consts.Mode)
 	return nil
 }

pkg/consts/const.go

@@ -159,6 +159,7 @@ const (
 	TypeTag       = "type_tag"
 	HexTag        = "hex"
 	SQLDir        = "sql_dir"
+	Mode          = "mode"
 )
 
 const (

pkg/model/generator_tpl.go

@@ -0,0 +1,91 @@
+package model
+
+import (
+	"bytes"
+	_ "embed"
+	"fmt"
+	"go/format"
+	"path/filepath"
+	"text/template"
+
+	"github.com/cloudwego/cwgo/config"
+	"github.com/cloudwego/cwgo/pkg/consts"
+
+	"gorm.io/gen"
+)
+
+//go:embed model_tpl.txt
+var mergedTemplate string
+
+type GenMethodTmpl struct {
+	GormOpen string
+	gen.Config
+	OnlyModel bool
+	// UseRawSQL indicates whether to use raw SQL for the database connection
+	UseRawSQL      bool
+	Tables         []string
+	StrategyParams struct {
+		ExcludeTables []string
+		Type          string
+	}
+}
+
+func execTmpl(c *config.ModelArgument) ([]byte, error) {
+	tpl, err := template.New("merged").Parse(mergedTemplate)
+	if err != nil {
+		return nil, fmt.Errorf("parse template fail: %w", err)
+	}
+
+	absOutPath, _ := filepath.Abs(c.OutPath)
+	data := GenMethodTmpl{
+		GormOpen:  buildGormOpen(c),
+		UseRawSQL: c.DSN == "" && c.SQLDir != "",
+		Tables:    c.Tables,
+		OnlyModel: c.OnlyModel,
+		Config: gen.Config{
+			OutPath:           absOutPath,
+			OutFile:           c.OutFile,
+			ModelPkgPath:      c.ModelPkgName,
+			WithUnitTest:      c.WithUnitTest,
+			FieldNullable:     c.FieldNullable,
+			FieldSignable:     c.FieldSignable,
+			FieldWithIndexTag: c.FieldWithIndexTag,
+			FieldWithTypeTag:  c.FieldWithTypeTag,
+			Mode:              buildGenMode(c.Mode),
+		},
+		StrategyParams: struct {
+			ExcludeTables []string
+			Type          string
+		}{
+			ExcludeTables: c.ExcludeTables,
+			Type:          c.Type,
+		},
+	}
+
+	var buf bytes.Buffer
+	err = tpl.Execute(&buf, data)
+	if err != nil {
+		return nil, fmt.Errorf("execute template fail: %w", err)
+	}
+	fmtCode, err := format.Source(buf.Bytes())
+	return fmtCode, err
+}
+
+func buildGormOpen(c *config.ModelArgument) string {
+	abs, _ := filepath.Abs(c.SQLDir)
+	switch {
+	case c.SQLDir != "":
+		return fmt.Sprintf(
+			"db, err := gorm.Open(rawsql.New(rawsql.Config{FilePath: []string{%q}}))",
+			abs,
+		)
+	case c.DSN != "" && c.Type != "":
+		return fmt.Sprintf(
+			"db, err := gorm.Open(%s.Open(%q))",
+			consts.DataBaseType(c.Type),
+			c.DSN,
+		)
+	default:
+		return ""
+	}
+}

pkg/model/model.go

@@ -18,6 +18,9 @@ package model
 
 import (
 	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
 	"strings"
 
 	"gorm.io/rawsql"
@@ -57,6 +60,7 @@ func Model(c *config.ModelArgument) error {
 		FieldSignable:     c.FieldSignable,
 		FieldWithIndexTag: c.FieldWithIndexTag,
 		FieldWithTypeTag:  c.FieldWithTypeTag,
+		Mode:              buildGenMode(c.Mode),
 	}
 
 	if len(c.ExcludeTables) > 0 || c.Type == string(consts.Sqlite) {
@@ -87,9 +91,12 @@ func Model(c *config.ModelArgument) error {
 	if !c.OnlyModel {
 		g.ApplyBasic(models...)
 	}
-
 	g.Execute()
-	return nil
+	getwd, _ := os.Getwd()
+	outPath := filepath.Join(getwd, c.OutPath)
+	genMainFileRootDir := filepath.Dir(outPath)
+	buf, err := execTmpl(c)
+	return os.WriteFile(filepath.Join(genMainFileRootDir, "gen.go"), buf, 0o644)
 }
 
 func genModels(g *gen.Generator, db *gorm.DB, tables []string) (models []interface{}, err error) {
@@ -109,3 +116,21 @@ func genModels(g *gen.Generator, db *gorm.DB, tables []string) (models []interfa
 	}
 	return models, nil
 }
+
+func buildGenMode(mode string) gen.GenerateMode {
+	var generateMode gen.GenerateMode
+	ss := regexp.MustCompile(`[,;\s]+`).Split(strings.ToLower(mode), -1)
+	for _, s := range ss {
+		s = strings.TrimSpace(s)
+		switch s {
+		case "withoutcontext", "noctx":
+			generateMode |= gen.WithoutContext
+		case "withdefaultquery", "defaultquery":
+			generateMode |= gen.WithDefaultQuery
+		case "withqueryinterface", "queryinterface":
+			generateMode |= gen.WithQueryInterface
+		default:
+		}
+	}
+	return generateMode
+}

pkg/model/model_tpl.txt

@@ -0,0 +1,88 @@
+{{ define "GORM_OPEN" }}
+    {{ .GormOpen }}
+    if err != nil {
+        log.Fatal("open db fail: ", err)
+    }
+{{ end }}
+
+{{ define "GEN_CONFIG" }}
+    genConfig := gen.Config{
+        OutPath:           {{ printf "%q" .OutPath }},
+        OutFile:           {{ printf "%q" .OutFile }},
+        ModelPkgPath:      {{ printf "%q" .ModelPkgPath }},
+        WithUnitTest:      {{ .WithUnitTest }},
+        FieldNullable:     {{ .FieldNullable }},
+        FieldSignable:     {{ .FieldSignable }},
+        FieldWithIndexTag: {{ .FieldWithIndexTag }},
+        FieldWithTypeTag:  {{ .FieldWithTypeTag }},
+        Mode:              {{ .Mode }},
+    }
+{{ end }}
+
+{{ define "TABLE_STRATEGY" }}
+    {{if or (gt (len .ExcludeTables) 0) (eq .Type "sqlite") }}
+    genConfig.WithTableNameStrategy(func(tableName string) string {
+        {{if eq .Type "sqlite" -}}
+        if strings.HasPrefix(tableName, "sqlite") {
+            return ""
+        }
+        {{end -}}
+        {{if gt (len .ExcludeTables) 0 -}}
+        switch tableName {
+            {{range $table := .ExcludeTables -}}
+            case "{{ $table }}":
+                return ""
+            {{end -}}
+        }
+        {{end -}}
+        return tableName
+    })
+    {{end}}
+{{ end }}
+
+package main
+
+import (
+    "fmt"
+    "log"
+
+    "gorm.io/gorm"
+    "gorm.io/gen"
+    {{if .UseRawSQL }} "gorm.io/rawsql" {{end}}
+)
+
+func main() {
+    {{ template "GORM_OPEN" . }}
+    {{ template "GEN_CONFIG" .Config }}
+    {{ template "TABLE_STRATEGY" .StrategyParams }}
+
+    g := gen.NewGenerator(genConfig)
+    g.UseDB(db)
+	models, err := genModels(g, db, []string{ {{- range $index, $element := .Tables }}
+	{{- if $index }}, {{ end }}"{{$element}}"{{- end }} })
+	if err != nil {
+		log.Fatal("gen models fail: ", err)
+	}
+	{{if not .OnlyModel}}
+		g.ApplyBasic(models...)
+	{{end}}
+	g.Execute()
+}
+
+func genModels(g *gen.Generator, db *gorm.DB, tables []string) (models []interface{}, err error) {
+	var tablesNameList []string
+	if len(tables) == 0 {
+		tablesNameList, err = db.Migrator().GetTables()
+		if err != nil {
+			return nil, fmt.Errorf("migrator get all tables fail: %w", err)
+		}
+	} else {
+		tablesNameList = tables
+	}
+
+	models = make([]interface{}, len(tablesNameList))
+	for i, tableName := range tablesNameList {
+		models[i] = g.GenerateModel(tableName)
+	}
+	return models, nil
+}