forked from schwartzmx/gremgo-neptune
/
client.go
148 lines (128 loc) · 3.85 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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package gremgo
import (
"errors"
"io/ioutil"
"log"
"sync"
"time"
)
// Client is a container for the gremgo client.
type Client struct {
conn dialer
requests chan []byte
responses chan []byte
results *sync.Map
responseNotifier *sync.Map // responseNotifier notifies the requester that a response has arrived for the request
respMutex *sync.Mutex
Errored bool
}
// NewDialer returns a WebSocket dialer to use when connecting to Gremlin Server
func NewDialer(host string, configs ...DialerConfig) (dialer *Ws) {
dialer = &Ws{
timeout: 5 * time.Second,
pingInterval: 60 * time.Second,
writingWait: 15 * time.Second,
readingWait: 15 * time.Second,
connected: false,
quit: make(chan struct{}),
}
for _, conf := range configs {
conf(dialer)
}
dialer.host = host
return dialer
}
func newClient() (c Client) {
c.requests = make(chan []byte, 3) // c.requests takes any request and delivers it to the WriteWorker for dispatch to Gremlin Server
c.responses = make(chan []byte, 3) // c.responses takes raw responses from ReadWorker and delivers it for sorting to handelResponse
c.results = &sync.Map{}
c.responseNotifier = &sync.Map{}
c.respMutex = &sync.Mutex{} // c.mutex ensures that sorting is thread safe
return
}
// Dial returns a gremgo client for interaction with the Gremlin Server specified in the host IP.
func Dial(conn dialer, errs chan error) (c Client, err error) {
c = newClient()
c.conn = conn
// Connects to Gremlin Server
err = conn.connect()
if err != nil {
return
}
quit := conn.(*Ws).quit
go c.writeWorker(errs, quit)
go c.readWorker(errs, quit)
go conn.ping(errs)
return
}
func (c *Client) executeRequest(query string, bindings, rebindings *map[string]string) (resp []Response, err error) {
var req request
var id string
if bindings != nil && rebindings != nil {
req, id, err = prepareRequestWithBindings(query, *bindings, *rebindings)
} else {
req, id, err = prepareRequest(query)
}
if err != nil {
return
}
msg, err := packageRequest(req)
if err != nil {
log.Println(err)
return
}
c.responseNotifier.Store(id, make(chan int, 1))
c.dispatchRequest(msg)
resp = c.retrieveResponse(id)
return
}
func (c *Client) authenticate(requestID string) (err error) {
auth := c.conn.getAuth()
req, err := prepareAuthRequest(requestID, auth.username, auth.password)
if err != nil {
return
}
msg, err := packageRequest(req)
if err != nil {
log.Println(err)
return
}
c.dispatchRequest(msg)
return
}
// ExecuteWithBindings formats a raw Gremlin query, sends it to Gremlin Server, and returns the result.
func (c *Client) ExecuteWithBindings(query string, bindings, rebindings map[string]string) (resp []Response, err error) {
if c.conn.isDisposed() {
return resp, errors.New("you cannot write on disposed connection")
}
resp, err = c.executeRequest(query, &bindings, &rebindings)
return
}
// Execute formats a raw Gremlin query, sends it to Gremlin Server, and returns the result.
func (c *Client) Execute(query string) (resp []Response, err error) {
if c.conn.isDisposed() {
return resp, errors.New("you cannot write on disposed connection")
}
resp, err = c.executeRequest(query, nil, nil)
return
}
// ExecuteFile takes a file path to a Gremlin script, sends it to Gremlin Server, and returns the result.
func (c *Client) ExecuteFile(path string, bindings, rebindings map[string]string) (resp []Response, err error) {
if c.conn.isDisposed() {
return resp, errors.New("you cannot write on disposed connection")
}
d, err := ioutil.ReadFile(path) // Read script from file
if err != nil {
log.Println(err)
return
}
query := string(d)
resp, err = c.executeRequest(query, &bindings, &rebindings)
return
}
// Close closes the underlying connection and marks the client as closed.
func (c *Client) Close() {
if c.conn != nil {
c.conn.close()
}
}