This repository has been archived by the owner on Feb 21, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 230
/
client.go
131 lines (115 loc) · 3.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
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
package fbcloud
import (
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
featurebase "github.com/featurebasedb/featurebase/v3"
"github.com/pkg/errors"
)
// TokenRefreshTimeout is currently hardcoded to be just under the
// Cognito token timeout for cloud which is 15 minutes (I think I
// heard that somewhere anyway). It seems to work.
const TokenRefreshTimeout = time.Minute * 13
type Queryer struct {
Host string
ClientID string
Region string
Email string
Password string
token string
lastRefresh time.Time
}
func (cq *Queryer) tokenRefresh() error {
token, err := authenticate(cq.ClientID, cq.Region, cq.Email, cq.Password)
if err != nil {
return errors.Wrap(err, "getting token")
}
cq.token = token
cq.lastRefresh = time.Now()
return nil
}
// Query issues a SQL query formatted for the FeatureBase cloud query endpoint.
func (cq *Queryer) Query(org string, db string, sql io.Reader) (*featurebase.WireQueryResponse, error) {
if time.Since(cq.lastRefresh) > TokenRefreshTimeout {
if err := cq.tokenRefresh(); err != nil {
return nil, errors.Wrap(err, "refreshing token")
}
}
url := fmt.Sprintf("%s/databases/%s/sql", cq.Host, db)
if db == "" {
url = fmt.Sprintf("%s/sql", cq.Host)
}
client := &http.Client{
Timeout: time.Second * 30,
}
req, err := http.NewRequest(http.MethodPost, url, sql)
if err != nil {
return nil, errors.Wrap(err, "creating new post request")
}
req.Header.Add("Content-Type", "text/plain")
req.Header.Add("Authorization", cq.token)
var resp *http.Response
if resp, err = client.Do(req); err != nil {
return nil, errors.Wrap(err, "executing post request")
}
fullbod, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Wrap(err, "reading cloud response")
}
if resp.StatusCode/100 != 2 {
return nil, errors.Errorf("unexpected status: %s, full body: '%s'", resp.Status, fullbod)
}
var sqlResponse featurebase.WireQueryResponse
if err := json.Unmarshal(fullbod, &sqlResponse); err != nil {
return nil, errors.Wrapf(err, "decoding cloud response, body:\n%s", fullbod)
}
return &sqlResponse, nil
}
// HTTPRequest can make an arbitrary http request to the host and
// tries to json unmarshal the response body into v if v is
// non-nil. This is handy for hitting cloud endpoints other than the
// query endpoint which is handled by Query. I don't think this is
// currently used, but I'd like to keep it around for debugging.
func (cq *Queryer) HTTPRequest(method, path, body string, v interface{}) ([]byte, error) {
if time.Since(cq.lastRefresh) > TokenRefreshTimeout {
if err := cq.tokenRefresh(); err != nil {
return nil, errors.Wrap(err, "refreshing token")
}
}
var bod io.Reader
if body == "" {
bod = nil
} else {
bod = strings.NewReader(body)
}
req, err := http.NewRequest(method, fmt.Sprintf("%s%s", cq.Host, path), bod)
if err != nil {
return nil, errors.Errorf("creating request: %v", err)
}
// fmt.Printf("%+v\n", req)
req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", cq.token))
if bod != nil {
req.Header.Add("Content-Type", "application/json")
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, errors.Errorf("doing request: %v", err)
}
bodbytes, err := io.ReadAll(resp.Body)
if err != nil {
return nil, errors.Errorf("reading response body: %v", err)
}
if resp.StatusCode/100 != 2 {
return nil, errors.Errorf("bad status: %s. body: '%s'", resp.Status, bodbytes)
}
if v != nil {
err = json.Unmarshal(bodbytes, v)
if err != nil {
return nil, errors.Errorf("unmarshaling: %v", err)
}
}
return bodbytes, nil
}