storage/client.go

// Copyright 2017 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package storage contains a client for the performance data storage server.
package storage

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"golang.org/x/perf/storage/benchfmt"
)

// A Client issues queries to a performance data storage server.
// It is safe to use from multiple goroutines simultaneously.
type Client struct {
	// BaseURL is the base URL of the storage server.
	BaseURL string
	// HTTPClient is the HTTP client for sending requests. If nil, http.DefaultClient will be used.
	HTTPClient *http.Client
}

// httpClient returns the http.Client to use for requests.
func (c *Client) httpClient() *http.Client {
	if c.HTTPClient != nil {
		return c.HTTPClient
	}
	return http.DefaultClient
}

// Query searches for results matching the given query string.
//
// The query string is first parsed into quoted words (as in the shell)
// and then each word must be formatted as one of the following:
// key:value - exact match on label "key" = "value"
// key>value - value greater than (useful for dates)
// key<value - value less than (also useful for dates)
func (c *Client) Query(q string) *Query {
	hc := c.httpClient()

	resp, err := hc.Get(c.BaseURL + "/search?" + url.Values{"q": []string{q}}.Encode())
	if err != nil {
		return &Query{err: err}
	}

	if resp.StatusCode != 200 {
		defer resp.Body.Close()
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return &Query{err: err}
		}
		return &Query{err: fmt.Errorf("%s", body)}
	}

	br := benchfmt.NewReader(resp.Body)

	return &Query{br: br, body: resp.Body}
}

// A Query allows iteration over the results of a search query.
// Use Next to advance through the results, making sure to call Close when done:
//
//   q := client.Query("key:value")
//   defer q.Close()
//   for q.Next() {
//     res := q.Result()
//     ...
//   }
//   if err = q.Err(); err != nil {
//     // handle error encountered during query
//   }
type Query struct {
	br   *benchfmt.Reader
	body io.ReadCloser
	err  error
}

// Next prepares the next result for reading with the Result
// method. It returns false when there are no more results, either by
// reaching the end of the input or an error.
func (q *Query) Next() bool {
	if q.err != nil {
		return false
	}
	return q.br.Next()
}

// Result returns the most recent result generated by a call to Next.
func (q *Query) Result() *benchfmt.Result {
	return q.br.Result()
}

// Err returns the first error encountered during the query.
func (q *Query) Err() error {
	if q.err != nil {
		return q.err
	}
	return q.br.Err()
}

// Close frees resources associated with the query.
func (q *Query) Close() error {
	if q.body != nil {
		q.body.Close()
		q.body = nil
	}
	return q.err
}

// UploadInfo represents an upload summary.
type UploadInfo struct {
	Count       int
	UploadID    string
	LabelValues benchfmt.Labels `json:",omitempty"`
}

// ListUploads searches for uploads containing results matching the given query string.
// The query may be empty, in which case all uploads will be returned.
// extraLabels specifies other labels to be retrieved.
// If limit is 0, no limit will be provided to the server.
// The uploads are returned starting with the most recent upload.
func (c *Client) ListUploads(q string, extraLabels []string, limit int) *UploadList {
	hc := c.httpClient()

	v := url.Values{"extra_label": extraLabels}
	if q != "" {
		v["q"] = []string{q}
	}
	if limit != 0 {
		v["limit"] = []string{fmt.Sprintf("%d", limit)}
	}

	u := c.BaseURL + "/uploads"
	if len(v) > 0 {
		u += "?" + v.Encode()
	}
	resp, err := hc.Get(u)
	if err != nil {
		return &UploadList{err: err}
	}
	if resp.StatusCode != 200 {
		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return &UploadList{err: err}
		}
		return &UploadList{err: fmt.Errorf("%s", body)}
	}
	return &UploadList{body: resp.Body, dec: json.NewDecoder(resp.Body)}
}

// UploadList is the result of ListUploads.
// Use Next to advance through the rows, making sure to call Close when done:
//
//   q := db.ListUploads("key:value")
//   defer q.Close()
//   for q.Next() {
//     id, count := q.Row()
//     labels := q.LabelValues()
//     ...
//   }
//   err = q.Err() // get any error encountered during iteration
//   ...
type UploadList struct {
	body io.Closer
	dec  *json.Decoder
	// from last call to Next
	ui  UploadInfo
	err error
}

// Next prepares the next result for reading with the Result
// method. It returns false when there are no more results, either by
// reaching the end of the input or an error.
func (ul *UploadList) Next() bool {
	if ul.err != nil {
		return false
	}

	// Clear UploadInfo before decoding new value.
	ul.ui = UploadInfo{}

	ul.err = ul.dec.Decode(&ul.ui)
	return ul.err == nil
}

// Info returns the most recent UploadInfo generated by a call to Next.
func (ul *UploadList) Info() UploadInfo {
	return ul.ui
}

// Err returns the error state of the query.
func (ul *UploadList) Err() error {
	if ul.err == io.EOF {
		return nil
	}
	return ul.err
}

// Close frees resources associated with the query.
func (ul *UploadList) Close() error {
	if ul.body != nil {
		return ul.body.Close()
		ul.body = nil
	}
	return ul.Err()
}

// TODO(quentin): Move upload code here from cmd/benchsave?