client: Allow creating a GoBGPClient using a user specific GRPC connection and client

This allows wrapping the dialer, implementing your own wrapped
GobgpApiClient or dialing with a context for cancelation on unreliable
endpoints.

1 files changed, 19 insertions, 2 deletions

diff --git a/client/client.go b/client/client.go
index b85360c2..7fa9d0ec 100644
--- a/client/client.go
+++ b/client/client.go
@@ -39,14 +39,23 @@ func defaultGRPCOptions() []grpc.DialOption {
 	return []grpc.DialOption{grpc.WithTimeout(time.Second), grpc.WithBlock(), grpc.WithInsecure()}
 }
 
-func NewGoBGPClient(target string, opts ...grpc.DialOption) (*GoBGPClient, error) {
+// New returns a new GoBGPClient using the given target and options for dialing
+// to the grpc server. If an error occurs during dialing it will be returned and
+// GoBGPClient will be nil.
+func New(target string, opts ...grpc.DialOption) (*GoBGPClient, error) {
+	return NewWith(context.Background(), target, opts...)
+}
+
+// NewWith is like New, but uses the given ctx to cancel or expire the current
+// attempt to connect if it becomes Done before the connection succeeds.
+func NewWith(ctx context.Context, target string, opts ...grpc.DialOption) (*GoBGPClient, error) {
 	if target == "" {
 		target = ":50051"
 	}
 	if len(opts) == 0 {
 		opts = defaultGRPCOptions()
 	}
-	conn, err := grpc.Dial(target, opts...)
+	conn, err := grpc.DialContext(ctx, target, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -54,6 +63,14 @@ func NewGoBGPClient(target string, opts ...grpc.DialOption) (*GoBGPClient, error
 	return &GoBGPClient{conn: conn, cli: cli}, nil
 }
 
+// NewFrom returns a new GoBGPClient, using the given conn and cli for the
+// underlying connection. The given grpc.ClientConn connection is expected to be
+// initialized and paired with the api client. See New to have the connection
+// dialed for you.
+func NewFrom(conn *grpc.ClientConn, cli api.GobgpApiClient) *GoBGPClient {
+	return &GoBGPClient{conn: conn, cli: cli}
+}
+
 func (cli *GoBGPClient) Close() error {
 	return cli.conn.Close()
 }