github.com/brycereitano/goa@v0.0.0-20170315073847-8ffa6c85e265/goagen/gen_main/generator.go

package genmain

import (
	"flag"
	"fmt"
	"net"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/goadesign/goa/design"
	"github.com/goadesign/goa/goagen/codegen"
	"github.com/goadesign/goa/goagen/utils"
)

//NewGenerator returns an initialized instance of a JavaScript Client Generator
func NewGenerator(options ...Option) *Generator {
	g := &Generator{}

	for _, option := range options {
		option(g)
	}

	return g
}

// Generator is the application code generator.
type Generator struct {
	API       *design.APIDefinition // The API definition
	OutDir    string                // Path to output directory
	DesignPkg string                // Path to design package, only used to mark generated files.
	Target    string                // Name of generated "app" package
	Force     bool                  // Whether to override existing files
	genfiles  []string              // Generated files
}

// Generate is the generator entry point called by the meta generator.
func Generate() (files []string, err error) {
	var (
		outDir, designPkg, target, ver string
		force                          bool
	)

	set := flag.NewFlagSet("main", flag.PanicOnError)
	set.StringVar(&outDir, "out", "", "")
	set.StringVar(&designPkg, "design", "", "")
	set.StringVar(&target, "pkg", "app", "")
	set.StringVar(&ver, "version", "", "")
	set.BoolVar(&force, "force", false, "")
	set.Bool("notest", false, "")
	set.Parse(os.Args[1:])

	if err := codegen.CheckVersion(ver); err != nil {
		return nil, err
	}

	target = codegen.Goify(target, false)
	g := &Generator{OutDir: outDir, DesignPkg: designPkg, Target: target, Force: force, API: design.Design}

	return g.Generate()
}

// GenerateController generates the controller corresponding to the given
// resource and returns the generated filename.
func GenerateController(force bool, appPkg, outDir, pkg, name string, r *design.ResourceDefinition) (string, error) {
	filename := filepath.Join(outDir, codegen.SnakeCase(name)+".go")
	if force {
		os.Remove(filename)
	}
	if _, e := os.Stat(filename); e == nil {
		return "", nil
	}
	if err := os.MkdirAll(outDir, 0755); err != nil {
		return "", err
	}
	file, err := codegen.SourceFileFor(filename)
	if err != nil {
		return "", err
	}

	elems := strings.Split(appPkg, "/")
	pkgName := elems[len(elems)-1]
	var imp string
	if _, err := codegen.PackageSourcePath(appPkg); err == nil {
		imp = appPkg
	} else {
		imp, err = codegen.PackagePath(outDir)
		if err != nil {
			return "", err
		}
		imp = path.Join(filepath.ToSlash(imp), appPkg)
	}

	imports := []*codegen.ImportSpec{
		codegen.SimpleImport("io"),
		codegen.SimpleImport("github.com/goadesign/goa"),
		codegen.SimpleImport(imp),
		codegen.SimpleImport("golang.org/x/net/websocket"),
	}

	file.WriteHeader("", pkg, imports)
	if err = file.ExecuteTemplate("controller", ctrlT, funcMap(pkgName), r); err != nil {
		return "", err
	}
	err = r.IterateActions(func(a *design.ActionDefinition) error {
		if a.WebSocket() {
			return file.ExecuteTemplate("actionWS", actionWST, funcMap(pkgName), a)
		}
		return file.ExecuteTemplate("action", actionT, funcMap(pkgName), a)
	})
	if err != nil {
		return "", err
	}
	if err = file.FormatCode(); err != nil {
		return "", err
	}

	return filename, nil
}

