microweb.go

package main

import (
	"text/template"

	"github.com/golang/protobuf/proto"
	"github.com/lyft/protoc-gen-star"
	"github.com/lyft/protoc-gen-star/lang/go"
	"google.golang.org/genproto/googleapis/api/annotations"
)

type handler struct {
	method  string
	pattern string
	body    string
}

type MicroWebModule struct {
	*pgs.ModuleBase
	ctx pgsgo.Context
	tpl *template.Template
}

func MicroWeb() *MicroWebModule {
	return &MicroWebModule{
		ModuleBase: &pgs.ModuleBase{},
	}
}

func (p *MicroWebModule) InitContext(ctx pgs.BuildContext) {
	p.ModuleBase.InitContext(ctx)
	p.ctx = pgsgo.InitContext(ctx.Parameters())

	tpl := template.New("microweb").Funcs(map[string]interface{}{
		"package":        p.ctx.PackageName,
		"name":           p.ctx.Name,
		"marshaler":      p.marshaler,
		"unmarshaler":    p.unmarshaler,
		"handlerPrefix":  p.handlerPrefix,
		"handlerMethod":  p.handlerMethod,
		"handlerPattern": p.handlerPattern,
		"handlerBody":    p.handlerBody,
	})

	p.tpl = template.Must(tpl.Parse(microwebTpl))
}

// Name satisfies the generator.Plugin interface.
func (p *MicroWebModule) Name() string {
	return "web"
}

func (p *MicroWebModule) Execute(targets map[string]pgs.File, pkgs map[string]pgs.Package) []pgs.Artifact {
	for _, t := range targets {
		p.generate(t)
	}

	return p.Artifacts()
}

func getHandler(m pgs.Method) *handler {
	opts := m.Descriptor().GetOptions()
	pext, _ := proto.GetExtension(opts, annotations.E_Http)
	ext, ok := pext.(*annotations.HttpRule)

	if !ok {
		return nil
	}

	switch {
	case ext.GetGet() != "":
		return &handler{method: "GET", pattern: ext.GetGet(), body: ext.GetBody()}
	case ext.GetPost() != "":
		return &handler{method: "POST", pattern: ext.GetPost(), body: ext.GetBody()}
	case ext.GetPut() != "":
		return &handler{method: "PUT", pattern: ext.GetPut(), body: ext.GetBody()}
	case ext.GetPatch() != "":
		return &handler{method: "PATCH", pattern: ext.GetPatch(), body: ext.GetBody()}
	case ext.GetDelete() != "":
		return &handler{method: "DELETE", pattern: ext.GetDelete(), body: ext.GetBody()}
	}

	return nil
}

func (p *MicroWebModule) generate(f pgs.File) {
	if len(f.Messages()) == 0 {
		return
	}

	name := p.ctx.OutputPath(f).SetExt(".web.go")
	p.AddGeneratorTemplateFile(name.String(), p.tpl, f)
}

func (p *MicroWebModule) handlerMethod(m pgs.Method) pgs.Name {
	return pgs.Name(getHandler(m).method)
}

func (p *MicroWebModule) handlerPattern(m pgs.Method) pgs.Name {
	return pgs.Name(getHandler(m).pattern)
}

func (p *MicroWebModule) handlerBody(m pgs.Method) pgs.Name {
	return pgs.Name(getHandler(m).body)
}

func (p *MicroWebModule) handlerPrefix(m pgs.Method) pgs.Name {
	return pgs.Name("")
}

func (p *MicroWebModule) marshaler(m pgs.Message) pgs.Name {
	return p.ctx.Name(m) + "JSONMarshaler"
}

func (p *MicroWebModule) unmarshaler(m pgs.Message) pgs.Name {
	return p.ctx.Name(m) + "JSONUnmarshaler"
}

