protoc-gen-kit/main.go

package main

import (
	"github.com/golang/protobuf/proto"
	"google.golang.org/genproto/googleapis/api/annotations"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/descriptorpb"
	"log"
)

//go build && protoc --unknow_out=./proto --go_out=./proto/ --go-grpc_out=./proto proto/user.proto -I ./proto
// protoc --unknow_out=./proto proto/user.proto
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("git.echinacities.com/mogfee/protoc-gen-kit/response")
	u.addImports("github.com/gin-gonic/gin")
	for _, f := range plugin.Files {
		if len(f.Services) == 0 {
			continue
		}
		fname := f.GeneratedFilenamePrefix + ".gin.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 {
			serverName := s.GoName
			t.P(`func Register`, serverName, `Handler(app *gin.Engine,srv `, serverName, `Server) {`)
			for _, m := range s.Methods {
				method, path := u.getMethod(m)
				if method == "" {
					continue
				}
				t.P(`app.`, method, `("`, path, `",http`, m.GoName, `Handler(srv))`)
			}
			t.P(`}`)
		}
		for _, s := range f.Services {
			serverName := s.GoName
			for _, m := range s.Methods {
				method, _ := u.getMethod(m)
				if method == "" {
					continue
				}
				switch method {
				case method_get:
					u.genGet(serverName, t, m)
				case method_post:
					u.genPost(serverName, t, m)
				case method_delete:
					u.genDelete(serverName, t, m)

				}
			}
		}
	}
	return nil
}

func (u *Kit) genGet(serverName string, t *protogen.GeneratedFile, m *protogen.Method) {
	t.P("func http", m.GoName, "Handler(srv ", serverName, "Server)func(c *gin.Context){")
	t.P(`return func(c *gin.Context) {
		post := `, m.Input.GoIdent.GoName, `{}
		resp := response.New(c)
		if err := resp.BindQuery(&post); err != nil {
			resp.Error(err)
			return
		}
		result, err := srv.`, m.GoName, `(c, &post)
		if err != nil {
			resp.Error(err)
			return
		}
		resp.Success(result)
	}`)
	t.P("}")
}
func (u *Kit) genPost(serverName string, t *protogen.GeneratedFile, m *protogen.Method) {
	t.P("func http", m.GoName, "Handler(srv ", serverName, "Server)func(c *gin.Context){")
	t.P(`return func(c *gin.Context) {
		post := `, m.Input.GoIdent.GoName, `{}
		resp := response.New(c)
		if err := resp.BindJSON(&post); err != nil {
			resp.Error(err)
			return
		}
		result, err := srv.`, m.GoName, `(c, &post)
		if err != nil {
			resp.Error(err)
			return
		}
		resp.Success(result)
	}`)
	t.P("}")
}
func (u *Kit) genDelete(serverName string, t *protogen.GeneratedFile, m *protogen.Method) {
	t.P("func http", m.GoName, "Handler(srv ", serverName, "Server)func(c *gin.Context){")
	t.P(`return func(c *gin.Context) {
		post := `, m.Input.GoIdent.GoName, `{}
		resp := response.New(c)
		if err := resp.BindJSON(&post); err != nil {
			resp.Error(err)
			return
		}
		result, err := srv.`, m.GoName, `(c, &post)
		if err != nil {
			resp.Error(err)
			return
		}
		resp.Success(result)
	}`)
	t.P("}")
}

//func (Kit) getAuth(m *protogen.Method) *auth.AuthInfo {
//	if op, ok := m.Desc.Options().(*descriptorpb.MethodOptions); ok {
//		if opts, err := proto.GetExtension(op, auth.E_Auth); err != nil {
//			log.Println(err)
//		} else {
//			if vv, ok := opts.(*auth.AuthInfo); ok {
//				return vv
//			}
//		}
//	}
//	return &auth.AuthInfo{}
//}
func (Kit) getMethod(m *protogen.Method) (method string, path string) {
	if op, ok := m.Desc.Options().(*descriptorpb.MethodOptions); ok {
		if opts, err := proto.GetExtension(op, annotations.E_Http); err != nil {
			log.Println(err)
		} else {
			if vv, ok := opts.(*annotations.HttpRule); ok {
				if vvv, ok := vv.Pattern.(*annotations.HttpRule_Get); ok {
					return method_get, vvv.Get
				} else if vvv, ok := vv.Pattern.(*annotations.HttpRule_Post); ok {
					return method_post, vvv.Post
				} else if vvv, ok := vv.Pattern.(*annotations.HttpRule_Delete); ok {
					return method_delete, vvv.Delete
				}
			}
		}
	}
	return "", ""
}

const (
	method_get    = "GET"
	method_post   = "POST"
	method_delete = "DELETE"
)