// Generate produces the skeleton main.
func (g *Generator) Generate() (_ []string, err error) {
	if g.API == nil {
		return nil, fmt.Errorf("missing API definition, make sure design is properly initialized")
	}

	go utils.Catch(nil, func() { g.Cleanup() })

	defer func() {
		if err != nil {
			g.Cleanup()
		}
	}()

	if g.Target == "" {
		g.Target = "app"
	}

	codegen.Reserved[g.Target] = true

	mainFile := filepath.Join(g.OutDir, "main.go")
	if g.Force {
		os.Remove(mainFile)
	}
	_, err = os.Stat(mainFile)
	if err != nil {
		// ensure that the output directory exists before creating a new main
		if err := os.MkdirAll(g.OutDir, 0755); err != nil {
			return nil, err
		}
		if err = g.createMainFile(mainFile, funcMap(g.Target)); err != nil {
			return nil, err
		}
	}

	err = g.API.IterateResources(func(r *design.ResourceDefinition) error {
		filename, err := GenerateController(g.Force, g.Target, g.OutDir, "main", r.Name, r)
		if err != nil {
			return err
		}

		g.genfiles = append(g.genfiles, filename)
		return nil
	})
	if err != nil {
		return
	}

	return g.genfiles, nil
}

// Cleanup removes all the files generated by this generator during the last invokation of Generate.
func (g *Generator) Cleanup() {
	for _, f := range g.genfiles {
		os.Remove(f)
	}
	g.genfiles = nil
}

func (g *Generator) createMainFile(mainFile string, funcs template.FuncMap) error {
	g.genfiles = append(g.genfiles, mainFile)
	file, err := codegen.SourceFileFor(mainFile)
	if err != nil {
		return err
	}
	funcs["getPort"] = func(hostport string) string {
		_, port, err := net.SplitHostPort(hostport)
		if err != nil {
			return "8080"
		}
		return port
	}
	outPkg, err := codegen.PackagePath(g.OutDir)
	if err != nil {
		return err
	}
	appPkg := path.Join(outPkg, "app")
	imports := []*codegen.ImportSpec{
		codegen.SimpleImport("time"),
		codegen.SimpleImport("github.com/goadesign/goa"),
		codegen.SimpleImport("github.com/goadesign/goa/middleware"),
		codegen.SimpleImport(appPkg),
	}
	file.Write([]byte("//go:generate goagen bootstrap -d " + g.DesignPkg + "\n\n"))
	file.WriteHeader("", "main", imports)
	data := map[string]interface{}{
		"Name": g.API.Name,
		"API":  g.API,
	}
	if err = file.ExecuteTemplate("main", mainT, funcs, data); err != nil {
		return err
	}
	return file.FormatCode()
}

// tempCount is the counter used to create unique temporary variable names.
var tempCount int

// tempvar generates a unique temp var name.
func tempvar() string {
	tempCount++
	if tempCount == 1 {
		return "c"
	}
	return fmt.Sprintf("c%d", tempCount)
}

func okResp(a *design.ActionDefinition, appPkg string) map[string]interface{} {
	var ok *design.ResponseDefinition
	for _, resp := range a.Responses {
		if resp.Status == 200 {
			ok = resp
			break
		}
	}
	if ok == nil {
		return nil
	}
	var mt *design.MediaTypeDefinition
	var ok2 bool
	if mt, ok2 = design.Design.MediaTypes[design.CanonicalIdentifier(ok.MediaType)]; !ok2 {
		return nil
	}
	view := ok.ViewName
	if view == "" {
		view = design.DefaultView
	}
	pmt, _, err := mt.Project(view)
	if err != nil {
		return nil
	}
	var typeref string
	if pmt.IsError() {
		typeref = `goa.ErrInternal("not implemented")`
	} else {
		name := codegen.GoTypeRef(pmt, pmt.AllRequired(), 1, false)
		var pointer string
		if strings.HasPrefix(name, "*") {
			name = name[1:]
			pointer = "*"
		}
		typeref = fmt.Sprintf("%s%s.%s", pointer, appPkg, name)
		if strings.HasPrefix(typeref, "*") {
			typeref = "&" + typeref[1:]
		}
		typeref += "{}"
	}
	var nameSuffix string
	if view != "default" {
		nameSuffix = codegen.Goify(view, true)
	}
	return map[string]interface{}{
		"Name":    ok.Name + nameSuffix,
		"GoType":  codegen.GoNativeType(pmt),
		"TypeRef": typeref,
	}
}

