influxdb/http/source_proxy_service.go

121 lines
3.5 KiB
Go

package http
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"github.com/influxdata/flux"
"github.com/influxdata/flux/lang"
"github.com/opentracing/opentracing-go"
platform "github.com/influxdata/influxdb"
"github.com/influxdata/influxdb/kit/tracing"
"github.com/influxdata/influxdb/query"
"github.com/influxdata/influxdb/query/influxql"
)
type SourceProxyQueryService struct {
Addr string
InsecureSkipVerify bool
platform.SourceFields
}
func (s *SourceProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
switch req.Request.Compiler.CompilerType() {
case influxql.CompilerType:
return s.queryInfluxQL(ctx, w, req)
case lang.FluxCompilerType:
return s.queryFlux(ctx, w, req)
}
return flux.Statistics{}, fmt.Errorf("compiler type not supported")
}
func (s *SourceProxyQueryService) queryFlux(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
span, ctx := opentracing.StartSpanFromContext(ctx, "SourceProxyQueryService.queryFlux")
defer span.Finish()
u, err := newURL(s.Addr, "/api/v2/query")
if err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(req); err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
hreq, err := http.NewRequest("POST", u.String(), &body)
if err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
hreq.Header.Set("Authorization", fmt.Sprintf("Token %s", s.Token))
hreq.Header.Set("Content-Type", "application/json")
hreq = hreq.WithContext(ctx)
tracing.InjectToHTTPRequest(span, hreq)
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
defer resp.Body.Close()
if err := CheckError(resp); err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
if _, err = io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
return flux.Statistics{}, nil
}
func (s *SourceProxyQueryService) queryInfluxQL(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
span, ctx := opentracing.StartSpanFromContext(ctx, "SourceProxyQueryService.queryInfluxQL")
defer span.Finish()
compiler, ok := req.Request.Compiler.(*influxql.Compiler)
if !ok {
return flux.Statistics{}, tracing.LogError(span, fmt.Errorf("compiler is not of type 'influxql'"))
}
u, err := newURL(s.Addr, "/query")
if err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
body := url.Values{}
body.Add("db", compiler.DB)
body.Add("org", compiler.Cluster)
body.Add("q", compiler.Query)
body.Add("rp", compiler.RP)
hreq, err := http.NewRequest("POST", u.String(), strings.NewReader(body.Encode()))
if err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
hreq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
hreq.Header.Set("Authorization", fmt.Sprintf("Token %s", s.Token))
hreq = hreq.WithContext(ctx)
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
defer resp.Body.Close()
if err := CheckError(resp); err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
if _, err = io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, tracing.LogError(span, err)
}
return flux.Statistics{}, nil
}