fix(generator): Refactor model file generation (#85)

Co-authored-by: Takafumi SEKIGUCHI <takkyuuplayer@gmail.com>

Author: keroxp

README.md

@@ -129,7 +129,7 @@ func GenerateModels() {
 And results are mostly like this:
 
 ```go
-// This file is generated by exql. DO NOT edit.
+// Code generated by exql. DO NOT EDIT.
 package model
 
 type Users struct {

generator.go

@@ -1,17 +1,13 @@
 package exql
 
 import (
-	"bytes"
 	"database/sql"
 	"fmt"
 	"go/format"
+	"log"
 	"os"
 	"path/filepath"
-	"strconv"
-	"strings"
-	"text/template"
-
-	"github.com/iancoleman/strcase"
+	"regexp"
 )
 
 type Generator interface {
@@ -24,19 +20,16 @@ type GenerateOptions struct {
 	OutDir  string
 	Package string
 	Exclude []string
+	// FileNameMap maps table names to output file names.
+	// Values must match [A-Za-z0-9_-]+.go.
+	FileNameMap map[string]string
 }
 
-type templateData struct {
-	Imports            string
-	Model              string
-	ModelLower         string
-	M                  string
-	Package            string
-	Fields             string
-	UpdaterFields      string
-	ScannedFields      string
-	TableName          string
-	TableNameGoLiteral string
+var safeModelFileNamePattern = regexp.MustCompile(`^[A-Za-z0-9_-]+\.go$`)
+
+type modelFileOutput struct {
+	path   string
+	source []byte
 }
 
 func NewGenerator(db *sql.DB) Generator {
@@ -80,95 +73,63 @@ func (d *generator) Generate(opts *GenerateOptions) error {
 	if err := rows.Err(); err != nil {
 		return err
 	}
+	var outputs []*modelFileOutput
+	seenPaths := map[string]string{}
 	for _, table := range tables {
-		if err := d.generateModelFile(table, opts); err != nil {
+		output, err := d.generateModelFile(table, opts)
+		if err != nil {
+			return err
+		}
+		if prevTable, ok := seenPaths[output.path]; ok {
+			return fmt.Errorf("duplicate generated model file %q for tables %q and %q", output.path, prevTable, table)
+		}
+		seenPaths[output.path] = table
+		outputs = append(outputs, output)
+	}
+	for _, output := range outputs {
+		if err := writeModelFile(output); err != nil {
 			return err
 		}
 	}
 	return nil
 }
 
-func (d *generator) generateModelFile(tableName string, opt *GenerateOptions) error {
-	tmpl := template.Must(template.New("model").Parse(modelTemplate))
+func (d *generator) generateModelFile(tableName string, opt *GenerateOptions) (*modelFileOutput, error) {
 	p := NewParser()
 	table, err := p.ParseTable(d.db, tableName)
 	if err != nil {
-		return err
-	}
-	var imports []string
-
-	if table.HasJsonField() {
-		imports = append(imports, `import "encoding/json"`)
-	}
-	if table.HasTimeField() {
-		imports = append(imports, `import "time"`)
+		return nil, err
 	}
-	if table.HasNullField() {
-		imports = append(imports, `import "github.com/loilo-inc/exql/v3/null"`)
+	modelFile, err := table.GenerateModelFile(opt.Package)
+	if err != nil {
+		return nil, err
 	}
-	fields := strings.Builder{}
-	updateFields := strings.Builder{}
-	scannedFields := strings.Builder{}
-	for i, col := range table.Columns {
-		scannedFields.WriteString(fmt.Sprintf(
-			"\t&%s.%s,", table.TableName[0:1], col.Field()),
-		)
-		fields.WriteString(fmt.Sprintf("\t%s", col.Field()))
-		updateFields.WriteString(fmt.Sprintf("\t%s", col.UpdateField()))
-		if i < len(table.Columns)-1 {
-			scannedFields.WriteString("\n")
-			fields.WriteString("\n")
-			updateFields.WriteString("\n")
+	outFileName := modelFile.Name
+	if mappedFileName, ok := opt.FileNameMap[tableName]; ok {
+		if err := validateMappedModelFileName(mappedFileName); err != nil {
+			return nil, err
 		}
+		outFileName = mappedFileName
 	}
-	data := &templateData{
-		Imports:            strings.Join(imports, "\n"),
-		Model:              strcase.ToCamel(table.TableName),
-		ModelLower:         strcase.ToLowerCamel(table.TableName),
-		M:                  table.TableName[0:1],
-		UpdaterFields:      updateFields.String(),
-		Package:            opt.Package,
-		Fields:             fields.String(),
-		TableName:          tableName,
-		TableNameGoLiteral: strconv.Quote(tableName),
-		ScannedFields:      scannedFields.String(),
-	}
-	outFile := filepath.Join(
-		opt.OutDir,
-		fmt.Sprintf("%s.go", strcase.ToSnake(table.TableName)),
-	)
-	var buf = &bytes.Buffer{}
-	if err := tmpl.Execute(buf, data); err != nil {
-		return err
-	}
-	if fmted, err := format.Source(buf.Bytes()); err != nil {
+	return &modelFileOutput{
+		path:   filepath.Join(opt.OutDir, outFileName),
+		source: modelFile.Source,
+	}, nil
+}
+
+func writeModelFile(output *modelFileOutput) error {
+	if fmted, err := format.Source(output.source); err != nil {
 		return err
-	} else if err := os.WriteFile(outFile, fmted, 0640); err != nil {
+	} else if err := os.WriteFile(output.path, fmted, 0640); err != nil {
 		return err
 	}
+	log.Printf("generated file: %s", output.path)
 	return nil
 }
 
-const modelTemplate = `// This file is generated by exql. DO NOT edit.
-package {{.Package}}
-
-{{.Imports}}
-
-type {{.Model}} struct {
-{{.Fields}}
-}
-
-func ({{.M}} *{{.Model}}) TableName() string {
-	return {{.Model}}TableName
-}
-
-type Update{{.Model}} struct {
-{{.UpdaterFields}}
-}
-
-func ({{.M}} *Update{{.Model}}) UpdateTableName() string {
-	return {{.Model}}TableName
+func validateMappedModelFileName(name string) error {
+	if !safeModelFileNamePattern.MatchString(name) {
+		return fmt.Errorf("invalid model file name %q: must match [A-Za-z0-9_-]+.go", name)
+	}
+	return nil
 }
-
-const {{.Model}}TableName = {{.TableNameGoLiteral}}
-`

generator_test.go

@@ -1,7 +1,9 @@
 package exql
 
 import (
+	"bytes"
 	"fmt"
+	"log"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -136,3 +138,142 @@ func TestGenerator_Generate_EscapesTableNameGoLiteral(t *testing.T) {
 	}
 	assert.NoError(t, mock.ExpectationsWereMet())
 }
+
+func TestGenerator_Generate_UsesFileNameMap(t *testing.T) {
+	mockDb, mock, err := sqlmock.New()
+	assert.NoError(t, err)
+	defer mockDb.Close()
+
+	table := "evil/foo"
+	mock.ExpectQuery(`show tables`).WillReturnRows(sqlmock.NewRows([]string{"tables"}).AddRow(table))
+	mock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf("show columns from `%s`", table))).WillReturnRows(
+		sqlmock.NewRows([]string{"Field", "Type", "Null", "Key", "Default", "Extra"}).
+			AddRow("id", "int(11)", "NO", "PRI", nil, ""),
+	)
+
+	dir := t.TempDir()
+	var logBuf bytes.Buffer
+	oldLogOutput := log.Writer()
+	oldLogFlags := log.Flags()
+	log.SetOutput(&logBuf)
+	log.SetFlags(0)
+	t.Cleanup(func() {
+		log.SetOutput(oldLogOutput)
+		log.SetFlags(oldLogFlags)
+	})
+
+	err = NewGenerator(mockDb).Generate(&GenerateOptions{
+		OutDir:  dir,
+		Package: "dist",
+		FileNameMap: map[string]string{
+			table: "evil.go",
+		},
+	})
+	assert.NoError(t, err)
+
+	outFile := filepath.Join(dir, "evil.go")
+	content, err := os.ReadFile(outFile)
+	assert.NoError(t, err)
+	assert.Contains(t, string(content), "package dist")
+	assert.Contains(t, string(content), `"evil/foo"`)
+	assert.Contains(t, logBuf.String(), outFile)
+	assert.NoError(t, mock.ExpectationsWereMet())
+}
+
+func TestGenerator_Generate_ReturnsErrorForDuplicateFileNames(t *testing.T) {
+	mockDb, mock, err := sqlmock.New()
+	assert.NoError(t, err)
+	defer mockDb.Close()
+
+	mock.ExpectQuery(`show tables`).WillReturnRows(
+		sqlmock.NewRows([]string{"tables"}).
+			AddRow("users").
+			AddRow("user_groups"),
+	)
+	mock.ExpectQuery("show columns from `users`").WillReturnRows(
+		sqlmock.NewRows([]string{"Field", "Type", "Null", "Key", "Default", "Extra"}).
+			AddRow("id", "int(11)", "NO", "PRI", nil, ""),
+	)
+	mock.ExpectQuery("show columns from `user_groups`").WillReturnRows(
+		sqlmock.NewRows([]string{"Field", "Type", "Null", "Key", "Default", "Extra"}).
+			AddRow("id", "int(11)", "NO", "PRI", nil, ""),
+	)
+
+	dir := t.TempDir()
+	err = NewGenerator(mockDb).Generate(&GenerateOptions{
+		OutDir:  dir,
+		Package: "dist",
+		FileNameMap: map[string]string{
+			"users":       "models.go",
+			"user_groups": "models.go",
+		},
+	})
+	assert.ErrorContains(t, err, "duplicate generated model file")
+	assert.ErrorContains(t, err, "users")
+	assert.ErrorContains(t, err, "user_groups")
+	entries, readErr := os.ReadDir(dir)
+	assert.NoError(t, readErr)
+	assert.Empty(t, entries)
+	assert.NoError(t, mock.ExpectationsWereMet())
+}
+
+func TestGenerator_Generate_PropagatesWriteError(t *testing.T) {
+	mockDb, mock, err := sqlmock.New()
+	assert.NoError(t, err)
+	defer mockDb.Close()
+
+	table := "users"
+	mock.ExpectQuery(`show tables`).WillReturnRows(sqlmock.NewRows([]string{"tables"}).AddRow(table))
+	mock.ExpectQuery("show columns from `users`").WillReturnRows(
+		sqlmock.NewRows([]string{"Field", "Type", "Null", "Key", "Default", "Extra"}).
+			AddRow("id", "int(11)", "NO", "PRI", nil, ""),
+	)
+
+	dir := t.TempDir()
+	err = os.Mkdir(filepath.Join(dir, "users.go"), 0750)
+	assert.NoError(t, err)
+
+	err = NewGenerator(mockDb).Generate(&GenerateOptions{
+		OutDir:  dir,
+		Package: "dist",
+		FileNameMap: map[string]string{
+			table: "users.go",
+		},
+	})
+	assert.Error(t, err)
+	assert.NoError(t, mock.ExpectationsWereMet())
+}
+
+func TestValidateMappedModelFileName(t *testing.T) {
+	for _, name := range []string{
+		"users.go",
+		"user_groups.go",
+		"user-groups_1.go",
+		"Users1.go",
+	} {
+		t.Run("valid "+name, func(t *testing.T) {
+			assert.NoError(t, validateMappedModelFileName(name))
+		})
+	}
+
+	for _, name := range []string{
+		"",
+		".",
+		"..",
+		".go",
+		"users",
+		"users.go.txt",
+		"user.name.go",
+		"user name.go",
+		"user;name.go",
+		"ユーザー.go",
+		"../users.go",
+		"/tmp/users.go",
+		"nested/users.go",
+		`nested\users.go`,
+	} {
+		t.Run("invalid "+name, func(t *testing.T) {
+			assert.Error(t, validateMappedModelFileName(name))
+		})
+	}
+}