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("context")
	u.addImports("git.diulo.com/mogfee/protoc-gen-kit/middleware")
	u.addImports("git.diulo.com/mogfee/protoc-gen-kit/response")
	u.addImports("git.diulo.com/mogfee/protoc-gen-kit/xerrors")
	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,m ...middleware.Middleware) {`)
			for _, m := range s.Methods {
				method, path := u.getMethod(m)
				if method == "" {
					continue
				}
				t.P(`app.`, method, `("`, path, `",http`, m.GoName, `Handler(srv,m...))`)
			}
			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,m ...middleware.Middleware)func(c *gin.Context){")
	t.P(`return func(c *gin.Context) {
		var post `, m.Input.GoIdent.GoName, `
		resp := response.New(c)
		if err := resp.BindQuery(&post); err != nil {
			resp.Error(err)
			return
		}
		h := func(ctx context.Context, a any) (any, error) {
			return srv.List(ctx, a.(*`, m.Input.GoIdent.GoName, `))
		}
		out, err := middleware.HttpMiddleware(c, h, m...)(c, &post)
		if err != nil {
			resp.Error(err)
		} else {
			if v, ok := out.(*`, m.Output.GoIdent.GoName, `); ok {
				resp.Success(v)
			} else {
				resp.Error(xerrors.InternalServer("RESULT_TYPE_ERROR", "`, m.Output.GoIdent.GoName, `"))
			}
		}
	}`)
	t.P("}")
}
func (u *Kit) genPost(serverName string, t *protogen.GeneratedFile, m *protogen.Method) {
	t.P("func http", m.GoName, "Handler(srv ", serverName, "Server,m ...middleware.Middleware)func(c *gin.Context){")
	t.P(`return func(c *gin.Context) {
		var post `, m.Input.GoIdent.GoName, `
		resp := response.New(c)
		if err := resp.BindJSON(&post); err != nil {
			resp.Error(err)
			return
		}
		h := func(ctx context.Context, a any) (any, error) {
			return srv.List(ctx, a.(*`, m.Input.GoIdent.GoName, `))
		}
		out, err := middleware.HttpMiddleware(c, h, m...)(c, &post)
		if err != nil {
			resp.Error(err)
		} else {
			if v, ok := out.(*`, m.Output.GoIdent.GoName, `); ok {
				resp.Success(v)
			} else {
				resp.Error(xerrors.InternalServer("RESULT_TYPE_ERROR", "`, m.Output.GoIdent.GoName, `"))
			}
		}
	}`)
	t.P("}")
}
func (u *Kit) genDelete(serverName string, t *protogen.GeneratedFile, m *protogen.Method) {
	t.P("func http", m.GoName, "Handler(srv ", serverName, "Server,m ...middleware.Middleware)func(c *gin.Context){")
	t.P(`return func(c *gin.Context) {
		var post `, m.Input.GoIdent.GoName, `
		resp := response.New(c)
		if err := resp.BindJSON(&post); err != nil {
			resp.Error(err)
			return
		}
		h := func(ctx context.Context, a any) (any, error) {
			return srv.List(ctx, a.(*`, m.Input.GoIdent.GoName, `))
		}
		out, err := middleware.HttpMiddleware(c, h, m...)(c, &post)
		if err != nil {
			resp.Error(err)
		} else {
			if v, ok := out.(*`, m.Output.GoIdent.GoName, `); ok {
				resp.Success(v)
			} else {
				resp.Error(xerrors.InternalServer("RESULT_TYPE_ERROR", "`, m.Output.GoIdent.GoName, `"))
			}
		}
	}`)
	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"
)