sqlboiler/output.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"go/format"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"text/template"

	"github.com/pkg/errors"
)

var testHarnessStdout io.Writer = os.Stdout
var testHarnessFileOpen = func(filename string) (io.WriteCloser, error) {
	file, err := os.Create(filename)
	return file, err
}

// generateOutput builds the file output and sends it to outHandler for saving
func generateOutput(state *State, data *templateData) error {
	return executeTemplates(executeTemplateData{
		state:                state,
		data:                 data,
		templates:            state.Templates,
		importSet:            defaultTemplateImports,
		combineImportsOnType: true,
		fileSuffix:           ".go",
	})
}

// generateTestOutput builds the test file output and sends it to outHandler for saving
func generateTestOutput(state *State, data *templateData) error {
	return executeTemplates(executeTemplateData{
		state:                state,
		data:                 data,
		templates:            state.TestTemplates,
		importSet:            defaultTestTemplateImports,
		combineImportsOnType: false,
		fileSuffix:           "_test.go",
	})
}

// generateSingletonOutput processes the templates that should only be run
// one time.
func generateSingletonOutput(state *State, data *templateData) error {
	return executeSingletonTemplates(executeTemplateData{
		state:          state,
		data:           data,
		templates:      state.SingletonTemplates,
		importNamedSet: defaultSingletonTemplateImports,
		fileSuffix:     ".go",
	})
}

// generateSingletonTestOutput processes the templates that should only be run
// one time.
func generateSingletonTestOutput(state *State, data *templateData) error {
	return executeSingletonTemplates(executeTemplateData{
		state:          state,
		data:           data,
		templates:      state.SingletonTestTemplates,
		importNamedSet: defaultSingletonTestTemplateImports,
		fileSuffix:     "_test.go",
	})
}

type executeTemplateData struct {
	state *State
	data  *templateData

	templates *templateList

	importSet      imports
	importNamedSet map[string]imports

	combineImportsOnType bool

	fileSuffix string
}

func executeTemplates(e executeTemplateData) error {
	if e.data.Table.IsJoinTable {
		return nil
	}

	var out [][]byte
	var imps imports

	imps.standard = e.importSet.standard
	imps.thirdParty = e.importSet.thirdParty

	for _, tplName := range e.templates.Templates() {
		if e.combineImportsOnType {
			imps = combineTypeImports(imps, importsBasedOnType, e.data.Table.Columns)
		}

		resp, err := executeTemplate(e.templates.Template, tplName, e.data)
		if err != nil {
			return errors.Wrapf(err, "Error generating template %s", tplName)
		}
		out = append(out, resp)
	}

	fName := e.data.Table.Name + e.fileSuffix
	err := outHandler(e.state.Config.OutFolder, fName, e.state.Config.PkgName, imps, out)
	if err != nil {
		return err
	}

	return nil
}

func executeSingletonTemplates(e executeTemplateData) error {
	if e.data.Table.IsJoinTable {
		return nil
	}

	rgxRemove := regexp.MustCompile(`[0-9]+_`)

	for _, tplName := range e.templates.Templates() {
		resp, err := executeTemplate(e.templates.Template, tplName, e.data)
		if err != nil {
			return errors.Wrapf(err, "Error generating template %s", tplName)
		}

		fName := tplName
		ext := filepath.Ext(fName)
		fName = rgxRemove.ReplaceAllString(fName[:len(fName)-len(ext)], "")

		imps := imports{
			standard:   e.importNamedSet[fName].standard,
			thirdParty: e.importNamedSet[fName].thirdParty,
		}

		err = outHandler(
			e.state.Config.OutFolder,
			fName+e.fileSuffix,
			e.state.Config.PkgName,
			imps,
			[][]byte{resp},
		)
		if err != nil {
			return err
		}
	}

	return nil
}

func generateTestMainOutput(state *State, data *templateData) error {
	if state.TestMainTemplate == nil {
		return errors.New("No TestMain template located for generation")
	}

	var out [][]byte
	var imps imports

	imps.standard = defaultTestMainImports[state.Config.DriverName].standard
	imps.thirdParty = defaultTestMainImports[state.Config.DriverName].thirdParty

	resp, err := executeTemplate(state.TestMainTemplate, state.TestMainTemplate.Name(), data)
	if err != nil {
		return err
	}
	out = append(out, resp)

	err = outHandler(state.Config.OutFolder, "main_test.go", state.Config.PkgName, imps, out)
	if err != nil {
		return err
	}

	return nil
}

func outHandler(outFolder string, fileName string, pkgName string, imps imports, contents [][]byte) error {
	out := testHarnessStdout

	path := filepath.Join(outFolder, fileName)

	outFile, err := testHarnessFileOpen(path)
	if err != nil {
		return errors.Wrapf(err, "Unable to create output file %s", path)
	}
	defer outFile.Close()
	out = outFile

	if _, err := fmt.Fprintf(out, "package %s\n\n", pkgName); err != nil {
		return errors.Errorf("Unable to write package name %s to file: %s", pkgName, path)
	}

	impStr := buildImportString(imps)
	if len(impStr) > 0 {
		if _, err := fmt.Fprintf(out, "%s\n", impStr); err != nil {
			return errors.Wrap(err, "Unable to write imports to file handle")
		}
	}

	for _, templateOutput := range contents {
		if _, err := fmt.Fprintf(out, "%s\n", templateOutput); err != nil {
			return errors.Wrap(err, "Unable to write template output to file handle")
		}
	}

	return nil
}

var rgxSyntaxError = regexp.MustCompile(`(\d+):\d+: `)

// executeTemplate takes a template and returns the output of the template
// execution.
func executeTemplate(t *template.Template, name string, data *templateData) ([]byte, error) {
	var buf bytes.Buffer
	if err := t.ExecuteTemplate(&buf, name, data); err != nil {
		return nil, errors.Wrap(err, "failed to execute template")
	}

	output, err := format.Source(buf.Bytes())
	if err != nil {
		matches := rgxSyntaxError.FindStringSubmatch(err.Error())
		if matches == nil {
			return nil, errors.Wrap(err, "failed to format template")
		}

		lineNum, _ := strconv.Atoi(matches[1])
		scanner := bufio.NewScanner(&buf)
		errBuf := &bytes.Buffer{}
		line := 0
		for ; scanner.Scan(); line++ {
			if delta := line - lineNum; delta < -5 || delta > 5 {
				continue
			}

			if line == lineNum {
				errBuf.WriteString(">>> ")
			} else {
				fmt.Fprintf(errBuf, "% 3d ", line)
			}
			errBuf.Write(scanner.Bytes())
			errBuf.WriteByte('\n')
		}

		return nil, errors.Wrapf(err, "failed to format template\n\n%s\n", errBuf.Bytes())
	}

	return output, nil
}