schema_columns_gen.go

package gen

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

	"github.com/kataras/pg"
)

// GenerateColumnsFromSchema generates Go code for the given schema.
// The generated code includes a struct for each table which contains
// struct fields for each column information.
//
// Example Code:
//
// schema := pg.NewSchema()
// schema.MustRegister("companies", Company{})
// schema.MustRegister("customers", Customer{})
//
//	opts := &ExportOptions{
//		RootDir: "./definition",
//	}
//
//	if err := GenerateColumnsFromSchema(schema, opts); err != nil {
//		t.Fatal(err)
//	}
//
// Usage:
// definition.Company.Name.String() // returns "name"
// definition.Customer.Email.String() // returns "email"
//
// Useful for type-safe query builders.
func GenerateColumnsFromSchema(s *pg.Schema, e ExportOptions) error {
	if err := e.apply(); err != nil {
		return err
	}

	tables := s.Tables()
	if len(tables) == 0 {
		return nil
	}

	// Create root file to store common structures and functions.
	data, err := generateRoot(e.GetPackageName(""))
	if err != nil {
		return fmt.Errorf("generate root: %s: %w", e.GetPackageName("columns"), err)
	}

	filename := e.GetFileName(e.RootDir, "")

	err = mkdir(filename)
	if err != nil {
		return fmt.Errorf("mkdir: %s: %w", e.RootDir, err)
	}

	err = os.WriteFile(filename, data, e.FileMode)
	if err != nil {
		return fmt.Errorf("write root: %s: %w", filename, err)
	}

	// Create each file for each table definition.
	for _, td := range tables {
		data, err = generateTableDefininion(e.GetPackageName(td.Name), td)
		if err != nil {
			return fmt.Errorf("generate table: %s: %w", td.Name, err)
		}

		filename := e.GetFileName(e.RootDir, td.Name)
		if filename == "" {
			continue
		}

		mkdir(filename)

		err = os.WriteFile(filename, data, e.FileMode)
		if err != nil {
			return fmt.Errorf("write table: %s: defininion file: %s: %w", td.Name, filename, err)
		}
	}

	return nil
}

var generateRootTmpl = template.Must(
	template.New("").Parse(`
// Code generated by pg. DO NOT EDIT.
package {{.PackageName}}

// Column is a struct that represents a column in a table.
type Column struct {
	Name string
}

// String returns the name of the column.
func (c Column) String() string {
	return c.Name
}
`))

func generateRoot(packageName string) ([]byte, error) {
	tmplData := generateTemplateData{
		PackageName: packageName,
	}
	var buf bytes.Buffer
	if err := generateRootTmpl.Execute(&buf, tmplData); err != nil {
		return nil, err
	}

	return format.Source(buf.Bytes())
}

type generateTemplateData struct {
	*pg.Table
	PackageName string
}

var generateTableDefininionTmpl = template.Must(
	template.New(""). // Note that we don't need to put import paths here as the only one type is the generated Column in the same package.
				Parse(`
// Code generated by pg. DO NOT EDIT.
package {{.PackageName}}

// {{.StructName}} is a struct value that represents a record in the {{.Name}} table.
var {{.StructName}} = struct {
    PG_TableName string
    {{range .Columns}} {{.FieldName}} Column
    {{end}}
}{
    PG_TableName: "{{.Name}}",
    {{range .Columns}} {{.FieldName}}: Column{
        Name: "{{.Name}}",
    },
    {{end }}
}`))

func generateTableDefininion(packageName string, td *pg.Table) ([]byte, error) {
	tmplData := generateTemplateData{
		Table:       td,
		PackageName: packageName,
	}
	var buf bytes.Buffer
	if err := generateTableDefininionTmpl.Execute(&buf, tmplData); err != nil {
		return nil, err
	}

	return format.Source(buf.Bytes())
}

func mkdir(path string) error {
	dir := filepath.Dir(path)
	return os.MkdirAll(dir, 0777)
}

func getCallerPackageName() string {
	pc, _, _, _ := runtime.Caller(2)
	funcName := runtime.FuncForPC(pc).Name()
	lastSlash := strings.LastIndexByte(funcName, '/')
	if lastSlash < 0 {
		lastSlash = 0
	}

	lastDot := strings.LastIndexByte(funcName[lastSlash:], '.')
	if lastDot == -1 {
		return ""
	}

	return funcName[:lastDot]
}