kit/cmd/gin-kit/main.go

package main

import (
	protogen2 "git.diulo.com/mogfee/kit/protogen"
	"google.golang.org/protobuf/compiler/protogen"
	"strings"
)

func main() {
	u := &Kit{
		imports: map[string]string{},
	}
	protogen.Options{}.Run(u.Generate)
}

type Kit struct {
	imports map[string]string
}

func (u *Kit) addImports(imp string) {
	u.imports[imp] = imp
}
func (u *Kit) Generate(plugin *protogen.Plugin) error {
	if len(plugin.Files) < 1 {
		return nil
	}
	u.addImports("context")
	u.addImports("git.diulo.com/mogfee/kit/transport/http")
	for _, f := range plugin.Files {
		if len(f.Services) == 0 {
			continue
		}
		fname := f.GeneratedFilenamePrefix + "_http.pb.go"
		t := plugin.NewGeneratedFile(fname, f.GoImportPath)
		t.P("package " + f.Desc.Name())
		t.P("import (")
		for _, v := range u.imports {
			t.P(`"` + v + `"`)
		}
		t.P(")")

		for _, s := range f.Services {
			t.P(`type `, s.GoName, `HTTPServer interface {`)
			for _, m := range s.Methods {
				t.P(m.GoName, `(context.Context, *`, m.Input.GoIdent.GoName, `) (*`, m.Output.GoIdent.GoName, `,error)`)
			}
			t.P(`}`)
		}

		for _, s := range f.Services {
			serverName := s.GoName
			t.P(`func Register`, serverName, `HTTPServer(s *http.Server,srv `, serverName, `Server) {`)
			t.P(`r:=s.Route("/")`)
			for _, m := range s.Methods {
				method, path := protogen2.GetProtoMethod(m)
				if method == "" {
					continue
				}
				t.P(`r.`, method, `("`, path, `",_`, s.GoName, `_`, m.GoName, `0_HTTP_Handler(srv))`)
			}
			t.P(`}`)
		}
		for _, s := range f.Services {
			for _, m := range s.Methods {
				method, _ := protogen2.GetProtoMethod(m)
				if method == "" {
					continue
				}
				u.genGet(f, s, t, m)
			}
		}
	}
	return nil
}

func (u *Kit) genGet(f *protogen.File, s *protogen.Service, t *protogen.GeneratedFile, m *protogen.Method) {
	method, path := protogen2.GetProtoMethod(m)
	if method == "" {
		return
	}

	t.P(`func _`, s.GoName, `_`, m.GoName, `0_HTTP_Handler(srv `, s.GoName, `HTTPServer) func(ctx http.Context) error {
	return func(ctx http.Context) error {
		var in `, m.Input.GoIdent.GoName)
	if method == protogen2.METHOD_GET {
		t.P(`if err := ctx.BindQuery(&in); err != nil {
			return err
		}`)
	} else if method == protogen2.METHOD_POST {
		t.P(`if err := ctx.Bind(&in); err != nil {
			return err
		}`)
	}
	if strings.LastIndexByte(path, '{') != -1 {
		t.P(`if err := ctx.BindVars(&in); err != nil {
			return err
		}`)
	}
	t.P(`http.SetOperation(ctx, "/`, f.Desc.Package(), `.`, s.Desc.Name(), `/`, m.Desc.Name(), `")
		h := ctx.Middleware(func(ctx context.Context, req interface{}) (interface{}, error) {
			return srv.`, m.GoName, `(ctx, req.(*`, m.Input.GoIdent.GoName, `))
		})
		out, err := h(ctx, &in)
		if err != nil {
			return err
		}
		reply := out.(*`, m.Output.GoIdent.GoName, `)
		return ctx.Result(200, reply)
	}
}`)
}