sqlboiler/cmds/output.go

232 lines
5.6 KiB
Go
Raw Normal View History

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
}