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
|
// Copyright 2019-present Facebook Inc. All rights reserved.
// This source code is licensed under the Apache 2.0 license found
// in the LICENSE file in the root directory of this source tree.
package gremlin
import (
"bytes"
"encoding/base64"
"time"
"github.com/google/uuid"
"github.com/pkg/errors"
)
type (
// A Request models a request message sent to the server.
Request struct {
RequestID string `json:"requestId" graphson:"g:UUID"`
Operation string `json:"op"`
Processor string `json:"processor"`
Arguments map[string]interface{} `json:"args"`
}
// RequestOption enables request customization.
RequestOption func(*Request)
// Credentials holds request plain auth credentials.
Credentials struct{ Username, Password string }
)
// NewEvalRequest returns a new evaluation request request.
func NewEvalRequest(query string, opts ...RequestOption) *Request {
r := &Request{
RequestID: uuid.New().String(),
Operation: OpsEval,
Arguments: map[string]interface{}{
ArgsGremlin: query,
ArgsLanguage: "gremlin-groovy",
},
}
for i := range opts {
opts[i](r)
}
return r
}
// NewAuthRequest returns a new auth request.
func NewAuthRequest(requestID, username, password string) *Request {
return &Request{
RequestID: requestID,
Operation: OpsAuthentication,
Arguments: map[string]interface{}{
ArgsSasl: Credentials{
Username: username,
Password: password,
},
ArgsSaslMechanism: "PLAIN",
},
}
}
// WithBindings sets request bindings.
func WithBindings(bindings map[string]interface{}) RequestOption {
return func(r *Request) {
r.Arguments[ArgsBindings] = bindings
}
}
// WithEvalTimeout sets script evaluation timeout.
func WithEvalTimeout(timeout time.Duration) RequestOption {
return func(r *Request) {
r.Arguments[ArgsEvalTimeout] = int64(timeout / time.Millisecond)
}
}
// MarshalText implements encoding.TextMarshaler interface.
func (c Credentials) MarshalText() ([]byte, error) {
var buf bytes.Buffer
buf.WriteByte(0)
buf.WriteString(c.Username)
buf.WriteByte(0)
buf.WriteString(c.Password)
enc := base64.StdEncoding
text := make([]byte, enc.EncodedLen(buf.Len()))
enc.Encode(text, buf.Bytes())
return text, nil
}
// UnmarshalText implements encoding.TextUnmarshaler interface.
func (c *Credentials) UnmarshalText(text []byte) error {
enc := base64.StdEncoding
data := make([]byte, enc.DecodedLen(len(text)))
n, err := enc.Decode(data, text)
if err != nil {
return err
}
data = data[:n]
parts := bytes.SplitN(data, []byte{0}, 3)
if len(parts) != 3 {
return errors.New("bad credentials data")
}
c.Username = string(parts[1])
c.Password = string(parts[2])
return nil
}
|
