diff --git a/Makefile b/Makefile index df7f7c9..15ef67a 100644 --- a/Makefile +++ b/Makefile @@ -21,6 +21,7 @@ PG_FIX_CANDIDATES=./protos/node/node.pb.go \ ./protos/status/status.pb.go \ ./protos/escrow/escrow.pb.go \ ./protos/guard/guard.pb.go \ + ./protos/online/online.pb.go \ install: brew trongogo diff --git a/utils/grpc/client.go b/utils/grpc/client.go index ad1c6b3..eafbf49 100644 --- a/utils/grpc/client.go +++ b/utils/grpc/client.go @@ -5,6 +5,7 @@ import ( "crypto/tls" "errors" "fmt" + "github.com/tron-us/go-btfs-common/protos/online" "net/url" "strconv" "strings" @@ -50,6 +51,8 @@ func (g *ClientBuilder) doWithContext(ctx context.Context, f interface{}) error return err } switch v := f.(type) { + case func(context.Context, online.OnlineServiceClient) error: + return wrapError("OnlineClient", v(ctx, online.NewOnlineServiceClient(conn))) case func(context.Context, statuspb.StatusServiceClient) error: return wrapError("StatusClient", v(ctx, statuspb.NewStatusServiceClient(conn))) case func(context.Context, hubpb.HubQueryServiceClient) error: diff --git a/utils/grpc/online.go b/utils/grpc/online.go new file mode 100644 index 0000000..75bdcec --- /dev/null +++ b/utils/grpc/online.go @@ -0,0 +1,19 @@ +package grpc + +import ( + "context" + "github.com/tron-us/go-btfs-common/protos/online" +) + +func OnlineClient(addr string) *OnlineClientBuilder { + return &OnlineClientBuilder{builder(addr)} +} + +type OnlineClientBuilder struct { + ClientBuilder +} + +func (g *OnlineClientBuilder) WithContext(ctx context.Context, f func(ctx context.Context, + client online.OnlineServiceClient) error) error { + return g.doWithContext(ctx, f) +} diff --git a/utils/grpc/setup_server.go b/utils/grpc/setup_server.go index ab2f380..cdae74e 100644 --- a/utils/grpc/setup_server.go +++ b/utils/grpc/setup_server.go @@ -3,6 +3,7 @@ package grpc import ( "context" "fmt" + "github.com/tron-us/go-btfs-common/protos/online" "net" "net/http" "os" @@ -55,6 +56,8 @@ type GrpcServer struct { func (s *GrpcServer) serverTypeToServerName(server interface{}) { switch t := server.(type) { + case online.OnlineServiceServer: + s.serverName = "online-server" case status.StatusServiceServer: s.serverName = "status-server" case escrow.EscrowServiceServer: @@ -173,6 +176,8 @@ func (s *GrpcServer) CreateHealthServer() *GrpcServer { func (s *GrpcServer) RegisterServer(server interface{}) *GrpcServer { switch server.(type) { + case online.OnlineServiceServer: + online.RegisterOnlineServiceServer(s.server, server.(online.OnlineServiceServer)) case status.StatusServiceServer: status.RegisterStatusServiceServer(s.server, server.(status.StatusServiceServer)) case escrow.EscrowServiceServer: