You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
1.9 KiB
Go
91 lines
1.9 KiB
Go
package zrpc
|
|
|
|
import (
|
|
"log"
|
|
"time"
|
|
|
|
"github.com/tal-tech/go-zero/core/discov"
|
|
"github.com/tal-tech/go-zero/zrpc/internal"
|
|
"github.com/tal-tech/go-zero/zrpc/internal/auth"
|
|
"google.golang.org/grpc"
|
|
)
|
|
|
|
var (
|
|
WithDialOption = internal.WithDialOption
|
|
WithTimeout = internal.WithTimeout
|
|
)
|
|
|
|
type (
|
|
ClientOption = internal.ClientOption
|
|
|
|
Client interface {
|
|
AddInterceptor(interceptor grpc.UnaryClientInterceptor)
|
|
Conn() *grpc.ClientConn
|
|
}
|
|
|
|
RpcClient struct {
|
|
client Client
|
|
}
|
|
)
|
|
|
|
func MustNewClient(c RpcClientConf, options ...ClientOption) Client {
|
|
cli, err := NewClient(c, options...)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
return cli
|
|
}
|
|
|
|
func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
|
|
var opts []ClientOption
|
|
if c.HasCredential() {
|
|
opts = append(opts, WithDialOption(grpc.WithPerRPCCredentials(&auth.Credential{
|
|
App: c.App,
|
|
Token: c.Token,
|
|
})))
|
|
}
|
|
if c.Timeout > 0 {
|
|
opts = append(opts, WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
|
|
}
|
|
opts = append(opts, options...)
|
|
|
|
var client Client
|
|
var err error
|
|
if len(c.Endpoints) > 0 {
|
|
client, err = internal.NewClient(internal.BuildDirectTarget(c.Endpoints), opts...)
|
|
} else if err = c.Etcd.Validate(); err == nil {
|
|
client, err = internal.NewClient(internal.BuildDiscovTarget(c.Etcd.Hosts, c.Etcd.Key), opts...)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &RpcClient{
|
|
client: client,
|
|
}, nil
|
|
}
|
|
|
|
func NewClientNoAuth(c discov.EtcdConf) (Client, error) {
|
|
client, err := internal.NewClient(internal.BuildDiscovTarget(c.Hosts, c.Key))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &RpcClient{
|
|
client: client,
|
|
}, nil
|
|
}
|
|
|
|
func NewClientWithTarget(target string, opts ...ClientOption) (Client, error) {
|
|
return internal.NewClient(target, opts...)
|
|
}
|
|
|
|
func (rc *RpcClient) AddInterceptor(interceptor grpc.UnaryClientInterceptor) {
|
|
rc.client.AddInterceptor(interceptor)
|
|
}
|
|
|
|
func (rc *RpcClient) Conn() *grpc.ClientConn {
|
|
return rc.client.Conn()
|
|
}
|