sqlboiler/output.go

package sqlboiler

import (
	"bytes"
	"errors"
	"fmt"
	"go/format"
	"io"
	"os"
	"path/filepath"
	"text/template"
)

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 {
	if len(state.Templates) == 0 {
		return errors.New("No template files located for generation")
	}
	var out [][]byte
	var imps imports

	imps.standard = defaultTemplateImports.standard
	imps.thirdParty = defaultTemplateImports.thirdParty

	for _, template := range state.Templates {
		imps = combineTypeImports(imps, importsBasedOnType, data.Table.Columns)
		resp, err := generateTemplate(template, data)
		if err != nil {
			return fmt.Errorf("Error generating template %s: %s", template.Name(), err)
		}
		out = append(out, resp)
	}

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

	return nil
}

// generateTestOutput builds the test file output and sends it to outHandler for saving
func generateTestOutput(state *State, data *templateData) error {
	if len(state.TestTemplates) == 0 {
		return errors.New("No template files located for generation")
	}

	var out [][]byte
	var imps imports

	imps.standard = defaultTestTemplateImports.standard
	imps.thirdParty = defaultTestTemplateImports.thirdParty

	for _, template := range state.TestTemplates {
		resp, err := generateTemplate(template, data)
		if err != nil {
			return fmt.Errorf("Error generating test template %s: %s", template.Name(), err)
		}
		out = append(out, resp)
	}

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

	return nil
}

// generateSingletonOutput processes the templates that should only be run
// one time.
func generateSingletonOutput(state *State) error {
	if state.SingletonTemplates == nil {
		return errors.New("No singleton templates located for generation")
	}

	templateData := &templateData{
		PkgName:    state.Config.PkgName,
		DriverName: state.Config.DriverName,
	}

	for _, template := range state.SingletonTemplates {
		var imps imports

		resp, err := generateTemplate(template, templateData)
		if err != nil {
			return fmt.Errorf("Error generating template %s: %s", template.Name(), err)
		}

		fName := template.Name()
		ext := filepath.Ext(fName)
		fName = fName[0 : len(fName)-len(ext)]

		imps.standard = defaultSingletonTemplateImports[fName].standard
		imps.thirdParty = defaultSingletonTemplateImports[fName].thirdParty

		fName = fName + ".go"

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

	return nil
}

// generateSingletonTestOutput processes the templates that should only be run
// one time.
func generateSingletonTestOutput(state *State) error {
	if state.SingletonTestTemplates == nil {
		return errors.New("No singleton test templates located for generation")
	}

	templateData := &templateData{
		PkgName:    state.Config.PkgName,
		DriverName: state.Config.DriverName,
	}

	for _, template := range state.SingletonTestTemplates {
		var imps imports

		resp, err := generateTemplate(template, templateData)
		if err != nil {
			return fmt.Errorf("Error generating test template %s: %s", template.Name(), err)
		}

		fName := template.Name()
		ext := filepath.Ext(fName)
		fName = fName[0 : len(fName)-len(ext)]

		imps.standard = defaultSingletonTestTemplateImports[fName].standard
		imps.thirdParty = defaultSingletonTestTemplateImports[fName].thirdParty

		fName = fName + "_test.go"

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

	return nil
}

func generateTestMainOutput(state *State) 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

	templateData := &templateData{
		Tables:     state.Tables,
		PkgName:    state.Config.PkgName,
		DriverName: state.Config.DriverName,
	}

	resp, err := generateTemplate(state.TestMainTemplate, templateData)
	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 fmt.Errorf("Unable to create output file %s: %s", path, err)
	}
	defer outFile.Close()
	out = outFile

	if _, err := fmt.Fprintf(out, "package %s\n\n", pkgName); err != nil {
		return fmt.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 fmt.Errorf("Unable to write imports to file handle: %v", err)
		}
	}

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

	return nil
}

// generateTemplate takes a template and returns the output of the template execution.
func generateTemplate(t *template.Template, data *templateData) ([]byte, error) {
	var buf bytes.Buffer
	if err := t.Execute(&buf, data); err != nil {
		return nil, err
	}

	output, err := format.Source(buf.Bytes())
	if err != nil {
		return nil, err
	}

	return output, nil
}