/influxdb/client.go

https://github.com/influxdb/kapacitor · Go · 548 lines · 414 code · 73 blank · 61 comment · 80 complexity · 0f6166d27a678d907608d8ba112fc133 MD5 · raw file

  1. package influxdb
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "io/ioutil"
  8. "net/http"
  9. "net/url"
  10. "strconv"
  11. "sync"
  12. "sync/atomic"
  13. "time"
  14. imodels "github.com/influxdata/influxdb/models"
  15. khttp "github.com/influxdata/kapacitor/http"
  16. "github.com/pkg/errors"
  17. )
  18. // Client is an interface for writing to and querying from an InfluxDB instance.
  19. type Client interface {
  20. // Ping checks that status of cluster
  21. // The provided context can be used to cancel the request.
  22. Ping(ctx context.Context) (time.Duration, string, error)
  23. // Write takes a BatchPoints object and writes all Points to InfluxDB.
  24. Write(bp BatchPoints) error
  25. // Query makes an InfluxDB Query on the database.
  26. // The response is checked for an error and the is returned
  27. // if it exists
  28. Query(q Query) (*Response, error)
  29. }
  30. type ClientUpdater interface {
  31. Client
  32. Update(new Config) error
  33. Close() error
  34. }
  35. // BatchPointsConfig is the config data needed to create an instance of the BatchPoints struct
  36. type BatchPointsConfig struct {
  37. // Precision is the write precision of the points, defaults to "ns"
  38. Precision string
  39. // Database is the database to write points to
  40. Database string
  41. // RetentionPolicy is the retention policy of the points
  42. RetentionPolicy string
  43. // Write consistency is the number of servers required to confirm write
  44. WriteConsistency string
  45. }
  46. // Query defines a query to send to the server
  47. type Query struct {
  48. Command string
  49. Database string
  50. Precision string
  51. }
  52. // HTTPConfig is the config data needed to create an HTTP Client
  53. type Config struct {
  54. // The URL of the InfluxDB server.
  55. URLs []string
  56. // Optional credentials for authenticating with the server.
  57. Credentials Credentials
  58. // UserAgent is the http User Agent, defaults to "KapacitorInfluxDBClient"
  59. UserAgent string
  60. // Timeout for requests, defaults to no timeout.
  61. Timeout time.Duration
  62. // Transport is the HTTP transport to use for requests
  63. // If nil, a default transport will be used.
  64. Transport *http.Transport
  65. }
  66. // AuthenticationMethod defines the type of authentication used.
  67. type AuthenticationMethod int
  68. // Supported authentication methods.
  69. const (
  70. NoAuthentication AuthenticationMethod = iota
  71. UserAuthentication
  72. BearerAuthentication
  73. )
  74. // Set of credentials depending on the authentication method
  75. type Credentials struct {
  76. Method AuthenticationMethod
  77. // UserAuthentication fields
  78. Username string
  79. Password string
  80. // BearerAuthentication fields
  81. Token string
  82. }
  83. // HTTPClient is safe for concurrent use.
  84. type HTTPClient struct {
  85. mu sync.RWMutex
  86. config Config
  87. urls []url.URL
  88. client *http.Client
  89. index int32
  90. }
  91. // NewHTTPClient returns a new Client from the provided config.
  92. // Client is safe for concurrent use by multiple goroutines.
  93. func NewHTTPClient(conf Config) (*HTTPClient, error) {
  94. if conf.UserAgent == "" {
  95. conf.UserAgent = "KapacitorInfluxDBClient"
  96. }
  97. urls, err := parseURLs(conf.URLs)
  98. if err != nil {
  99. return nil, errors.Wrap(err, "invalid URLs")
  100. }
  101. if conf.Transport == nil {
  102. conf.Transport = khttp.NewDefaultTransport()
  103. }
  104. c := &HTTPClient{
  105. config: conf,
  106. urls: urls,
  107. client: &http.Client{
  108. Timeout: conf.Timeout,
  109. Transport: conf.Transport,
  110. },
  111. }
  112. return c, nil
  113. }
  114. func parseURLs(urlStrs []string) ([]url.URL, error) {
  115. urls := make([]url.URL, len(urlStrs))
  116. for i, urlStr := range urlStrs {
  117. u, err := url.Parse(urlStr)
  118. if err != nil {
  119. return nil, err
  120. } else if u.Scheme != "http" && u.Scheme != "https" {
  121. return nil, fmt.Errorf(
  122. "Unsupported protocol scheme: %s, your address must start with http:// or https://",
  123. u.Scheme,
  124. )
  125. }
  126. urls[i] = *u
  127. }
  128. return urls, nil
  129. }
  130. func (c *HTTPClient) loadConfig() Config {
  131. c.mu.RLock()
  132. config := c.config
  133. c.mu.RUnlock()
  134. return config
  135. }
  136. func (c *HTTPClient) loadURLs() []url.URL {
  137. c.mu.RLock()
  138. urls := c.urls
  139. c.mu.RUnlock()
  140. return urls
  141. }
  142. func (c *HTTPClient) loadHTTPClient() *http.Client {
  143. c.mu.RLock()
  144. client := c.client
  145. c.mu.RUnlock()
  146. return client
  147. }
  148. func (c *HTTPClient) Close() error {
  149. return nil
  150. }
  151. // UpdateURLs updates the running list of URLs.
  152. func (c *HTTPClient) Update(new Config) error {
  153. if new.UserAgent == "" {
  154. new.UserAgent = "KapacitorInfluxDBClient"
  155. }
  156. c.mu.Lock()
  157. defer c.mu.Unlock()
  158. old := c.config
  159. c.config = new
  160. // Replace urls
  161. urls, err := parseURLs(new.URLs)
  162. if err != nil {
  163. return err
  164. }
  165. c.urls = urls
  166. if old.Credentials != new.Credentials ||
  167. old.Timeout != new.Timeout ||
  168. old.Transport != new.Transport {
  169. //Replace the client
  170. tr := new.Transport
  171. if tr == nil {
  172. tr = old.Transport
  173. }
  174. c.client = &http.Client{
  175. Timeout: new.Timeout,
  176. Transport: tr,
  177. }
  178. }
  179. return nil
  180. }
  181. func (c *HTTPClient) url() url.URL {
  182. urls := c.loadURLs()
  183. i := atomic.LoadInt32(&c.index)
  184. i = (i + 1) % int32(len(urls))
  185. atomic.StoreInt32(&c.index, i)
  186. return urls[i]
  187. }
  188. func (c *HTTPClient) do(req *http.Request, result interface{}, codes ...int) (*http.Response, error) {
  189. // Get current config
  190. config := c.loadConfig()
  191. // Set auth credentials
  192. cred := config.Credentials
  193. switch cred.Method {
  194. case NoAuthentication:
  195. case UserAuthentication:
  196. req.SetBasicAuth(cred.Username, cred.Password)
  197. case BearerAuthentication:
  198. req.Header.Set("Authorization", "Bearer "+cred.Token)
  199. default:
  200. return nil, errors.New("unknown authentication method set")
  201. }
  202. // Set user agent
  203. req.Header.Set("User-Agent", config.UserAgent)
  204. // Get client
  205. client := c.loadHTTPClient()
  206. // Do request
  207. resp, err := client.Do(req)
  208. if err != nil {
  209. return nil, err
  210. }
  211. defer resp.Body.Close()
  212. valid := false
  213. for _, code := range codes {
  214. if resp.StatusCode == code {
  215. valid = true
  216. break
  217. }
  218. }
  219. if !valid {
  220. body, err := ioutil.ReadAll(resp.Body)
  221. if err != nil {
  222. return nil, err
  223. }
  224. d := json.NewDecoder(bytes.NewReader(body))
  225. rp := struct {
  226. Error string `json:"error"`
  227. }{}
  228. if err := d.Decode(&rp); err != nil {
  229. return nil, err
  230. }
  231. if rp.Error != "" {
  232. return nil, errors.New(rp.Error)
  233. }
  234. return nil, fmt.Errorf("invalid response: code %d: body: %s", resp.StatusCode, string(body))
  235. }
  236. if result != nil {
  237. d := json.NewDecoder(resp.Body)
  238. d.UseNumber()
  239. err := d.Decode(result)
  240. if err != nil {
  241. return nil, errors.Wrap(err, "failed to decode JSON")
  242. }
  243. }
  244. return resp, nil
  245. }
  246. // Ping will check to see if the server is up with an optional timeout on waiting for leader.
  247. // Ping returns how long the request took, the version of the server it connected to, and an error if one occurred.
  248. func (c *HTTPClient) Ping(ctx context.Context) (time.Duration, string, error) {
  249. now := time.Now()
  250. u := c.url()
  251. u.Path = "ping"
  252. if ctx != nil {
  253. if dl, ok := ctx.Deadline(); ok {
  254. v := url.Values{}
  255. v.Set("wait_for_leader", fmt.Sprintf("%.0fs", time.Now().Sub(dl).Seconds()))
  256. u.RawQuery = v.Encode()
  257. }
  258. }
  259. req, err := http.NewRequest("GET", u.String(), nil)
  260. if err != nil {
  261. return 0, "", err
  262. }
  263. if ctx != nil {
  264. req = req.WithContext(ctx)
  265. }
  266. resp, err := c.do(req, nil, http.StatusNoContent)
  267. if err != nil {
  268. return 0, "", err
  269. }
  270. version := resp.Header.Get("X-Influxdb-Version")
  271. return time.Since(now), version, nil
  272. }
  273. func (c *HTTPClient) Write(bp BatchPoints) error {
  274. var b bytes.Buffer
  275. precision := bp.Precision()
  276. for _, p := range bp.Points() {
  277. if _, err := b.Write(p.Bytes(precision)); err != nil {
  278. return err
  279. }
  280. if err := b.WriteByte('\n'); err != nil {
  281. return err
  282. }
  283. }
  284. u := c.url()
  285. u.Path = "write"
  286. v := url.Values{}
  287. v.Set("db", bp.Database())
  288. v.Set("rp", bp.RetentionPolicy())
  289. v.Set("precision", bp.Precision())
  290. v.Set("consistency", bp.WriteConsistency())
  291. u.RawQuery = v.Encode()
  292. req, err := http.NewRequest("POST", u.String(), &b)
  293. if err != nil {
  294. return err
  295. }
  296. req.Header.Set("Content-Type", "application/octet-stream")
  297. _, err = c.do(req, nil, http.StatusNoContent, http.StatusOK)
  298. return err
  299. }
  300. // Response represents a list of statement results.
  301. type Response struct {
  302. Results []Result
  303. Err string `json:"error,omitempty"`
  304. }
  305. // Error returns the first error from any statement.
  306. // Returns nil if no errors occurred on any statements.
  307. func (r *Response) Error() error {
  308. if r.Err != "" {
  309. return fmt.Errorf(r.Err)
  310. }
  311. for _, result := range r.Results {
  312. if result.Err != "" {
  313. return fmt.Errorf(result.Err)
  314. }
  315. }
  316. return nil
  317. }
  318. // Message represents a user message.
  319. type Message struct {
  320. Level string
  321. Text string
  322. }
  323. // Result represents a resultset returned from a single statement.
  324. type Result struct {
  325. Series []imodels.Row
  326. Messages []*Message
  327. Err string `json:"error,omitempty"`
  328. }
  329. // Query sends a command to the server and returns the Response
  330. func (c *HTTPClient) Query(q Query) (*Response, error) {
  331. u := c.url()
  332. u.Path = "query"
  333. v := url.Values{}
  334. v.Set("q", q.Command)
  335. v.Set("db", q.Database)
  336. if q.Precision != "" {
  337. v.Set("epoch", q.Precision)
  338. }
  339. u.RawQuery = v.Encode()
  340. req, err := http.NewRequest("POST", u.String(), nil)
  341. if err != nil {
  342. return nil, err
  343. }
  344. response := &Response{}
  345. _, err = c.do(req, response, http.StatusOK)
  346. if err != nil {
  347. return nil, err
  348. }
  349. if err := response.Error(); err != nil {
  350. return nil, err
  351. }
  352. return response, nil
  353. }
  354. // BatchPoints is an interface into a batched grouping of points to write into
  355. // InfluxDB together. BatchPoints is NOT thread-safe, you must create a separate
  356. // batch for each goroutine.
  357. type BatchPoints interface {
  358. // AddPoint adds the given point to the Batch of points
  359. AddPoint(p Point)
  360. // AddPoints adds the given points to the Batch of points
  361. AddPoints(ps []Point)
  362. // Points lists the points in the Batch
  363. Points() []Point
  364. // Precision returns the currently set precision of this Batch
  365. Precision() string
  366. // SetPrecision sets the precision of this batch.
  367. SetPrecision(s string) error
  368. // Database returns the currently set database of this Batch
  369. Database() string
  370. // SetDatabase sets the database of this Batch
  371. SetDatabase(s string)
  372. // WriteConsistency returns the currently set write consistency of this Batch
  373. WriteConsistency() string
  374. // SetWriteConsistency sets the write consistency of this Batch
  375. SetWriteConsistency(s string)
  376. // RetentionPolicy returns the currently set retention policy of this Batch
  377. RetentionPolicy() string
  378. // SetRetentionPolicy sets the retention policy of this Batch
  379. SetRetentionPolicy(s string)
  380. }
  381. // NewBatchPoints returns a BatchPoints interface based on the given config.
  382. func NewBatchPoints(conf BatchPointsConfig) (BatchPoints, error) {
  383. if conf.Precision == "" {
  384. conf.Precision = "ns"
  385. }
  386. if _, err := time.ParseDuration("1" + conf.Precision); err != nil {
  387. return nil, err
  388. }
  389. bp := &batchpoints{
  390. database: conf.Database,
  391. precision: conf.Precision,
  392. retentionPolicy: conf.RetentionPolicy,
  393. writeConsistency: conf.WriteConsistency,
  394. }
  395. return bp, nil
  396. }
  397. type batchpoints struct {
  398. points []Point
  399. database string
  400. precision string
  401. retentionPolicy string
  402. writeConsistency string
  403. }
  404. func (bp *batchpoints) AddPoint(p Point) {
  405. bp.points = append(bp.points, p)
  406. }
  407. func (bp *batchpoints) AddPoints(ps []Point) {
  408. bp.points = append(bp.points, ps...)
  409. }
  410. func (bp *batchpoints) Points() []Point {
  411. return bp.points
  412. }
  413. func (bp *batchpoints) Precision() string {
  414. return bp.precision
  415. }
  416. func (bp *batchpoints) Database() string {
  417. return bp.database
  418. }
  419. func (bp *batchpoints) WriteConsistency() string {
  420. return bp.writeConsistency
  421. }
  422. func (bp *batchpoints) RetentionPolicy() string {
  423. return bp.retentionPolicy
  424. }
  425. func (bp *batchpoints) SetPrecision(p string) error {
  426. if _, err := time.ParseDuration("1" + p); err != nil {
  427. return err
  428. }
  429. bp.precision = p
  430. return nil
  431. }
  432. func (bp *batchpoints) SetDatabase(db string) {
  433. bp.database = db
  434. }
  435. func (bp *batchpoints) SetWriteConsistency(wc string) {
  436. bp.writeConsistency = wc
  437. }
  438. func (bp *batchpoints) SetRetentionPolicy(rp string) {
  439. bp.retentionPolicy = rp
  440. }
  441. type Point struct {
  442. Name string
  443. Tags map[string]string
  444. Fields map[string]interface{}
  445. Time time.Time
  446. }
  447. // Returns byte array of a line protocol representation of the point
  448. func (p Point) Bytes(precision string) []byte {
  449. key := imodels.MakeKey([]byte(p.Name), imodels.NewTags(p.Tags))
  450. fields := imodels.Fields(p.Fields).MarshalBinary()
  451. kl := len(key)
  452. fl := len(fields)
  453. var bytes []byte
  454. if p.Time.IsZero() {
  455. bytes = make([]byte, fl+kl+1)
  456. copy(bytes, key)
  457. bytes[kl] = ' '
  458. copy(bytes[kl+1:], fields)
  459. } else {
  460. timeStr := strconv.FormatInt(p.Time.UnixNano()/imodels.GetPrecisionMultiplier(precision), 10)
  461. tl := len(timeStr)
  462. bytes = make([]byte, fl+kl+tl+2)
  463. copy(bytes, key)
  464. bytes[kl] = ' '
  465. copy(bytes[kl+1:], fields)
  466. bytes[kl+fl+1] = ' '
  467. copy(bytes[kl+fl+2:], []byte(timeStr))
  468. }
  469. return bytes
  470. }
  471. // Simple type to create github.com/influxdata/kapacitor/influxdb clients.
  472. type ClientCreator struct{}
  473. func (ClientCreator) Create(config Config) (ClientUpdater, error) {
  474. return NewHTTPClient(config)
  475. }