Merge pull request #11 from pepabo/with-close

Add `with_close` option to Add `Close()` method

README.md

@@ -17,3 +17,4 @@ Generate [client.go](example/gen/go/myapp/client.go) that bundles gRPC clients g
 | --- | --- | --- |
 | `package` | string | Specify package name (ex `--go-client_opt=package=xxxx` ) |
 | `same_package` | bool | Make the package the same as the package generated by proto-gen-go (ex `--go-client_opt=same_package` ) |
+| `with_close` | bool | Add `Close()` method for closing *grpc.ClientConn from client |

example/buf.gen.yaml

@@ -12,4 +12,4 @@ plugins:
     opt: paths=source_relative
   - plugin: go-client
     out: gen/go
-    opt: paths=source_relative,same_package
+    opt: paths=source_relative,same_package,with_close

example/example_test.go

@@ -18,6 +18,9 @@ func TestClient(t *testing.T) {
 	})
 	ts.ResponseDynamic()
 	client := myapp.New(ts.Conn())
+	t.Cleanup(func() {
+		client.Close()
+	})
 	if _, err := client.UserService().CreateUser(ctx, &myapp.CreateUserRequest{
 		Name:  "alice",
 		Email: "[email protected]",

gen/gen.go

@@ -15,11 +15,13 @@ import (
 const (
 	optPackage     = "package"
 	optSamePackage = "same_package"
+	optWithClose   = "with_close"
 )
 
 type Generator struct {
 	genp        *protogen.Plugin
 	samePackage bool
+	withClose   bool
 	packageName string
 }
 
@@ -56,13 +58,27 @@ func (gen *Generator) Generate() error {
 	}
 	g.P("")
 	g.P(`import (`)
+	if gen.withClose {
+		g.P(`"context"`)
+		g.P("")
+	}
 	if !gen.samePackage {
 		g.P(fmt.Sprintf("%s %s", tmppf.GoPackageName, tmppf.GoImportPath.String()))
 	}
 	g.P(`"google.golang.org/grpc"`)
 	g.P(`)`)
 
-	// type ClientInterface interface
+	if gen.withClose {
+		// type ClientConn interface
+		g.P(`type ClientConnInterface interface {`)
+		g.P(`Invoke(ctx context.Context, method string, args interface{}, reply interface{}, opts ...grpc.CallOption) error`)
+		g.P(`NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error)`)
+		g.P("Close() error")
+		g.P(`}`)
+		g.P("")
+	}
+
+	// type Client interface
 	g.P(`type Client interface {`)
 	for _, pf := range gen.genp.Files {
 		if !pf.Generate {
@@ -76,6 +92,9 @@ func (gen *Generator) Generate() error {
 			}
 		}
 	}
+	if gen.withClose {
+		g.P("Close() error")
+	}
 	g.P(`}`)
 	g.P("")
 
@@ -93,11 +112,18 @@ func (gen *Generator) Generate() error {
 			}
 		}
 	}
+	if gen.withClose {
+		g.P("cc ClientConnInterface")
+	}
 	g.P(`}`)
 	g.P("")
 
-	// func New(cc grpc.ClientConnInterface) *Client
-	g.P(`func New(cc grpc.ClientConnInterface) Client {`)
+	// func New(cc *.ClientConnInterface) *Client
+	if gen.withClose {
+		g.P(`func New(cc ClientConnInterface) Client {`)
+	} else {
+		g.P(`func New(cc grpc.ClientConnInterface) Client {`)
+	}
 	g.P(`return &client{`)
 	for _, pf := range gen.genp.Files {
 		if !pf.Generate {
@@ -111,6 +137,9 @@ func (gen *Generator) Generate() error {
 			}
 		}
 	}
+	if gen.withClose {
+		g.P("cc: cc,")
+	}
 	g.P(`}`)
 	g.P(`}`)
 	g.P("")
@@ -131,6 +160,12 @@ func (gen *Generator) Generate() error {
 			g.P("")
 		}
 	}
+	if gen.withClose {
+		g.P("func (c *client) Close() error {")
+		g.P("return c.cc.Close()")
+		g.P("}")
+		g.P("")
+	}
 	return nil
 }
 
@@ -141,6 +176,8 @@ func (gen *Generator) parseOpts() error {
 		switch {
 		case o == optSamePackage:
 			gen.samePackage = true
+		case o == optWithClose:
+			gen.withClose = true
 		case strings.HasPrefix(o, fmt.Sprintf("%s=", optPackage)):
 			gen.packageName = strings.TrimPrefix(o, fmt.Sprintf("%s=", optPackage))
 		}