forked from samuel/go-thrift
/
client.go
95 lines (84 loc) · 2.55 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package thrift
import (
"io"
"net"
"net/rpc"
)
// Implements rpc.ClientCodec
type clientCodec struct {
transport io.ReadWriteCloser
protocol Protocol
}
// Dial connects to a Thrift RPC server at the specified network address using the specified protocol.
func Dial(network, address string, framed bool, protocol Protocol) (*rpc.Client, error) {
conn, err := net.Dial(network, address)
if err != nil {
return nil, err
}
codec := &clientCodec{
transport: conn,
protocol: protocol,
}
if framed {
codec.transport = NewFramedReadWriteCloser(conn, DefaultMaxFrameSize)
}
return rpc.NewClientWithCodec(codec), nil
}
// NewClient returns a new rpc.Client to handle requests to the set of
// services at the other end of the connection.
func NewClient(conn io.ReadWriteCloser, protocol Protocol) *rpc.Client {
return rpc.NewClientWithCodec(NewClientCodec(conn, protocol))
}
// NewClientCodec returns a new rpc.ClientCodec using Thrift RPC on conn using the specified protocol.
func NewClientCodec(conn io.ReadWriteCloser, protocol Protocol) rpc.ClientCodec {
return &clientCodec{
transport: conn,
protocol: protocol,
}
}
func (c *clientCodec) WriteRequest(request *rpc.Request, thriftStruct interface{}) error {
if err := c.protocol.WriteMessageBegin(c.transport, request.ServiceMethod, messageTypeCall, int32(request.Seq)); err != nil {
return err
}
if err := EncodeStruct(c.transport, c.protocol, thriftStruct); err != nil {
return err
}
if err := c.protocol.WriteMessageEnd(c.transport); err != nil {
return err
}
if flusher, ok := c.transport.(Flusher); ok {
return flusher.Flush()
}
return nil
}
func (c *clientCodec) ReadResponseHeader(response *rpc.Response) error {
name, messageType, seq, err := c.protocol.ReadMessageBegin(c.transport)
if err != nil {
return err
}
response.ServiceMethod = name
response.Seq = uint64(seq)
if messageType == messageTypeException {
exception := &ApplicationException{}
if err := DecodeStruct(c.transport, c.protocol, exception); err != nil {
return err
}
response.Error = exception.String()
return c.protocol.ReadMessageEnd(c.transport)
}
return nil
}
func (c *clientCodec) ReadResponseBody(thriftStruct interface{}) error {
if thriftStruct == nil {
// Should only get called if ReadResponseHeader set the Error value in
// which case we've already read the body (ApplicationException)
return nil
}
if err := DecodeStruct(c.transport, c.protocol, thriftStruct); err != nil {
return err
}
return c.protocol.ReadMessageEnd(c.transport)
}
func (c *clientCodec) Close() error {
return nil
}