diff --git a/http.go b/http.go index 9a18680..2fd99a6 100644 --- a/http.go +++ b/http.go @@ -4,7 +4,7 @@ import ( "google.golang.org/protobuf/compiler/protogen" ) -func (g *Generator) httpGenerate(component string, plugin *protogen.Plugin) error { +func (g *Generator) httpGenerate(component string, plugin *protogen.Plugin, genClient bool, genServer bool) error { for _, file := range plugin.Files { if !file.Generate { continue @@ -27,16 +27,24 @@ func (g *Generator) httpGenerate(component string, plugin *protogen.Plugin) erro gfile.Import(contextPackage) gfile.Import(microApiPackage) - gfile.Import(microClientPackage) - gfile.Import(microClientHttpPackage) - gfile.Import(microServerPackage) + if genClient { + gfile.Import(microClientPackage) + gfile.Import(microClientHttpPackage) + } + if genServer { + gfile.Import(microServerPackage) + } for _, service := range file.Services { - generateServiceClient(gfile, service) - generateServiceClientMethods(gfile, service, true) - generateServiceServer(gfile, service) - generateServiceServerMethods(gfile, service) - generateServiceRegister(gfile, service) + if genClient { + generateServiceClient(gfile, service) + generateServiceClientMethods(gfile, service, true) + } + if genServer { + generateServiceServer(gfile, service) + generateServiceServerMethods(gfile, service) + generateServiceRegister(gfile, service) + } } } diff --git a/main.go b/main.go index c6c0b97..18411ca 100644 --- a/main.go +++ b/main.go @@ -12,7 +12,7 @@ import ( var ( flagDebug = flag.Bool("debug", false, "") flagStandalone = flag.Bool("standalone", false, "") - flagComponents = flag.String("components", "micro|rpc", "") + flagComponents = flag.String("components", "micro|rpc|client|server", "") ) func main() { @@ -39,15 +39,27 @@ func (g *Generator) Generate(plugin *protogen.Plugin) error { g.components = *flagComponents plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL) + var genClient bool + var genServer bool + + if strings.Contains(g.components, "server") { + genServer = true + } + if strings.Contains(g.components, "client") { + genClient = true + } + // Protoc passes a slice of File structs for us to process for _, component := range strings.Split(g.components, "|") { switch component { + case "server", "client": + continue case "micro": - err = g.microGenerate(component, plugin) + err = g.microGenerate(component, plugin, genClient, genServer) case "http": - err = g.httpGenerate(component, plugin) + err = g.httpGenerate(component, plugin, genClient, genServer) case "grpc", "rpc": - err = g.rpcGenerate("rpc", plugin) + err = g.rpcGenerate("rpc", plugin, genClient, genServer) case "gorilla": err = g.gorillaGenerate(component, plugin) case "chi": diff --git a/micro.go b/micro.go index 98f70f4..fa534a3 100644 --- a/micro.go +++ b/micro.go @@ -4,7 +4,7 @@ import ( "google.golang.org/protobuf/compiler/protogen" ) -func (g *Generator) microGenerate(component string, plugin *protogen.Plugin) error { +func (g *Generator) microGenerate(component string, plugin *protogen.Plugin, genClient bool, genServer bool) error { for _, file := range plugin.Files { if !file.Generate { continue @@ -28,15 +28,20 @@ func (g *Generator) microGenerate(component string, plugin *protogen.Plugin) err gfile.Import(contextPackage) gfile.Import(microApiPackage) - gfile.Import(microClientPackage) - + if genClient { + gfile.Import(microClientPackage) + } // generate services for _, service := range file.Services { generateServiceEndpoints(gfile, service) - generateServiceClientInterface(gfile, service) - generateServiceClientStreamInterface(gfile, service) - generateServiceServerInterface(gfile, service) - generateServiceServerStreamInterface(gfile, service) + if genClient { + generateServiceClientInterface(gfile, service) + generateServiceClientStreamInterface(gfile, service) + } + if genServer { + generateServiceServerInterface(gfile, service) + generateServiceServerStreamInterface(gfile, service) + } } } diff --git a/rpc.go b/rpc.go index ba5e4dd..75fa385 100644 --- a/rpc.go +++ b/rpc.go @@ -4,7 +4,7 @@ import ( "google.golang.org/protobuf/compiler/protogen" ) -func (g *Generator) rpcGenerate(component string, plugin *protogen.Plugin) error { +func (g *Generator) rpcGenerate(component string, plugin *protogen.Plugin, genClient bool, genServer bool) error { for _, file := range plugin.Files { if !file.Generate { continue @@ -27,15 +27,22 @@ func (g *Generator) rpcGenerate(component string, plugin *protogen.Plugin) error gfile.Import(contextPackage) gfile.Import(microApiPackage) - gfile.Import(microClientPackage) - gfile.Import(microServerPackage) - + if genClient { + gfile.Import(microClientPackage) + } + if genServer { + gfile.Import(microServerPackage) + } for _, service := range file.Services { - generateServiceClient(gfile, service) - generateServiceClientMethods(gfile, service, false) - generateServiceServer(gfile, service) - generateServiceServerMethods(gfile, service) - generateServiceRegister(gfile, service) + if genClient { + generateServiceClient(gfile, service) + generateServiceClientMethods(gfile, service, false) + } + if genServer { + generateServiceServer(gfile, service) + generateServiceServerMethods(gfile, service) + generateServiceRegister(gfile, service) + } } } diff --git a/util.go b/util.go index a28b749..94a640c 100644 --- a/util.go +++ b/util.go @@ -283,10 +283,11 @@ func generateServiceRegister(gfile *protogen.GeneratedFile, service *protogen.Se gfile.P(unexport(serviceName)) gfile.P("}") gfile.P("h := &", unexport(serviceName), "Server{sh}") + gfile.P("var nopts []", microServerPackage.Ident("HandlerOption")) gfile.P("for _, endpoint := range New", serviceName, "Endpoints() {") - gfile.P("opts = append(opts, ", microApiPackage.Ident("WithEndpoint"), "(endpoint))") + gfile.P("nopts = append(nopts, ", microApiPackage.Ident("WithEndpoint"), "(endpoint))") gfile.P("}") - gfile.P("return s.Handle(s.NewHandler(&", serviceName, "{h}, opts...))") + gfile.P("return s.Handle(s.NewHandler(&", serviceName, "{h}, append(nopts, opts...)...))") gfile.P("}") } @@ -559,11 +560,11 @@ func generateEndpoint(gfile *protogen.GeneratedFile, serviceName string, methodN path, meth, body := getEndpoint(rule) gfile.P("Name:", fmt.Sprintf(`"%s.%s",`, serviceName, methodName)) gfile.P("Path:", fmt.Sprintf(`[]string{"%s"},`, path)) - if vmethod, ok := httpMethodMap[meth]; ok { - gfile.P("Method:", `[]string{`, httpPackage.Ident(vmethod), `},`) - } else { - gfile.P("Method:", fmt.Sprintf(`[]string{"%s"},`, meth)) - } + //if vmethod, ok := httpMethodMap[meth]; ok { + // gfile.P("Method:", `[]string{`, httpPackage.Ident(vmethod), `},`) + //} else { + gfile.P("Method:", fmt.Sprintf(`[]string{"%s"},`, meth)) + // } if len(rule.GetGet()) == 0 && body != "" { gfile.P("Body:", fmt.Sprintf(`"%s",`, body)) }