generator.go

package generator

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"

	"github.com/tagenasec/go-queryset/internal/parser"
	"github.com/pkg/errors"
	"golang.org/x/tools/imports"
)

type Generator struct {
	StructsParser *parser.Structs
}

// Generate generates output file with querysets
func (g Generator) Generate(ctx context.Context, inFilePath, outFilePath string) error {
	parsedFile, err := g.StructsParser.ParseFile(ctx, inFilePath)
	if err != nil {
		return errors.Wrapf(err, "can't parse file %s to get structs", inFilePath)
	}

	var r io.Reader
	r, err = GenerateQuerySetsForStructs(parsedFile.Types, parsedFile.Structs)
	if err != nil {
		return errors.Wrap(err, "can't generate query sets")
	}

	if r == nil {
		return fmt.Errorf("no structs to generate query set in %s", inFilePath)
	}

	if err = g.writeQuerySetsToOutput(r, parsedFile.PackageName, outFilePath); err != nil {
		return errors.Wrapf(err, "can't save query sets to out file %s", outFilePath)
	}

	var absOutPath string
	absOutPath, err = filepath.Abs(outFilePath)
	if err != nil {
		absOutPath = outFilePath
	}

	log.Printf("successfully wrote querysets to %s", absOutPath)
	return nil
}

func (g Generator) writeQuerySetsToOutput(r io.Reader, packageName, outFile string) error {
	const hdrTmpl = `%s
	package %s

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"gorm.io/gorm"
)
`

	// https://golang.org/s/generatedcode
	const genHdr = `// Code generated by go-queryset. DO NOT EDIT.`

	var buf bytes.Buffer
	pkgName := fmt.Sprintf(hdrTmpl, genHdr, packageName)
	if _, err := buf.WriteString(pkgName); err != nil {
		return errors.Wrap(err, "can't write hdr string into buf")
	}
	if _, err := io.Copy(&buf, r); err != nil {
		return errors.Wrap(err, "can't write to buf")
	}

	formattedRes, err := imports.Process(outFile, buf.Bytes(), nil)
	if err != nil {
		if os.Getenv("GOQUERYSET_DEBUG_IMPORTS_ERRORS") == "1" {
			log.Printf("Can't format generated file: %s", err)
			formattedRes = buf.Bytes()
		} else {
			return errors.Wrap(err, "can't format generated file")
		}
	}

	var outF *os.File
	outF, err = os.OpenFile(outFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0640) // nolint: gas
	if err != nil {
		return fmt.Errorf("can't open out file: %s", err)
	}
	defer func() {
		if e := outF.Close(); e != nil {
			log.Printf("can't close file: %s", e)
		}
	}()

	if _, err = outF.Write(formattedRes); err != nil {
		return errors.Wrap(err, "can't write to out file")
	}

	return nil
}