// funcMap creates the funcMap used to render the controller code.
func funcMap(appPkg string) template.FuncMap {
	return template.FuncMap{
		"tempvar":   tempvar,
		"okResp":    okResp,
		"targetPkg": func() string { return appPkg },
	}
}

const ctrlT = `// {{ $ctrlName := printf "%s%s" (goify .Name true) "Controller" }}{{ $ctrlName }} implements the {{ .Name }} resource.
type {{ $ctrlName }} struct {
	*goa.Controller
}

// New{{ $ctrlName }} creates a {{ .Name }} controller.
func New{{ $ctrlName }}(service *goa.Service) *{{ $ctrlName }} {
	return &{{ $ctrlName }}{Controller: service.NewController("{{ $ctrlName }}")}
}
`

const actionT = `{{ $ctrlName := printf "%s%s" (goify .Parent.Name true) "Controller" }}// {{ goify .Name true }} runs the {{ .Name }} action.
func (c *{{ $ctrlName }}) {{ goify .Name true }}(ctx *{{ targetPkg }}.{{ goify .Name true }}{{ goify .Parent.Name true }}Context) error {
	// {{ $ctrlName }}_{{ goify .Name true }}: start_implement

	// Put your logic here

	// {{ $ctrlName }}_{{ goify .Name true }}: end_implement
{{ $ok := okResp . targetPkg }}{{ if $ok }} res := {{ $ok.TypeRef }}
{{ end }} return {{ if $ok }}ctx.{{ $ok.Name }}(res){{ else }}nil{{ end }}
}
`

const actionWST = `{{ $ctrlName := printf "%s%s" (goify .Parent.Name true) "Controller" }}// {{ goify .Name true }} runs the {{ .Name }} action.
func (c *{{ $ctrlName }}) {{ goify .Name true }}(ctx *{{ targetPkg }}.{{ goify .Name true }}{{ goify .Parent.Name true }}Context) error {
	c.{{ goify .Name true }}WSHandler(ctx).ServeHTTP(ctx.ResponseWriter, ctx.Request)
	return nil
}

// {{ goify .Name true }}WSHandler establishes a websocket connection to run the {{ .Name }} action.
func (c *{{ $ctrlName }}) {{ goify .Name true }}WSHandler(ctx *{{ targetPkg }}.{{ goify .Name true }}{{ goify .Parent.Name true }}Context) websocket.Handler {
	return func(ws *websocket.Conn) {
		// {{ $ctrlName }}_{{ goify .Name true }}: start_implement

		// Put your logic here

		// {{ $ctrlName }}_{{ goify .Name true }}: end_implement
		ws.Write([]byte("{{ .Name }} {{ .Parent.Name }}"))
		// Dummy echo websocket server
		io.Copy(ws, ws)
	}
}`

const mainT = `
func main() {
	// Create service
	service := goa.New({{ printf "%q" .Name }})

	// Mount middleware
	service.Use(middleware.RequestID())
	service.Use(middleware.LogRequest(true))
	service.Use(middleware.ErrorHandler(service, true))
	service.Use(middleware.Recover())
{{ $api := .API }}
{{ range $name, $res := $api.Resources }}{{ $name := goify $res.Name true }} // Mount "{{$res.Name}}" controller
	{{ $tmp := tempvar }}{{ $tmp }} := New{{ $name }}Controller(service)
	{{ targetPkg }}.Mount{{ $name }}Controller(service, {{ $tmp }})
{{ end }}

	// Start service
	if err := service.ListenAndServe(":{{ getPort .API.Host }}"); err != nil {
		service.LogError("startup", "err", err)
	}
}
`