sqlboiler/cmds/output.go

package cmds

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

	imps.standard = sqlBoilerImports.standard
	imps.thirdparty = sqlBoilerImports.thirdparty

	for _, template := range cmdData.Templates {
		imps = combineTypeImports(imps, sqlBoilerTypeImports, 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(cmdData.OutFolder, fName, cmdData.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(cmdData *CmdData, data *tplData) error {
	if len(cmdData.TestTemplates) == 0 {
		return errors.New("No template files located for generation")
	}

	var out [][]byte
	var imps imports

	imps.standard = sqlBoilerTestImports.standard
	imps.thirdparty = sqlBoilerTestImports.thirdparty

	for _, template := range cmdData.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(cmdData.OutFolder, fName, cmdData.PkgName, imps, out)
	if err != nil {
		return err
	}

	return nil
}

func generateSinglesOutput(cmdData *CmdData) error {
	if cmdData.SingleTemplates == nil {
		return errors.New("No single templates located for generation")
	}

	tplData := &tplData{
		PkgName:    cmdData.PkgName,
		DriverName: cmdData.DriverName,
	}

	for _, template := range cmdData.SingleTemplates {
		var imps imports

		resp, err := generateTemplate(template, tplData)
		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 = sqlBoilerSinglesImports[fName].standard
		imps.thirdparty = sqlBoilerSinglesImports[fName].thirdparty

		fName = fName + ".go"

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

	return nil
}

func generateSinglesTestOutput(cmdData *CmdData) error {
	if cmdData.SingleTestTemplates == nil {
		return errors.New("No single test templates located for generation")
	}

	tplData := &tplData{
		PkgName:    cmdData.PkgName,
		DriverName: cmdData.DriverName,
	}

	for _, template := range cmdData.SingleTestTemplates {
		var imps imports

		resp, err := generateTemplate(template, tplData)
		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 = sqlBoilerSinglesTestImports[fName].standard
		imps.thirdparty = sqlBoilerSinglesTestImports[fName].thirdparty

		fName = fName + "_test.go"

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

	return nil
}

func generateTestMainOutput(cmdData *CmdData) error {
	if cmdData.TestMainTemplate == nil {
		return errors.New("No TestMain template located for generation")
	}

	var out [][]byte
	var imps imports

	imps.standard = sqlBoilerTestMainImports[cmdData.DriverName].standard
	imps.thirdparty = sqlBoilerTestMainImports[cmdData.DriverName].thirdparty

	var tables []string
	for _, v := range cmdData.Tables {
		tables = append(tables, v.Name)
	}

	tplData := &tplData{
		PkgName:    cmdData.PkgName,
		DriverName: cmdData.DriverName,
		Tables:     tables,
	}

	resp, err := generateTemplate(cmdData.TestMainTemplate, tplData)
	if err != nil {
		return err
	}
	out = append(out, resp)

	err = outHandler(cmdData.OutFolder, "main_test.go", cmdData.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 *tplData) ([]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
}