const microwebTpl = `// Code generated by protoc-gen-microweb. DO NOT EDIT.
// source: {{ name .}}.proto

package {{ package . }}

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"

	"github.com/golang/protobuf/jsonpb"
	"github.com/go-chi/render"
	"github.com/go-chi/chi"
)

{{ range $idx, $svc := .Services }}
type web{{ $svc.Name }}Handler struct {
	r chi.Router
	h {{ $svc.Name }}Handler
}

func (h *web{{ $svc.Name }}Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.r.ServeHTTP(w, r)
}

{{ range .Methods }}
func (h *web{{ $svc.Name }}Handler) {{ name . }}(w http.ResponseWriter, r *http.Request) {
	req := &{{ .Input.Name }}{}
	{{ if ne .Output.Name "Empty" -}}
	resp := &{{ .Output.Name }}{}
	{{- end }}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusPreconditionFailed)
		return
	}

	if err := h.h.{{ name . }}(
		context.Background(),
		req,
		{{ if ne .Output.Name "Empty" -}}
		resp,
		{{- end }}
	); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	{{ if ne .Output.Name "Empty" }}
	{{- if eq (handlerMethod .) "POST" }}
	render.Status(r, http.StatusCreated)
	render.JSON(w, r, resp)
	{{- end }}
	{{- if eq (handlerMethod .) "DELETE" }}
	render.Status(r, http.StatusOK)
	render.JSON(w, r, resp)
	{{- end }}
	{{- if eq (handlerMethod .) "PATCH" }}
	render.Status(r, http.StatusOK)
	render.JSON(w, r, resp)
	{{- end }}
	{{- if eq (handlerMethod .) "PUT" }}
	render.Status(r, http.StatusOK)
	render.JSON(w, r, resp)
	{{- end }}
	{{- if eq (handlerMethod .) "GET" }}
	render.Status(r, http.StatusOK)
	render.JSON(w, r, resp)
	{{- end }}
	{{- else }}
	render.Status(r, http.StatusNoContent)
	render.NoContent(w, r)
	{{- end }}
}
{{ end }}

func Register{{ .Name }}Web(r chi.Router, i {{ .Name }}Handler, middlewares ...func(http.Handler) http.Handler) {
	handler := &web{{ .Name }}Handler{
		r: r,
		h: i,
	}

	{{ range .Methods }}{{ if .ServerStreaming }}{{- else -}}
	r.MethodFunc("{{ handlerMethod . }}", "{{ handlerPrefix . }}{{ handlerPattern . }}", handler.{{ name .}})
	{{ end }}{{- end -}}
}

{{ end }}

{{ range .AllMessages }}

// {{ marshaler . }} describes the default jsonpb.Marshaler used by all
// instances of {{ name . }}. This struct is safe to replace or modify but
// should not be done so concurrently.
var {{ marshaler . }} = new(jsonpb.Marshaler)

// MarshalJSON satisfies the encoding/json Marshaler interface. This method
// uses the more correct jsonpb package to correctly marshal the message.
func (m *{{ name . }}) MarshalJSON() ([]byte, error) {
	if m == nil {
		return json.Marshal(nil)
	}

	buf := &bytes.Buffer{}

	if err := {{ marshaler . }}.Marshal(buf, m); err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

var _ json.Marshaler = (*{{ name . }})(nil)

// {{ unmarshaler . }} describes the default jsonpb.Unmarshaler used by all
// instances of {{ name . }}. This struct is safe to replace or modify but
// should not be done so concurrently.
var {{ unmarshaler . }} = new(jsonpb.Unmarshaler)

// UnmarshalJSON satisfies the encoding/json Unmarshaler interface. This method
// uses the more correct jsonpb package to correctly unmarshal the message.
func (m *{{ name . }}) UnmarshalJSON(b []byte) error {
	return {{ unmarshaler . }}.Unmarshal(bytes.NewReader(b), m)
}

var _ json.Unmarshaler = (*{{ name . }})(nil)

{{ end }}
`