generator.go

package exql

import (
	"bytes"
	"database/sql"
	"fmt"
	"go/format"
	"os"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/iancoleman/strcase"
)

type Generator interface {
	Generate(opts *GenerateOptions) error
}
type generator struct {
	db *sql.DB
}
type GenerateOptions struct {
	OutDir  string
	Package string
	Exclude []string
}

type templateData struct {
	Imports       string
	Model         string
	ModelLower    string
	M             string
	Package       string
	Fields        string
	UpdaterFields string
	ScannedFields string
	TableName     string
}

func NewGenerator(db *sql.DB) Generator {
	return &generator{db: db}
}

func (d *generator) Generate(opts *GenerateOptions) error {
	rows, err := d.db.Query(`show tables`)
	if err != nil {
		return err
	}
	if opts.OutDir == "" {
		opts.OutDir = "model"
	}
	if opts.Package == "" {
		opts.Package = "model"
	}
	if _, err := os.Stat(opts.OutDir); os.IsNotExist(err) {
		err := os.Mkdir(opts.OutDir, 0777)
		if err != nil {
			return err
		}
	} else if err != nil {
		return err
	}
	defer rows.Close()
	var tables []string
	for rows.Next() {
		var table string
		if err := rows.Scan(&table); err != nil {
			return err
		}
		for _, e := range opts.Exclude {
			if e == table {
				goto EOL
			}
		}
		tables = append(tables, table)
	EOL:
	}
	if err := rows.Err(); err != nil {
		return err
	}
	for _, table := range tables {
		if err := d.generateModelFile(table, opts); err != nil {
			return err
		}
	}
	return nil
}

func (d *generator) generateModelFile(tableName string, opt *GenerateOptions) error {
	tmpl := template.Must(template.New("model").Parse(modelTemplate))
	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"`)
	}
	if table.HasNullField() {
		imports = append(imports, `import "github.com/volatiletech/null"`)
	}
	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")
		}
	}
	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,
		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 err
	} else if err := os.WriteFile(outFile, fmted, os.ModePerm); err != nil {
		return err
	}
	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
}

const {{.Model}}TableName = "{{.TableName}}"
`