influxdb/cmd/influx/query.go

435 lines
9.4 KiB
Go

package main
import (
"bytes"
"encoding/json"
"fmt"
"io"
"io/ioutil"
"net/http"
"net/url"
"os"
"sort"
"strconv"
"strings"
"github.com/influxdata/flux"
"github.com/influxdata/flux/csv"
"github.com/influxdata/flux/values"
ihttp "github.com/influxdata/influxdb/v2/http"
"github.com/spf13/cobra"
)
var queryFlags struct {
org organization
file string
raw bool
}
func cmdQuery(f *globalFlags, opts genericCLIOpts) *cobra.Command {
cmd := opts.newCmd("query [query literal or -f /path/to/query.flux]", fluxQueryF, true)
cmd.Short = "Execute a Flux query"
cmd.Long = `Execute a Flux query provided via the first argument or a file or stdin`
cmd.Args = cobra.MaximumNArgs(1)
f.registerFlags(cmd)
queryFlags.org.register(cmd, true)
cmd.Flags().StringVarP(&queryFlags.file, "file", "f", "", "Path to Flux query file")
cmd.Flags().BoolVarP(&queryFlags.raw, "raw", "r", false, "Display raw query results")
return cmd
}
// readFluxQuery returns first argument, file contents or stdin
func readFluxQuery(args []string, file string) (string, error) {
// backward compatibility
if len(args) > 0 {
if strings.HasPrefix(args[0], "@") {
file = args[0][1:]
args = args[:0]
} else if args[0] == "-" {
file = ""
args = args[:0]
}
}
var query string
switch {
case len(args) > 0:
query = args[0]
case len(file) > 0:
content, err := ioutil.ReadFile(file)
if err != nil {
return "", err
}
query = string(content)
default:
content, err := ioutil.ReadAll(os.Stdin)
if err != nil {
return "", err
}
query = string(content)
}
return query, nil
}
func fluxQueryF(cmd *cobra.Command, args []string) error {
if err := queryFlags.org.validOrgFlags(&flags); err != nil {
return err
}
q, err := readFluxQuery(args, queryFlags.file)
if err != nil {
return fmt.Errorf("failed to load query: %v", err)
}
u, err := url.Parse(flags.config().Host)
if err != nil {
return fmt.Errorf("unable to parse host: %s", err)
}
if !strings.HasSuffix(u.Path, "/") {
u.Path += "/"
}
u.Path += "api/v2/query"
params := url.Values{}
if queryFlags.org.id != "" {
params.Set("orgID", queryFlags.org.id)
} else {
params.Set("org", queryFlags.org.name)
}
u.RawQuery = params.Encode()
body, _ := json.Marshal(map[string]interface{}{
"query": q,
"type": "flux",
"dialect": map[string]interface{}{
"annotations": []string{"group", "datatype", "default"},
"delimiter": ",",
"header": true,
},
})
req, _ := http.NewRequest("POST", u.String(), bytes.NewReader(body))
req.Header.Set("Authorization", "Token "+flags.config().Token)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return err
}
defer func() { _ = resp.Body.Close() }()
if err := ihttp.CheckError(resp); err != nil {
return err
}
if queryFlags.raw {
io.Copy(os.Stdout, resp.Body)
return nil
}
dec := csv.NewMultiResultDecoder(csv.ResultDecoderConfig{})
results, err := dec.Decode(resp.Body)
if err != nil {
return fmt.Errorf("query decode error: %s", err)
}
defer results.Release()
for results.More() {
res := results.Next()
fmt.Println("Result:", res.Name())
if err := res.Tables().Do(func(tbl flux.Table) error {
_, err := newFormatter(tbl).WriteTo(os.Stdout)
return err
}); err != nil {
return err
}
}
// It is safe and appropriate to call Release multiple times and must be
// called before checking the error on the next line.
results.Release()
return results.Err()
}
// Below is a copy and trimmed version of the execute/format.go file from flux.
// It is copied here to avoid requiring a dependency on the execute package which
// may pull in the flux runtime as a dependency.
// In the future, the formatters and other primitives such as the csv parser should
// probably be separated out into user libraries anyway.
const fixedWidthTimeFmt = "2006-01-02T15:04:05.000000000Z"
// formatter writes a table to a Writer.
type formatter struct {
tbl flux.Table
widths []int
maxWidth int
newWidths []int
pad []byte
dash []byte
// fmtBuf is used to format values
fmtBuf [64]byte
cols orderedCols
}
var eol = []byte{'\n'}
// newFormatter creates a formatter for a given table.
func newFormatter(tbl flux.Table) *formatter {
return &formatter{
tbl: tbl,
}
}
type writeToHelper struct {
w io.Writer
n int64
err error
}
func (w *writeToHelper) write(data []byte) {
if w.err != nil {
return
}
n, err := w.w.Write(data)
w.n += int64(n)
w.err = err
}
var minWidthsByType = map[flux.ColType]int{
flux.TBool: 12,
flux.TInt: 26,
flux.TUInt: 27,
flux.TFloat: 28,
flux.TString: 22,
flux.TTime: len(fixedWidthTimeFmt),
flux.TInvalid: 10,
}
// WriteTo writes the formatted table data to w.
func (f *formatter) WriteTo(out io.Writer) (int64, error) {
w := &writeToHelper{w: out}
// Sort cols
cols := f.tbl.Cols()
f.cols = newOrderedCols(cols, f.tbl.Key())
sort.Sort(f.cols)
// Compute header widths
f.widths = make([]int, len(cols))
for j, c := range cols {
// Column header is "<label>:<type>"
l := len(c.Label) + len(c.Type.String()) + 1
min := minWidthsByType[c.Type]
if min > l {
l = min
}
if l > f.widths[j] {
f.widths[j] = l
}
if l > f.maxWidth {
f.maxWidth = l
}
}
// Write table header
w.write([]byte("Table: keys: ["))
labels := make([]string, len(f.tbl.Key().Cols()))
for i, c := range f.tbl.Key().Cols() {
labels[i] = c.Label
}
w.write([]byte(strings.Join(labels, ", ")))
w.write([]byte("]"))
w.write(eol)
// Check err and return early
if w.err != nil {
return w.n, w.err
}
// Write rows
r := 0
w.err = f.tbl.Do(func(cr flux.ColReader) error {
if r == 0 {
l := cr.Len()
for i := 0; i < l; i++ {
for oj, c := range f.cols.cols {
j := f.cols.Idx(oj)
buf := f.valueBuf(i, j, c.Type, cr)
l := len(buf)
if l > f.widths[j] {
f.widths[j] = l
}
if l > f.maxWidth {
f.maxWidth = l
}
}
}
f.makePaddingBuffers()
f.writeHeader(w)
f.writeHeaderSeparator(w)
f.newWidths = make([]int, len(f.widths))
copy(f.newWidths, f.widths)
}
l := cr.Len()
for i := 0; i < l; i++ {
for oj, c := range f.cols.cols {
j := f.cols.Idx(oj)
buf := f.valueBuf(i, j, c.Type, cr)
l := len(buf)
padding := f.widths[j] - l
if padding >= 0 {
w.write(f.pad[:padding])
w.write(buf)
} else {
//TODO make unicode friendly
w.write(buf[:f.widths[j]-3])
w.write([]byte{'.', '.', '.'})
}
w.write(f.pad[:2])
if l > f.newWidths[j] {
f.newWidths[j] = l
}
if l > f.maxWidth {
f.maxWidth = l
}
}
w.write(eol)
r++
}
return w.err
})
return w.n, w.err
}
func (f *formatter) makePaddingBuffers() {
if len(f.pad) != f.maxWidth {
f.pad = make([]byte, f.maxWidth)
for i := range f.pad {
f.pad[i] = ' '
}
}
if len(f.dash) != f.maxWidth {
f.dash = make([]byte, f.maxWidth)
for i := range f.dash {
f.dash[i] = '-'
}
}
}
func (f *formatter) writeHeader(w *writeToHelper) {
for oj, c := range f.cols.cols {
j := f.cols.Idx(oj)
buf := append(append([]byte(c.Label), ':'), []byte(c.Type.String())...)
w.write(f.pad[:f.widths[j]-len(buf)])
w.write(buf)
w.write(f.pad[:2])
}
w.write(eol)
}
func (f *formatter) writeHeaderSeparator(w *writeToHelper) {
for oj := range f.cols.cols {
j := f.cols.Idx(oj)
w.write(f.dash[:f.widths[j]])
w.write(f.pad[:2])
}
w.write(eol)
}
func (f *formatter) valueBuf(i, j int, typ flux.ColType, cr flux.ColReader) []byte {
buf := []byte("")
switch typ {
case flux.TBool:
if cr.Bools(j).IsValid(i) {
buf = strconv.AppendBool(f.fmtBuf[0:0], cr.Bools(j).Value(i))
}
case flux.TInt:
if cr.Ints(j).IsValid(i) {
buf = strconv.AppendInt(f.fmtBuf[0:0], cr.Ints(j).Value(i), 10)
}
case flux.TUInt:
if cr.UInts(j).IsValid(i) {
buf = strconv.AppendUint(f.fmtBuf[0:0], cr.UInts(j).Value(i), 10)
}
case flux.TFloat:
if cr.Floats(j).IsValid(i) {
// TODO allow specifying format and precision
buf = strconv.AppendFloat(f.fmtBuf[0:0], cr.Floats(j).Value(i), 'f', -1, 64)
}
case flux.TString:
if cr.Strings(j).IsValid(i) {
buf = []byte(cr.Strings(j).ValueString(i))
}
case flux.TTime:
if cr.Times(j).IsValid(i) {
buf = []byte(values.Time(cr.Times(j).Value(i)).String())
}
}
return buf
}
// orderedCols sorts a list of columns:
//
// * time
// * common tags sorted by label
// * other tags sorted by label
// * value
//
type orderedCols struct {
indexMap []int
cols []flux.ColMeta
key flux.GroupKey
}
func newOrderedCols(cols []flux.ColMeta, key flux.GroupKey) orderedCols {
indexMap := make([]int, len(cols))
for i := range indexMap {
indexMap[i] = i
}
cpy := make([]flux.ColMeta, len(cols))
copy(cpy, cols)
return orderedCols{
indexMap: indexMap,
cols: cpy,
key: key,
}
}
func (o orderedCols) Idx(oj int) int {
return o.indexMap[oj]
}
func (o orderedCols) Len() int { return len(o.cols) }
func (o orderedCols) Swap(i int, j int) {
o.cols[i], o.cols[j] = o.cols[j], o.cols[i]
o.indexMap[i], o.indexMap[j] = o.indexMap[j], o.indexMap[i]
}
func (o orderedCols) Less(i int, j int) bool {
ki := colIdx(o.cols[i].Label, o.key.Cols())
kj := colIdx(o.cols[j].Label, o.key.Cols())
if ki >= 0 && kj >= 0 {
return ki < kj
} else if ki >= 0 {
return true
} else if kj >= 0 {
return false
}
return i < j
}
func colIdx(label string, cols []flux.ColMeta) int {
for j, c := range cols {
if c.Label == label {
return j
}
}
return -1
}