feat(codegen/golang): allow exporting models file to separate package

berk-karaal · berk-karaal · commit d248d0ff2f6c · 2025-03-05T23:53:31.000+03:00
diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go
@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"go/format"
+	"path/filepath"
 	"strings"
 	"text/template"
 
@@ -122,7 +123,7 @@ func Generate(ctx context.Context, req *plugin.GenerateRequest) (*plugin.Generat
 	}
 
 	if options.OmitUnusedStructs {
-		enums, structs = filterUnusedStructs(enums, structs, queries)
+		enums, structs = filterUnusedStructs(options, enums, structs, queries)
 	}
 
 	if err := validate(options, enums, structs, queries); err != nil {
@@ -211,6 +212,7 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 		"imports":    i.Imports,
 		"hasImports": i.HasImports,
 		"hasPrefix":  strings.HasPrefix,
+		"trimPrefix": strings.TrimPrefix,
 
 		// These methods are Go specific, they do not belong in the codegen package
 		// (as that is language independent)
@@ -232,14 +234,15 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 
 	output := map[string]string{}
 
-	execute := func(name, templateName string) error {
+	execute := func(name, packageName, templateName string) error {
 		imports := i.Imports(name)
 		replacedQueries := replaceConflictedArg(imports, queries)
 
 		var b bytes.Buffer
 		w := bufio.NewWriter(&b)
 		tctx.SourceName = name
 		tctx.GoQueries = replacedQueries
+		tctx.Package = packageName
 		err := tmpl.ExecuteTemplate(w, templateName, &tctx)
 		w.Flush()
 		if err != nil {
@@ -251,8 +254,13 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 			return fmt.Errorf("source error: %w", err)
 		}
 
-		if templateName == "queryFile" && options.OutputFilesSuffix != "" {
-			name += options.OutputFilesSuffix
+		if templateName == "queryFile" {
+			if options.OutputQueryFilesDirectory != "" {
+				name = filepath.Join(options.OutputQueryFilesDirectory, name)
+			}
+			if options.OutputFilesSuffix != "" {
+				name += options.OutputFilesSuffix
+			}
 		}
 
 		if !strings.HasSuffix(name, ".go") {
@@ -284,24 +292,29 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 		batchFileName = options.OutputBatchFileName
 	}
 
-	if err := execute(dbFileName, "dbFile"); err != nil {
+	modelsPackageName := options.Package
+	if options.OutputModelsPackage != "" {
+		modelsPackageName = options.OutputModelsPackage
+	}
+
+	if err := execute(dbFileName, options.Package, "dbFile"); err != nil {
 		return nil, err
 	}
-	if err := execute(modelsFileName, "modelsFile"); err != nil {
+	if err := execute(modelsFileName, modelsPackageName, "modelsFile"); err != nil {
 		return nil, err
 	}
 	if options.EmitInterface {
-		if err := execute(querierFileName, "interfaceFile"); err != nil {
+		if err := execute(querierFileName, options.Package, "interfaceFile"); err != nil {
 			return nil, err
 		}
 	}
 	if tctx.UsesCopyFrom {
-		if err := execute(copyfromFileName, "copyfromFile"); err != nil {
+		if err := execute(copyfromFileName, options.Package, "copyfromFile"); err != nil {
 			return nil, err
 		}
 	}
 	if tctx.UsesBatch {
-		if err := execute(batchFileName, "batchFile"); err != nil {
+		if err := execute(batchFileName, options.Package, "batchFile"); err != nil {
 			return nil, err
 		}
 	}
@@ -312,7 +325,7 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 	}
 
 	for source := range files {
-		if err := execute(source, "queryFile"); err != nil {
+		if err := execute(source, options.Package, "queryFile"); err != nil {
 			return nil, err
 		}
 	}
@@ -362,7 +375,7 @@ func checkNoTimesForMySQLCopyFrom(queries []Query) error {
 	return nil
 }
 
-func filterUnusedStructs(enums []Enum, structs []Struct, queries []Query) ([]Enum, []Struct) {
+func filterUnusedStructs(options *opts.Options, enums []Enum, structs []Struct, queries []Query) ([]Enum, []Struct) {
 	keepTypes := make(map[string]struct{})
 
 	for _, query := range queries {
@@ -389,16 +402,23 @@ func filterUnusedStructs(enums []Enum, structs []Struct, queries []Query) ([]Enu
 
 	keepEnums := make([]Enum, 0, len(enums))
 	for _, enum := range enums {
-		_, keep := keepTypes[enum.Name]
-		_, keepNull := keepTypes["Null"+enum.Name]
+		var enumType string
+		if options.ModelsPackageImportPath != "" {
+			enumType = options.OutputModelsPackage + "." + enum.Name
+		} else {
+			enumType = enum.Name
+		}
+
+		_, keep := keepTypes[enumType]
+		_, keepNull := keepTypes["Null"+enumType]
 		if keep || keepNull {
 			keepEnums = append(keepEnums, enum)
 		}
 	}
 
 	keepStructs := make([]Struct, 0, len(structs))
 	for _, st := range structs {
-		if _, ok := keepTypes[st.Name]; ok {
+		if _, ok := keepTypes[st.Type()]; ok {
 			keepStructs = append(keepStructs, st)
 		}
 	}
diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go
@@ -160,7 +160,7 @@ var pqtypeTypes = map[string]struct{}{
 	"pqtype.NullRawMessage": {},
 }
 
-func buildImports(options *opts.Options, queries []Query, uses func(string) bool) (map[string]struct{}, map[ImportSpec]struct{}) {
+func buildImports(options *opts.Options, queries []Query, outputFile OutputFile, uses func(string) bool) (map[string]struct{}, map[ImportSpec]struct{}) {
 	pkg := make(map[ImportSpec]struct{})
 	std := make(map[string]struct{})
 
@@ -243,11 +243,52 @@ func buildImports(options *opts.Options, queries []Query, uses func(string) bool
 		}
 	}
 
+	requiresModelsPackageImport := func() bool {
+		if options.ModelsPackageImportPath == "" {
+			return false
+		}
+
+		for _, q := range queries {
+			// Check if the return type is from models package (possibly a model struct or an enum)
+			if q.hasRetType() && strings.HasPrefix(q.Ret.Type(), options.OutputModelsPackage+".") {
+				return true
+			}
+
+			// Check if the return type struct contains a type from models package (possibly an enum field or an embedded struct)
+			if outputFile != OutputFileInterface && q.hasRetType() && q.Ret.IsStruct() {
+				for _, f := range q.Ret.Struct.Fields {
+					if strings.HasPrefix(f.Type, options.OutputModelsPackage+".") {
+						return true
+					}
+				}
+			}
+
+			// Check if the argument type is from models package (possibly an enum)
+			if !q.Arg.isEmpty() && strings.HasPrefix(q.Arg.Type(), options.OutputModelsPackage+".") {
+				return true
+			}
+
+			// Check if the argument struct contains a type from models package (possibly an enum field)
+			if outputFile != OutputFileInterface && !q.Arg.isEmpty() && q.Arg.IsStruct() {
+				for _, f := range q.Arg.Struct.Fields {
+					if strings.HasPrefix(f.Type, options.OutputModelsPackage+".") {
+						return true
+					}
+				}
+			}
+
+		}
+		return false
+	}
+	if requiresModelsPackageImport() {
+		pkg[ImportSpec{Path: options.ModelsPackageImportPath}] = struct{}{}
+	}
+
 	return std, pkg
 }
 
 func (i *importer) interfaceImports() fileImports {
-	std, pkg := buildImports(i.Options, i.Queries, func(name string) bool {
+	std, pkg := buildImports(i.Options, i.Queries, OutputFileInterface, func(name string) bool {
 		for _, q := range i.Queries {
 			if q.hasRetType() {
 				if usesBatch([]Query{q}) {
@@ -272,7 +313,7 @@ func (i *importer) interfaceImports() fileImports {
 }
 
 func (i *importer) modelImports() fileImports {
-	std, pkg := buildImports(i.Options, nil, i.usesType)
+	std, pkg := buildImports(i.Options, nil, OutputFileModel, i.usesType)
 
 	if len(i.Enums) > 0 {
 		std["fmt"] = struct{}{}
@@ -311,7 +352,7 @@ func (i *importer) queryImports(filename string) fileImports {
 		}
 	}
 
-	std, pkg := buildImports(i.Options, gq, func(name string) bool {
+	std, pkg := buildImports(i.Options, gq, OutputFileQuery, func(name string) bool {
 		for _, q := range gq {
 			if q.hasRetType() {
 				if q.Ret.EmitStruct() {
@@ -412,7 +453,7 @@ func (i *importer) copyfromImports() fileImports {
 			copyFromQueries = append(copyFromQueries, q)
 		}
 	}
-	std, pkg := buildImports(i.Options, copyFromQueries, func(name string) bool {
+	std, pkg := buildImports(i.Options, copyFromQueries, OutputFileCopyfrom, func(name string) bool {
 		for _, q := range copyFromQueries {
 			if q.hasRetType() {
 				if strings.HasPrefix(q.Ret.Type(), name) {
@@ -447,7 +488,7 @@ func (i *importer) batchImports() fileImports {
 			batchQueries = append(batchQueries, q)
 		}
 	}
-	std, pkg := buildImports(i.Options, batchQueries, func(name string) bool {
+	std, pkg := buildImports(i.Options, batchQueries, OutputFileBatch, func(name string) bool {
 		for _, q := range batchQueries {
 			if q.hasRetType() {
 				if q.Ret.EmitStruct() {
diff --git a/internal/codegen/golang/opts/options.go b/internal/codegen/golang/opts/options.go
@@ -35,8 +35,11 @@ type Options struct {
 	OutputBatchFileName         string            `json:"output_batch_file_name,omitempty" yaml:"output_batch_file_name"`
 	OutputDbFileName            string            `json:"output_db_file_name,omitempty" yaml:"output_db_file_name"`
 	OutputModelsFileName        string            `json:"output_models_file_name,omitempty" yaml:"output_models_file_name"`
+	OutputModelsPackage         string            `json:"output_models_package,omitempty" yaml:"output_models_package"`
+	ModelsPackageImportPath     string            `json:"models_package_import_path,omitempty" yaml:"models_package_import_path"`
 	OutputQuerierFileName       string            `json:"output_querier_file_name,omitempty" yaml:"output_querier_file_name"`
 	OutputCopyfromFileName      string            `json:"output_copyfrom_file_name,omitempty" yaml:"output_copyfrom_file_name"`
+	OutputQueryFilesDirectory   string            `json:"output_query_files_directory,omitempty" yaml:"output_query_files_directory"`
 	OutputFilesSuffix           string            `json:"output_files_suffix,omitempty" yaml:"output_files_suffix"`
 	InflectionExcludeTableNames []string          `json:"inflection_exclude_table_names,omitempty" yaml:"inflection_exclude_table_names"`
 	QueryParameterLimit         *int32            `json:"query_parameter_limit,omitempty" yaml:"query_parameter_limit"`
@@ -150,6 +153,11 @@ func ValidateOpts(opts *Options) error {
 	if *opts.QueryParameterLimit < 0 {
 		return fmt.Errorf("invalid options: query parameter limit must not be negative")
 	}
-
+	if opts.OutputModelsPackage != "" && opts.ModelsPackageImportPath == "" {
+		return fmt.Errorf("invalid options: models_package_import_path must be set when output_models_package is used")
+	}
+	if opts.ModelsPackageImportPath != "" && opts.OutputModelsPackage == "" {
+		return fmt.Errorf("invalid options: output_models_package must be set when models_package_import_path is used")
+	}
 	return nil
 }
diff --git a/internal/codegen/golang/output_file.go b/internal/codegen/golang/output_file.go
@@ -0,0 +1,12 @@
+package golang
+
+type OutputFile string
+
+const (
+	OutputFileModel     OutputFile = "modelFile"
+	OutputFileQuery     OutputFile = "queryFile"
+	OutputFileDb        OutputFile = "dbFile"
+	OutputFileInterface OutputFile = "interfaceFile"
+	OutputFileCopyfrom  OutputFile = "copyfromFile"
+	OutputFileBatch     OutputFile = "batchFile"
+)
diff --git a/internal/codegen/golang/postgresql_type.go b/internal/codegen/golang/postgresql_type.go
@@ -571,17 +571,24 @@ func postgresType(req *plugin.GenerateRequest, options *opts.Options, col *plugi
 
 			for _, enum := range schema.Enums {
 				if rel.Name == enum.Name && rel.Schema == schema.Name {
+					enumName := ""
 					if notNull {
 						if schema.Name == req.Catalog.DefaultSchema {
-							return StructName(enum.Name, options)
+							enumName = StructName(enum.Name, options)
+						} else {
+							enumName = StructName(schema.Name+"_"+enum.Name, options)
 						}
-						return StructName(schema.Name+"_"+enum.Name, options)
 					} else {
 						if schema.Name == req.Catalog.DefaultSchema {
-							return "Null" + StructName(enum.Name, options)
+							enumName = "Null" + StructName(enum.Name, options)
+						} else {
+							enumName = "Null" + StructName(schema.Name+"_"+enum.Name, options)
 						}
-						return "Null" + StructName(schema.Name+"_"+enum.Name, options)
 					}
+					if options.ModelsPackageImportPath != "" {
+						return options.OutputModelsPackage + "." + enumName
+					}
+					return enumName
 				}
 			}
 
diff --git a/internal/codegen/golang/query.go b/internal/codegen/golang/query.go
@@ -88,7 +88,7 @@ func (v QueryValue) Type() string {
 		return v.Typ
 	}
 	if v.Struct != nil {
-		return v.Struct.Name
+		return v.Struct.Type()
 	}
 	panic("no type for QueryValue: " + v.Name)
 }
diff --git a/internal/codegen/golang/result.go b/internal/codegen/golang/result.go
@@ -83,6 +83,7 @@ func buildStructs(req *plugin.GenerateRequest, options *opts.Options) []Struct {
 			s := Struct{
 				Table:   &plugin.Identifier{Schema: schema.Name, Name: table.Rel.Name},
 				Name:    StructName(structName, options),
+				Package: options.OutputModelsPackage,
 				Comment: table.Comment,
 			}
 			for _, column := range table.Columns {
@@ -146,7 +147,7 @@ func newGoEmbed(embed *plugin.Identifier, structs []Struct, defaultSchema string
 		}
 
 		return &goEmbed{
-			modelType: s.Name,
+			modelType: s.Type(),
 			modelName: s.Name,
 			fields:    fields,
 		}
diff --git a/internal/codegen/golang/struct.go b/internal/codegen/golang/struct.go
@@ -12,10 +12,18 @@ import (
 type Struct struct {
 	Table   *plugin.Identifier
 	Name    string
+	Package string
 	Fields  []Field
 	Comment string
 }
 
+func (s Struct) Type() string {
+	if s.Package != "" {
+		return s.Package + "." + s.Name
+	}
+	return s.Name
+}
+
 func StructName(name string, options *opts.Options) string {
 	if rename := options.Rename[name]; rename != "" {
 		return rename
diff --git a/internal/codegen/golang/templates/template.tmpl b/internal/codegen/golang/templates/template.tmpl
@@ -156,7 +156,7 @@ type {{.Name}} struct { {{- range .Fields}}
   {{- if .Comment}}
   {{comment .Comment}}{{else}}
   {{- end}}
-  {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
+  {{.Name}} {{trimPrefix .Type (printf "%s%s" $.Package ".") }} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
   {{- end}}
 }
 {{end}}

Original file line number	Diff line number	Diff line change
`@@ -571,17 +571,24 @@ func postgresType(req plugin.GenerateRequest, options opts.Options, col *plugi`
`571`	`571`
`572`	`572`	`for _, enum := range schema.Enums {`
`573`	`573`	`if rel.Name == enum.Name && rel.Schema == schema.Name {`
	`574`	`+ enumName := ""`
`574`	`575`	`if notNull {`
`575`	`576`	`if schema.Name == req.Catalog.DefaultSchema {`
`576`		`- return StructName(enum.Name, options)`
	`577`	`+ enumName = StructName(enum.Name, options)`
	`578`	`+ } else {`
	`579`	`+ enumName = StructName(schema.Name+"_"+enum.Name, options)`
`577`	`580`	`}`
`578`		`- return StructName(schema.Name+"_"+enum.Name, options)`
`579`	`581`	`} else {`
`580`	`582`	`if schema.Name == req.Catalog.DefaultSchema {`
`581`		`- return "Null" + StructName(enum.Name, options)`
	`583`	`+ enumName = "Null" + StructName(enum.Name, options)`
	`584`	`+ } else {`
	`585`	`+ enumName = "Null" + StructName(schema.Name+"_"+enum.Name, options)`
`582`	`586`	`}`
`583`		`- return "Null" + StructName(schema.Name+"_"+enum.Name, options)`
`584`	`587`	`}`
	`588`	`+ if options.ModelsPackageImportPath != "" {`
	`589`	`+ return options.OutputModelsPackage + "." + enumName`
	`590`	`+ }`
	`591`	`+ return enumName`
`585`	`592`	`}`
`586`	`593`	`}`
`587`	`594`
Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ func (v QueryValue) Type() string {`
`88`	`88`	`return v.Typ`
`89`	`89`	`}`
`90`	`90`	`if v.Struct != nil {`
`91`		`- return v.Struct.Name`
	`91`	`+ return v.Struct.Type()`
`92`	`92`	`}`
`93`	`93`	`panic("no type for QueryValue: " + v.Name)`
`94`	`94`	`}`
Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,7 @@ func buildStructs(req plugin.GenerateRequest, options opts.Options) []Struct {`
`83`	`83`	`s := Struct{`
`84`	`84`	`Table: &plugin.Identifier{Schema: schema.Name, Name: table.Rel.Name},`
`85`	`85`	`Name: StructName(structName, options),`
	`86`	`+ Package: options.OutputModelsPackage,`
`86`	`87`	`Comment: table.Comment,`
`87`	`88`	`}`
`88`	`89`	`for _, column := range table.Columns {`
`@@ -146,7 +147,7 @@ func newGoEmbed(embed *plugin.Identifier, structs []Struct, defaultSchema string`
`146`	`147`	`}`
`147`	`148`
`148`	`149`	`return &goEmbed{`
`149`		`- modelType: s.Name,`
	`150`	`+ modelType: s.Type(),`
`150`	`151`	`modelName: s.Name,`
`151`	`152`	`fields: fields,`
`152`	`153`	`}`
Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,7 @@ type {{.Name}} struct { {{- range .Fields}}`
`156`	`156`	`{{- if .Comment}}`
`157`	`157`	`{{comment .Comment}}{{else}}`
`158`	`158`	`{{- end}}`
`159`		`- {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}`
	`159`	`+ {{.Name}} {{trimPrefix .Type (printf "%s%s" $.Package ".") }} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}`
`160`	`160`	`{{- end}}`
`161`	`161`	`}`
`162`	`162`	`{{end}}`