From cbb053786aa28c724e147e2a4029bfba0d65ac9c Mon Sep 17 00:00:00 2001 From: abc-valera Date: Fri, 24 Apr 2026 18:57:56 +0300 Subject: [PATCH] Added support for same package generation --- examples/samepackage/config.go | 7 ++ examples/samepackage/query.go | 19 ++++ examples/samepackage/struct.go | 10 ++ genconfig/config.go | 8 ++ internal/gen/gen.go | 13 ++- internal/gen/generator.go | 169 +++++++++++++++++++++++++++++++-- internal/gen/generator_test.go | 101 ++++++++++++++++++++ internal/gen/template.go | 17 ++-- 8 files changed, 323 insertions(+), 21 deletions(-) create mode 100644 examples/samepackage/config.go create mode 100644 examples/samepackage/query.go create mode 100644 examples/samepackage/struct.go diff --git a/examples/samepackage/config.go b/examples/samepackage/config.go new file mode 100644 index 0000000..693e599 --- /dev/null +++ b/examples/samepackage/config.go @@ -0,0 +1,7 @@ +package samepackage + +import "gorm.io/cli/gorm/genconfig" + +var _ = genconfig.Config{ + IsSamePackage: true, +} diff --git a/examples/samepackage/query.go b/examples/samepackage/query.go new file mode 100644 index 0000000..d6b9938 --- /dev/null +++ b/examples/samepackage/query.go @@ -0,0 +1,19 @@ +package samepackage + +import "time" + +// Example query interfaces with raw SQL annotations + +// IProductQuery demonstrates raw SQL query generation +type IProductQuery interface { + // SELECT * FROM @@table WHERE price > @minPrice + FindExpensiveProducts(minPrice float64) ([]Product, error) + + // UPDATE @@table + // {{set}} + // {{if discountFactor > 0}} price = price * @discountFactor, {{end}} + // {{if !beforeDate.IsZero()}} updated_at = NOW() {{end}} + // {{end}} + // WHERE created_at < @beforeDate + DiscountOldProducts(discountFactor float64, beforeDate time.Time) error +} diff --git a/examples/samepackage/struct.go b/examples/samepackage/struct.go new file mode 100644 index 0000000..abce64b --- /dev/null +++ b/examples/samepackage/struct.go @@ -0,0 +1,10 @@ +package samepackage + +import "time" + +type Product struct { + ID int + Name string + Price float64 + CreatedAt time.Time +} diff --git a/genconfig/config.go b/genconfig/config.go index eaca204..7701605 100644 --- a/genconfig/config.go +++ b/genconfig/config.go @@ -60,4 +60,12 @@ type Config struct { // ExcludeStructs is an optional blacklist for struct types to skip. // Applied after IncludeStructs filtering. Same selector rules as IncludeStructs. ExcludeStructs []any + + // IsSamePackage generates code in the same package as the source file. + IsSamePackage bool + + // SamePackageSuffix is appended to both generated file names and generated type names. + // Example: "gen" generates "model_gen.go" and "UserGen". + // Defaults to "gen" when IsSamePackage is true and no suffix is specified. + SamePackageSuffix string } diff --git a/internal/gen/gen.go b/internal/gen/gen.go index a24a7a6..94eeb30 100644 --- a/internal/gen/gen.go +++ b/internal/gen/gen.go @@ -10,16 +10,19 @@ var defaultOutPath = "./g" func New() *cobra.Command { var typed bool - var input, output string + var samePackage bool + var input, output, suffix string cmd := &cobra.Command{ Use: "gen", Short: "Generate GORM query code from raw SQL interfaces", RunE: func(cmd *cobra.Command, args []string) error { g := Generator{ - Typed: typed, - Files: map[string]*File{}, - outPath: output, + Typed: typed, + Files: map[string]*File{}, + outPath: output, + samePackage: samePackage, + suffix: suffix, } err := g.Process(input) @@ -39,6 +42,8 @@ func New() *cobra.Command { cmd.Flags().BoolVarP(&typed, "typed", "t", true, "Generated Typed API") cmd.Flags().StringVarP(&output, "output", "o", defaultOutPath, "Directory to place generated code") cmd.Flags().StringVarP(&input, "input", "i", "", "Path to Go interface file with raw SQL annotations") + cmd.Flags().BoolVar(&samePackage, "is-same-package", false, "Generate code in the same package as the source") + cmd.Flags().StringVar(&suffix, "same-package-suffix", "", "Suffix to append to generated file names and type names (e.g., 'gen' to get 'model_gen.go' and 'UserGen')") cobra.CheckErr(cmd.MarkFlagRequired("input")) return cmd diff --git a/internal/gen/generator.go b/internal/gen/generator.go index f67eca0..21cf906 100644 --- a/internal/gen/generator.go +++ b/internal/gen/generator.go @@ -23,9 +23,11 @@ import ( type ( Generator struct { - Typed bool - Files map[string]*File - outPath string + Typed bool + Files map[string]*File + outPath string + samePackage bool + suffix string } File struct { ToPackage string @@ -51,6 +53,7 @@ type ( Doc string Methods []*Method err error + file *File } Method struct { Name string @@ -68,6 +71,7 @@ type ( Name string Doc string Fields []Field + file *File } Field struct { Name string @@ -197,8 +201,43 @@ func (g *Generator) Gen() error { continue } - outPath = filepath.Join(outPath, file.relPath) - file.ToPackage = filepath.Base(filepath.Dir(outPath)) + // Determine output file path and check if config overrides are available. + // Config values override defaults but CLI flags take precedence. + samePackage := g.samePackage + suffix := g.suffix + for _, cfg := range file.applicableConfigs { + if !g.samePackage && cfg.IsSamePackage { + samePackage = cfg.IsSamePackage + } + if g.suffix == "" && cfg.SamePackageSuffix != "" { + suffix = cfg.SamePackageSuffix + } + } + + if samePackage { + // Generate in the same directory as the source file + dir := filepath.Dir(file.inputPath) + baseName := filepath.Base(file.inputPath) + ext := filepath.Ext(baseName) + nameWithoutExt := strings.TrimSuffix(baseName, ext) + + // Use the default suffix value if not provided + if suffix == "" { + suffix = "gen" + } + nameWithoutExt = nameWithoutExt + "_" + suffix + + outPath = filepath.Join(dir, nameWithoutExt+ext) + file.ToPackage = file.Package + } else { + outPath = filepath.Join(outPath, file.relPath) + file.ToPackage = filepath.Base(filepath.Dir(outPath)) + } + + // Update file's generator suffix for template access + if suffix != "" { + file.Generator.suffix = suffix + } var results bytes.Buffer if err := tmpl.Execute(&results, file); err != nil { @@ -281,11 +320,69 @@ func (p Import) ImportPath() string { return fmt.Sprintf("%s %q", p.Name, p.Path) } +// TemplateImports omits self imports for the same package generation. +func (f *File) TemplateImports() []Import { + var filtered []Import + for _, imp := range f.Imports { + // Skip self-import when generating in the same package + if f.ToPackage == f.Package && imp.Path == f.PackagePath { + continue + } + filtered = append(filtered, imp) + } + return filtered +} + // GoFullType returns the complete Go type string for a parameter func (p Param) GoFullType() string { return p.Type } +func (m Method) stripSamePackagePrefix(typeName string) string { + if m.Interface.file == nil { + return typeName + } + + file := m.Interface.file + // Only strip prefix when generating in the same package + if file.ToPackage != file.Package || file.PackagePath == "" { + return typeName + } + + // Strip the package prefix if it matches the current package + // Handle both short name (e.g., "samepackage.Product") and full path prefixes + packagePrefix := file.Package + "." + fullPathPrefix := file.PackagePath + "." + + // Try to strip short package prefix first + if after, found := strings.CutPrefix(typeName, packagePrefix); found { + return after + } + + // Try to strip full path prefix + if after, found := strings.CutPrefix(typeName, fullPathPrefix); found { + return after + } + + // Handle slice types: []packagename.Type -> []Type + if strings.HasPrefix(typeName, "[]") { + stripped := m.stripSamePackagePrefix(typeName[2:]) + if stripped != typeName[2:] { + return "[]" + stripped + } + } + + // Handle pointer types: *packagename.Type -> *Type + if strings.HasPrefix(typeName, "*") { + stripped := m.stripSamePackagePrefix(typeName[1:]) + if stripped != typeName[1:] { + return "*" + stripped + } + } + + return typeName +} + // ParamsString formats method parameters as a string for code generation func (m Method) ParamsString() string { var parts []string @@ -297,7 +394,8 @@ func (m Method) ParamsString() string { p.Name = "ctx" } - parts = append(parts, fmt.Sprintf("%s %s", p.Name, p.GoFullType())) + typeStr := m.stripSamePackagePrefix(p.GoFullType()) + parts = append(parts, fmt.Sprintf("%s %s", p.Name, typeStr)) } if !hasCtx { @@ -312,12 +410,13 @@ func (m Method) ResultString() string { if m.SQL.Raw != "" { var rets []string for _, r := range m.Result { - rets = append(rets, r.GoFullType()) + typeStr := m.stripSamePackagePrefix(r.GoFullType()) + rets = append(rets, typeStr) } return strings.Join(rets, ", ") } - return fmt.Sprintf("%sInterface[T]", m.Interface.IfaceName) + return fmt.Sprintf("%s[T]", m.Interface.GeneratedInterfaceName()) } // Body generates the method body code for templates @@ -350,7 +449,7 @@ return e.Exec(ctx, sb.String(), _params...)`, sqlSnippet) return fmt.Sprintf(`%s var result %s err := e.Raw(sb.String(), _params...).Scan(ctx, &result) -return result, err`, sqlSnippet, m.Result[0].GoFullType()) +return result, err`, sqlSnippet, m.stripSamePackagePrefix(m.Result[0].GoFullType())) } // chainMethodBody generates method body for chaining SQL operations that return interface @@ -470,6 +569,50 @@ func (f Field) Type() string { return fmt.Sprintf("field.Field[%s]", filepath.Base(goType)) } +// capitalizedSuffix returns the suffix with its first letter uppercased. +func capitalizedSuffix(suffix string) string { + if suffix == "" { + return "" + } + r := []rune(suffix) + r[0] = unicode.ToUpper(r[0]) + return string(r) +} + +// GeneratedName appends the suffix to the struct name if configured +func (s Struct) GeneratedName() string { + if s.file != nil && s.file.Generator.suffix != "" { + return s.Name + capitalizedSuffix(s.file.Generator.suffix) + } + return s.Name +} + +// GeneratedConstructorName appends the suffix to the constructor func if configured +func (i Interface) GeneratedConstructorName() string { + if i.file != nil && i.file.Generator.suffix != "" { + return i.Name + capitalizedSuffix(i.file.Generator.suffix) + } + return i.Name +} + +// GeneratedInterfaceName appends the suffix to the generated interface name if configured +func (i Interface) GeneratedInterfaceName() string { + if i.file != nil && i.file.Generator.suffix != "" { + baseName := strings.TrimPrefix(i.IfaceName, "_") + return "_" + baseName + capitalizedSuffix(i.file.Generator.suffix) + "Interface" + } + return i.IfaceName + "Interface" +} + +// GeneratedImplName appends the suffix to the generated implementation struct name if configured +func (i Interface) GeneratedImplName() string { + if i.file != nil && i.file.Generator.suffix != "" { + baseName := strings.TrimPrefix(i.IfaceName, "_") + return "_" + baseName + capitalizedSuffix(i.file.Generator.suffix) + } + return i.IfaceName + "Impl" +} + // Value returns the field value string with column name for template generation func (f Field) Value() string { fieldType := f.Type() @@ -599,6 +742,12 @@ func (p *File) parseConfigLiteral(cl *ast.CompositeLit) *genconfig.Config { cfg.IncludeStructs = append(cfg.IncludeStructs, collect(kv.Value)...) case "ExcludeStructs": cfg.ExcludeStructs = append(cfg.ExcludeStructs, collect(kv.Value)...) + case "IsSamePackage": + if ident, ok := kv.Value.(*ast.Ident); ok { + cfg.IsSamePackage = ident.Name == "true" + } + case "SamePackageSuffix": + cfg.SamePackageSuffix = strLit(kv.Value) } } return cfg @@ -610,6 +759,7 @@ func (p *File) processInterfaceType(n *ast.TypeSpec, data *ast.InterfaceType) In Name: n.Name.Name, IfaceName: "_" + n.Name.Name, Doc: n.Doc.Text(), + file: p, } methods := data.Methods.List @@ -659,6 +809,7 @@ func (p *File) processInterfaceType(n *ast.TypeSpec, data *ast.InterfaceType) In func (p *File) processStructType(typeSpec *ast.TypeSpec, data *ast.StructType, pkgName string) Struct { s := Struct{ Name: typeSpec.Name.Name, + file: p, } for _, field := range data.Fields.List { diff --git a/internal/gen/generator_test.go b/internal/gen/generator_test.go index 8534990..ab32a71 100644 --- a/internal/gen/generator_test.go +++ b/internal/gen/generator_test.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "reflect" + "strings" "testing" "text/template" ) @@ -131,6 +132,106 @@ type Entity interface { } } +func TestSamePackageWithSuffixes(t *testing.T) { + samePackageDir, err := filepath.Abs("../../examples/samepackage") + if err != nil { + t.Fatalf("failed to get absolute path: %v", err) + } + + tests := []struct { + name string + cliSuffix string + wantSuffix string // suffix in generated file names, e.g. "gen" → "query_gen.go" + wantTypeSuffix string // suffix appended to generated type names, e.g. "Gen" → "IProductQueryGen" + }{ + { + name: "default suffix from config", + cliSuffix: "", + wantSuffix: "gen", + wantTypeSuffix: "Gen", + }, + { + name: "cli suffix overrides config", + cliSuffix: "custom", + wantSuffix: "custom", + wantTypeSuffix: "Custom", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + g := &Generator{ + Files: map[string]*File{}, + outPath: defaultOutPath, + samePackage: false, + suffix: tc.cliSuffix, + Typed: true, + } + + if err := g.Process(samePackageDir); err != nil { + t.Fatalf("Process error: %v", err) + } + if err := g.Gen(); err != nil { + t.Fatalf("Gen error: %v", err) + } + + // Collect generated files and register cleanup. + generatedFiles, err := filepath.Glob(filepath.Join(samePackageDir, fmt.Sprintf("*_%s.go", tc.wantSuffix))) + if err != nil { + t.Fatalf("glob error: %v", err) + } + t.Cleanup(func() { + for _, f := range generatedFiles { + os.Remove(f) + } + }) + + if len(generatedFiles) == 0 { + t.Fatalf("no files generated with suffix _%s in %s", tc.wantSuffix, samePackageDir) + } + + for _, genFile := range generatedFiles { + content, err := os.ReadFile(genFile) + if err != nil { + t.Fatalf("failed to read generated file %s: %v", genFile, err) + } + + if _, err := parser.ParseFile(token.NewFileSet(), genFile, content, parser.AllErrors); err != nil { + t.Errorf("generated file %s has invalid Go syntax: %v", genFile, err) + } + + contentStr := string(content) + + // File must be in the same package as the source. + if !strings.Contains(contentStr, "package samepackage") { + t.Errorf("generated file %s should declare package samepackage", filepath.Base(genFile)) + } + + // Self-package import must be absent (same-package generation). + if strings.Contains(contentStr, `"gorm.io/cli/gorm/examples/samepackage"`) { + t.Errorf("generated file %s must not import its own package", filepath.Base(genFile)) + } + + // Type references must be un-prefixed (e.g. "[]Product", not "[]samepackage.Product"). + if strings.Contains(contentStr, "samepackage.") { + t.Errorf("generated file %s must not reference own package by name, got: %s", filepath.Base(genFile), contentStr) + } + } + + // Verify that the query file's generated constructor uses the expected suffix. + queryGen := filepath.Join(samePackageDir, fmt.Sprintf("query_%s.go", tc.wantSuffix)) + queryContent, err := os.ReadFile(queryGen) + if err != nil { + t.Fatalf("expected generated query file %s: %v", queryGen, err) + } + expectedConstructor := fmt.Sprintf("func IProductQuery%s[", tc.wantTypeSuffix) + if !strings.Contains(string(queryContent), expectedConstructor) { + t.Errorf("expected constructor %q in %s, got:\n%s", expectedConstructor, queryGen, string(queryContent)) + } + }) + } +} + func TestProcessStructType(t *testing.T) { fileset := token.NewFileSet() file, err := parser.ParseFile(fileset, "../../examples/models/user.go", nil, parser.AllErrors) diff --git a/internal/gen/template.go b/internal/gen/template.go index 2ce4988..762599b 100644 --- a/internal/gen/template.go +++ b/internal/gen/template.go @@ -13,39 +13,40 @@ import ( {{- if .UsedTypedAPI }} "gorm.io/cli/gorm/typed" {{- end }} - {{range .Imports -}} + {{range .TemplateImports -}} {{.ImportPath}} {{end -}} ) {{range .Interfaces}} -{{$IfaceName := .IfaceName}} -func {{.Name}}[T any](db *gorm.DB, opts ...clause.Expression) {{$IfaceName}}Interface[T] { - return {{$IfaceName}}Impl[T]{ +{{$InterfaceName := .GeneratedInterfaceName}} +{{$ImplName := .GeneratedImplName}} +func {{.GeneratedConstructorName}}[T any](db *gorm.DB, opts ...clause.Expression) {{$InterfaceName}}[T] { + return {{$ImplName}}[T]{ Interface: {{if $.UsedTypedAPI}}typed{{else}}gorm{{end}}.G[T](db, opts...), } } -type {{$IfaceName}}Interface[T any] interface { +type {{$InterfaceName}}[T any] interface { {{if $.UsedTypedAPI}}typed{{else}}gorm{{end}}.Interface[T] {{range .Methods -}} {{.Name}}({{.ParamsString}}) ({{.ResultString}}) {{end}} } -type {{$IfaceName}}Impl[T any] struct { +type {{$ImplName}}[T any] struct { {{if $.UsedTypedAPI}}typed{{else}}gorm{{end}}.Interface[T] } {{range .Methods}} -func (e {{$IfaceName}}Impl[T]) {{.Name}}({{.ParamsString}}) ({{.ResultString}}) { +func (e {{$ImplName}}[T]) {{.Name}}({{.ParamsString}}) ({{.ResultString}}) { {{.Body}} } {{end}} {{end}} {{range .Structs}} -var {{.Name}} = struct { +var {{.GeneratedName}} = struct { {{range .Fields -}} {{.Name}} {{.Type}} {{end}}