refactor(query): make queryd present ProxyQueryService (#12360)

Fixes influxdata/idpe#2014.
pull/12422/head
Christopher M. Wolff 2019-03-07 07:32:13 -08:00 committed by GitHub
parent 27970f8ee9
commit e28ecdc0e9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 236 additions and 352 deletions

View File

@ -16,7 +16,6 @@ type Service struct {
*OrganizationService
*UserService
*BucketService
*QueryService
*VariableService
*DashboardService
}
@ -43,10 +42,6 @@ func NewService(addr, token string) *Service {
Addr: addr,
Token: token,
},
QueryService: &QueryService{
Addr: addr,
Token: token,
},
DashboardService: &DashboardService{
Addr: addr,
Token: token,

View File

@ -25,7 +25,7 @@ type SourceProxyQueryService struct {
platform.V1SourceFields
}
func (s *SourceProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
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.influxQuery(ctx, w, req)
@ -33,10 +33,10 @@ func (s *SourceProxyQueryService) Query(ctx context.Context, w io.Writer, req *q
return s.fluxQuery(ctx, w, req)
}
return 0, fmt.Errorf("compiler type not supported")
return flux.Statistics{}, fmt.Errorf("compiler type not supported")
}
func (s *SourceProxyQueryService) fluxQuery(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
func (s *SourceProxyQueryService) fluxQuery(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
request := struct {
Spec *flux.Spec `json:"spec"`
Query string `json:"query"`
@ -52,7 +52,7 @@ func (s *SourceProxyQueryService) fluxQuery(ctx context.Context, w io.Writer, re
request.Spec = c.Spec
request.Type = lang.SpecCompilerType
default:
return 0, fmt.Errorf("compiler type not supported: %s", c.CompilerType())
return flux.Statistics{}, fmt.Errorf("compiler type not supported: %s", c.CompilerType())
}
request.Dialect = req.Dialect
@ -68,7 +68,7 @@ func (s *SourceProxyQueryService) fluxQuery(ctx context.Context, w io.Writer, re
u, err := newURL(s.URL, "/api/v2/query")
if err != nil {
return 0, err
return flux.Statistics{}, err
}
qp := u.Query()
@ -77,12 +77,12 @@ func (s *SourceProxyQueryService) fluxQuery(ctx context.Context, w io.Writer, re
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(request); err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq, err := http.NewRequest("POST", u.String(), &body)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq.Header.Set("Authorization", s.Token)
hreq.Header.Set("Content-Type", "application/json")
@ -91,28 +91,32 @@ func (s *SourceProxyQueryService) fluxQuery(ctx context.Context, w io.Writer, re
hc := newTraceClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
defer resp.Body.Close()
if err := platformhttp.CheckError(resp); err != nil {
return 0, err
return flux.Statistics{}, err
}
return io.Copy(w, resp.Body)
if _, err = io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, err
}
return flux.Statistics{}, nil
}
func (s *SourceProxyQueryService) influxQuery(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
func (s *SourceProxyQueryService) influxQuery(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
if len(s.URL) == 0 {
return 0, fmt.Errorf("URL from source cannot be empty if the compiler type is influxql")
return flux.Statistics{}, fmt.Errorf("URL from source cannot be empty if the compiler type is influxql")
}
u, err := newURL(s.URL, "/query")
if err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq, err := http.NewRequest("POST", u.String(), nil)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
// TODO(fntlnz): configure authentication methods username/password and stuff
@ -121,7 +125,7 @@ func (s *SourceProxyQueryService) influxQuery(ctx context.Context, w io.Writer,
params := hreq.URL.Query()
compiler, ok := req.Request.Compiler.(*influxql.Compiler)
if !ok {
return 0, fmt.Errorf("passed compiler is not of type 'influxql'")
return flux.Statistics{}, fmt.Errorf("passed compiler is not of type 'influxql'")
}
params.Set("q", compiler.Query)
params.Set("db", compiler.DB)
@ -132,22 +136,25 @@ func (s *SourceProxyQueryService) influxQuery(ctx context.Context, w io.Writer,
hc := newTraceClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
defer resp.Body.Close()
if err := platformhttp.CheckError(resp); err != nil {
return 0, err
return flux.Statistics{}, err
}
res := &influxql.Response{}
if err := json.NewDecoder(resp.Body).Decode(res); err != nil {
return 0, err
return flux.Statistics{}, err
}
csvDialect, ok := req.Dialect.(csv.Dialect)
if !ok {
return 0, fmt.Errorf("unsupported dialect %T", req.Dialect)
return flux.Statistics{}, fmt.Errorf("unsupported dialect %T", req.Dialect)
}
return csv.NewMultiResultEncoder(csvDialect.ResultEncoderConfig).Encode(w, influxql.NewResponseIterator(res))
if _, err = csv.NewMultiResultEncoder(csvDialect.ResultEncoderConfig).Encode(w, influxql.NewResponseIterator(res)); err != nil {
return flux.Statistics{}, err
}
return flux.Statistics{}, nil
}

View File

@ -9,6 +9,8 @@ import (
"net/http"
"github.com/influxdata/flux"
"github.com/influxdata/flux/iocounter"
icontext "github.com/influxdata/influxdb/context"
"github.com/influxdata/influxdb/query"
"github.com/julienschmidt/httprouter"
"github.com/prometheus/client_golang/prometheus"
@ -16,7 +18,8 @@ import (
)
const (
proxyQueryPath = "/api/v2/queryproxysvc"
proxyQueryPath = "/queryproxysvc"
QueryStatsTrailer = "Influx-Query-Statistics"
)
type ProxyQueryHandler struct {
@ -24,6 +27,8 @@ type ProxyQueryHandler struct {
Logger *zap.Logger
Name string
ProxyQueryService query.ProxyQueryService
CompilerMappings flux.CompilerMappings
@ -31,15 +36,23 @@ type ProxyQueryHandler struct {
}
// NewProxyQueryHandler returns a new instance of ProxyQueryHandler.
func NewProxyQueryHandler() *ProxyQueryHandler {
func NewProxyQueryHandler(name string) *ProxyQueryHandler {
h := &ProxyQueryHandler{
Router: NewRouter(),
Name: name,
}
h.HandlerFunc("GET", "/ping", h.handlePing)
h.HandlerFunc("POST", proxyQueryPath, h.handlePostQuery)
return h
}
// handlePing returns a simple response to let the client know the server is running.
func (h *ProxyQueryHandler) handlePing(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusNoContent)
}
// HTTPDialect is an encoding dialect that can write metadata to HTTP headers
type HTTPDialect interface {
SetHeaders(w http.ResponseWriter)
@ -57,6 +70,8 @@ func (h *ProxyQueryHandler) handlePostQuery(w http.ResponseWriter, r *http.Reque
return
}
w.Header().Set("Trailer", QueryStatsTrailer)
hd, ok := req.Dialect.(HTTPDialect)
if !ok {
EncodeError(ctx, fmt.Errorf("unsupported dialect over HTTP %T", req.Dialect), w)
@ -64,18 +79,29 @@ func (h *ProxyQueryHandler) handlePostQuery(w http.ResponseWriter, r *http.Reque
}
hd.SetHeaders(w)
n, err := h.ProxyQueryService.Query(ctx, w, &req)
cw := iocounter.Writer{Writer: w}
stats, err := h.ProxyQueryService.Query(ctx, &cw, &req)
if err != nil {
if n == 0 {
if cw.Count() == 0 {
// Only record the error headers IFF nothing has been written to w.
EncodeError(ctx, err, w)
return
}
h.Logger.Info("Error writing response to client",
zap.String("handler", "transpilerde"),
zap.String("handler", h.Name),
zap.Error(err),
)
return
}
// Write statistics trailer
data, err := json.Marshal(stats)
if err != nil {
h.Logger.Info("Failed to encode statistics", zap.Error(err))
return
}
w.Header().Set(QueryStatsTrailer, string(data))
}
// PrometheusCollectors satisifies the prom.PrometheusCollector interface.
@ -90,19 +116,16 @@ type ProxyQueryService struct {
InsecureSkipVerify bool
}
func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
u, err := newURL(s.Addr, proxyQueryPath)
// Ping checks to see if the server is responding to a ping request.
func (s *ProxyQueryService) Ping(ctx context.Context) error {
u, err := newURL(s.Addr, "/ping")
if err != nil {
return 0, err
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(req); err != nil {
return 0, err
return err
}
hreq, err := http.NewRequest("POST", u.String(), &body)
hreq, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return 0, err
return err
}
SetToken(s.Token, hreq)
hreq = hreq.WithContext(ctx)
@ -110,11 +133,57 @@ func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.P
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return 0, err
return err
}
return CheckError(resp)
}
func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
u, err := newURL(s.Addr, proxyQueryPath)
if err != nil {
return flux.Statistics{}, err
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(req); err != nil {
return flux.Statistics{}, err
}
hreq, err := http.NewRequest("POST", u.String(), &body)
if err != nil {
return flux.Statistics{}, err
}
token := s.Token
if token == "" {
token, err = icontext.GetToken(ctx)
if err != nil {
return flux.Statistics{}, err
}
}
SetToken(token, hreq)
hreq = hreq.WithContext(ctx)
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return flux.Statistics{}, err
}
defer resp.Body.Close()
if err := CheckError(resp); err != nil {
return 0, err
return flux.Statistics{}, err
}
return io.Copy(w, resp.Body)
if _, err = io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, err
}
data := []byte(resp.Trailer.Get(QueryStatsTrailer))
var stats flux.Statistics
if err := json.Unmarshal(data, &stats); err != nil {
return stats, err
}
return stats, nil
}

View File

@ -15,6 +15,7 @@ import (
"github.com/influxdata/flux/ast"
"github.com/influxdata/flux/complete"
"github.com/influxdata/flux/csv"
"github.com/influxdata/flux/iocounter"
"github.com/influxdata/flux/parser"
platform "github.com/influxdata/influxdb"
pcontext "github.com/influxdata/influxdb/context"
@ -100,9 +101,9 @@ func (h *FluxHandler) handleQuery(w http.ResponseWriter, r *http.Request) {
}
hd.SetHeaders(w)
n, err := h.ProxyQueryService.Query(ctx, w, req)
if err != nil {
if n == 0 {
cw := iocounter.Writer{Writer: w}
if _, err := h.ProxyQueryService.Query(ctx, &cw, req); err != nil {
if cw.Count() == 0 {
// Only record the error headers IFF nothing has been written to w.
EncodeError(ctx, err, w)
return
@ -307,24 +308,24 @@ type FluxService struct {
// Query runs a flux query against a influx server and sends the results to the io.Writer.
// Will use the token from the context over the token within the service struct.
func (s *FluxService) Query(ctx context.Context, w io.Writer, r *query.ProxyRequest) (int64, error) {
func (s *FluxService) Query(ctx context.Context, w io.Writer, r *query.ProxyRequest) (flux.Statistics, error) {
u, err := newURL(s.Addr, fluxPath)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
qreq, err := QueryRequestFromProxyRequest(r)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(qreq); err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq, err := http.NewRequest("POST", u.String(), &body)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
SetToken(s.Token, hreq)
@ -336,14 +337,18 @@ func (s *FluxService) Query(ctx context.Context, w io.Writer, r *query.ProxyRequ
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
defer resp.Body.Close()
if err := CheckError(resp); err != nil {
return 0, err
return flux.Statistics{}, err
}
return io.Copy(w, resp.Body)
if _, err := io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, err
}
return flux.Statistics{}, nil
}
var _ query.QueryService = (*FluxQueryService)(nil)

View File

@ -11,6 +11,8 @@ import (
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/influxdata/flux"
"github.com/influxdata/flux/csv"
"github.com/influxdata/flux/lang"
platform "github.com/influxdata/influxdb"
@ -24,7 +26,7 @@ func TestFluxService_Query(t *testing.T) {
ctx context.Context
r *query.ProxyRequest
status int
want int64
want flux.Statistics
wantW string
wantErr bool
}{
@ -41,7 +43,7 @@ func TestFluxService_Query(t *testing.T) {
Dialect: csv.DefaultDialect(),
},
status: http.StatusOK,
want: 6,
want: flux.Statistics{},
wantW: "howdy\n",
},
{
@ -78,8 +80,8 @@ func TestFluxService_Query(t *testing.T) {
t.Errorf("FluxService.Query() error = %v, wantErr %v", err, tt.wantErr)
return
}
if got != tt.want {
t.Errorf("FluxService.Query() = %v, want %v", got, tt.want)
if diff := cmp.Diff(tt.want, got); diff != "" {
t.Errorf("FluxService.Query() = -want/+got: %v", diff)
}
if gotW := w.String(); gotW != tt.wantW {
t.Errorf("FluxService.Query() = %v, want %v", gotW, tt.wantW)

View File

@ -1,242 +0,0 @@
package http
import (
"bytes"
"context"
"encoding/json"
"net/http"
"github.com/influxdata/flux"
"github.com/influxdata/flux/csv"
icontext "github.com/influxdata/influxdb/context"
"github.com/influxdata/influxdb/query"
"github.com/julienschmidt/httprouter"
"github.com/prometheus/client_golang/prometheus"
"go.uber.org/zap"
)
const (
queryPath = "/api/v2/querysvc"
queryStatisticsTrailer = "Influx-Query-Statistics"
)
// QueryHandler is a low-level query service handler.
type QueryHandler struct {
*httprouter.Router
Logger *zap.Logger
csvDialect csv.Dialect
QueryService query.QueryService
CompilerMappings flux.CompilerMappings
}
// NewQueryHandler returns a new instance of QueryHandler.
func NewQueryHandler() *QueryHandler {
h := &QueryHandler{
Router: NewRouter(),
csvDialect: csv.Dialect{
ResultEncoderConfig: csv.DefaultEncoderConfig(),
},
}
h.HandlerFunc("GET", "/ping", h.handlePing)
h.HandlerFunc("POST", queryPath, h.handlePostQuery)
return h
}
// handlePing returns a simple response to let the client know the server is running.
func (h *QueryHandler) handlePing(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusNoContent)
}
// handlePostQuery is the HTTP handler for the POST /api/v2/query route.
func (h *QueryHandler) handlePostQuery(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
var req query.Request
req.WithCompilerMappings(h.CompilerMappings)
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
EncodeError(ctx, err, w)
return
}
results, err := h.QueryService.Query(ctx, &req)
if err != nil {
EncodeError(ctx, err, w)
return
}
// Always cancel the results to free resources.
// If all results were consumed cancelling does nothing.
defer results.Release()
// Setup headers
w.Header().Set("Trailer", queryStatisticsTrailer)
// NOTE: We do not write out the headers here.
// It is possible that if the encoding step fails
// that we can write an error header so long as
// the encoder did not write anything.
// As such we rely on the http.ResponseWriter behavior
// to write an StatusOK header with the first write.
switch r.Header.Get("Accept") {
case "text/csv":
fallthrough
default:
h.csvDialect.SetHeaders(w)
encoder := h.csvDialect.Encoder()
n, err := encoder.Encode(w, results)
if err != nil {
if n == 0 {
// If the encoder did not write anything, we can write an error header.
EncodeError(ctx, err, w)
} else {
h.Logger.Info("Failed to encode client response",
zap.Error(err),
)
}
}
}
data, err := json.Marshal(results.Statistics())
if err != nil {
h.Logger.Info("Failed to encode statistics", zap.Error(err))
return
}
// Write statisitcs trailer
w.Header().Set(queryStatisticsTrailer, string(data))
}
// PrometheusCollectors satisifies the prom.PrometheusCollector interface.
func (h *QueryHandler) PrometheusCollectors() []prometheus.Collector {
// TODO: gather and return relevant metrics.
return nil
}
// QueryService is a basic client to interact with the QueryHandler.
type QueryService struct {
Addr string
Token string
InsecureSkipVerify bool
}
// Ping checks to see if the server is responding to a ping request.
func (s *QueryService) Ping(ctx context.Context) error {
u, err := newURL(s.Addr, "/ping")
if err != nil {
return err
}
hreq, err := http.NewRequest("GET", u.String(), nil)
if err != nil {
return err
}
hreq = hreq.WithContext(ctx)
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return err
}
return CheckError(resp)
}
// Query calls the query route with the requested query and returns a result iterator.
func (s *QueryService) Query(ctx context.Context, req *query.Request) (flux.ResultIterator, error) {
u, err := newURL(s.Addr, queryPath)
if err != nil {
return nil, err
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(req); err != nil {
return nil, err
}
hreq, err := http.NewRequest("POST", u.String(), &body)
if err != nil {
return nil, err
}
token := s.Token
if token == "" {
token, err = icontext.GetToken(ctx)
if err != nil {
return nil, err
}
}
SetToken(token, hreq)
hreq = hreq.WithContext(ctx)
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return nil, err
}
if err := CheckError(resp); err != nil {
return nil, err
}
var decoder flux.MultiResultDecoder
switch resp.Header.Get("Content-Type") {
case "text/csv":
fallthrough
default:
decoder = csv.NewMultiResultDecoder(csv.ResultDecoderConfig{})
}
results, err := decoder.Decode(resp.Body)
if err != nil {
return nil, err
}
statResults := &statsResultIterator{
results: results,
resp: resp,
}
return statResults, nil
}
// statsResultIterator implements flux.ResultIterator and flux.Statisticser by reading the HTTP trailers.
type statsResultIterator struct {
results flux.ResultIterator
resp *http.Response
statisitcs flux.Statistics
err error
}
func (s *statsResultIterator) More() bool {
return s.results.More()
}
func (s *statsResultIterator) Next() flux.Result {
return s.results.Next()
}
func (s *statsResultIterator) Release() {
s.results.Release()
s.readStats()
}
func (s *statsResultIterator) Err() error {
err := s.results.Err()
if err != nil {
return err
}
return s.err
}
func (s *statsResultIterator) Statistics() flux.Statistics {
return s.statisitcs
}
// readStats reads the query statisitcs off the response trailers.
func (s *statsResultIterator) readStats() {
data := s.resp.Trailer.Get(queryStatisticsTrailer)
if data != "" {
s.err = json.Unmarshal([]byte(data), &s.statisitcs)
}
}

View File

@ -10,6 +10,7 @@ import (
"net/url"
"strings"
"github.com/influxdata/flux"
"github.com/influxdata/flux/lang"
platform "github.com/influxdata/influxdb"
"github.com/influxdata/influxdb/query"
@ -22,29 +23,29 @@ type SourceProxyQueryService struct {
platform.SourceFields
}
func (s *SourceProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
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 0, fmt.Errorf("compiler type not supported")
return flux.Statistics{}, fmt.Errorf("compiler type not supported")
}
func (s *SourceProxyQueryService) queryFlux(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
func (s *SourceProxyQueryService) queryFlux(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
u, err := newURL(s.Addr, "/api/v2/query")
if err != nil {
return 0, err
return flux.Statistics{}, err
}
var body bytes.Buffer
if err := json.NewEncoder(&body).Encode(req); err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq, err := http.NewRequest("POST", u.String(), &body)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq.Header.Set("Authorization", fmt.Sprintf("Token %s", s.Token))
hreq.Header.Set("Content-Type", "application/json")
@ -53,25 +54,30 @@ func (s *SourceProxyQueryService) queryFlux(ctx context.Context, w io.Writer, re
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
defer resp.Body.Close()
if err := CheckError(resp); err != nil {
return 0, err
return flux.Statistics{}, err
}
return io.Copy(w, resp.Body)
if _, err = io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, err
}
return flux.Statistics{}, nil
}
func (s *SourceProxyQueryService) queryInfluxQL(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
func (s *SourceProxyQueryService) queryInfluxQL(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
compiler, ok := req.Request.Compiler.(*influxql.Compiler)
if !ok {
return 0, fmt.Errorf("compiler is not of type 'influxql'")
return flux.Statistics{}, fmt.Errorf("compiler is not of type 'influxql'")
}
u, err := newURL(s.Addr, "/query")
if err != nil {
return 0, err
return flux.Statistics{}, err
}
body := url.Values{}
@ -81,7 +87,7 @@ func (s *SourceProxyQueryService) queryInfluxQL(ctx context.Context, w io.Writer
body.Add("rp", compiler.RP)
hreq, err := http.NewRequest("POST", u.String(), strings.NewReader(body.Encode()))
if err != nil {
return 0, err
return flux.Statistics{}, err
}
hreq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
hreq.Header.Set("Authorization", fmt.Sprintf("Token %s", s.Token))
@ -90,12 +96,17 @@ func (s *SourceProxyQueryService) queryInfluxQL(ctx context.Context, w io.Writer
hc := newClient(u.Scheme, s.InsecureSkipVerify)
resp, err := hc.Do(hreq)
if err != nil {
return 0, err
return flux.Statistics{}, err
}
defer resp.Body.Close()
if err := CheckError(resp); err != nil {
return 0, err
return flux.Statistics{}, err
}
return io.Copy(w, resp.Body)
if _, err = io.Copy(w, resp.Body); err != nil {
return flux.Statistics{}, err
}
return flux.Statistics{}, nil
}

View File

@ -4,6 +4,7 @@ import (
"context"
"io"
"github.com/influxdata/flux"
"github.com/influxdata/influxdb/query"
)
@ -11,17 +12,19 @@ var _ query.ProxyQueryService = (*ProxyQueryService)(nil)
// ProxyQueryService is a mock implementation of a query.ProxyQueryService.
type ProxyQueryService struct {
QueryFn func(context.Context, io.Writer, *query.ProxyRequest) (int64, error)
QueryFn func(context.Context, io.Writer, *query.ProxyRequest) (flux.Statistics, error)
}
// NewProxyQueryService returns a mock of ProxyQueryService where its methods will return zero values.
func NewProxyQueryService() *ProxyQueryService {
return &ProxyQueryService{
QueryFn: func(context.Context, io.Writer, *query.ProxyRequest) (int64, error) { return 0, nil },
QueryFn: func(context.Context, io.Writer, *query.ProxyRequest) (flux.Statistics, error) {
return flux.Statistics{}, nil
},
}
}
// Query performs the requested query and encodes the results into w.
func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
return s.QueryFn(ctx, w, req)
}

View File

@ -2,11 +2,10 @@ package query
import (
"context"
"encoding/json"
"io"
"net/http"
"github.com/influxdata/flux"
"github.com/influxdata/flux/csv"
platform "github.com/influxdata/influxdb"
)
@ -23,35 +22,72 @@ func (b QueryServiceBridge) Query(ctx context.Context, req *Request) (flux.Resul
return flux.NewResultIteratorFromQuery(query), nil
}
// ProxyQueryServiceBridge implements ProxyQueryService while consuming a QueryService interface.
type ProxyQueryServiceBridge struct {
QueryService QueryService
// QueryServiceProxyBridge implements QueryService while consuming a ProxyQueryService interface.
type QueryServiceProxyBridge struct {
ProxyQueryService ProxyQueryService
}
func (b ProxyQueryServiceBridge) Query(ctx context.Context, w io.Writer, req *ProxyRequest) (int64, error) {
results, err := b.QueryService.Query(ctx, &req.Request)
if err != nil {
return 0, err
func (b QueryServiceProxyBridge) Query(ctx context.Context, req *Request) (flux.ResultIterator, error) {
d := csv.Dialect{ResultEncoderConfig: csv.DefaultEncoderConfig()}
preq := &ProxyRequest{
Request: *req,
Dialect: d,
}
r, w := io.Pipe()
statsChan := make(chan flux.Statistics, 1)
go func() {
stats, err := b.ProxyQueryService.Query(ctx, w, preq)
_ = w.CloseWithError(err)
statsChan <- stats
}()
dec := csv.NewMultiResultDecoder(csv.ResultDecoderConfig{})
ri, err := dec.Decode(r)
return asyncStatsResultIterator{
ResultIterator: ri,
statsChan: statsChan,
}, err
}
type asyncStatsResultIterator struct {
flux.ResultIterator
statsChan chan flux.Statistics
stats flux.Statistics
}
func (i asyncStatsResultIterator) Release() {
i.ResultIterator.Release()
i.stats = <-i.statsChan
}
func (i asyncStatsResultIterator) Statistics() flux.Statistics {
return i.stats
}
// ProxyQueryServiceAsyncBridge implements ProxyQueryService while consuming an AsyncQueryService
type ProxyQueryServiceAsyncBridge struct {
AsyncQueryService AsyncQueryService
}
func (b ProxyQueryServiceAsyncBridge) Query(ctx context.Context, w io.Writer, req *ProxyRequest) (flux.Statistics, error) {
q, err := b.AsyncQueryService.Query(ctx, &req.Request)
if err != nil {
return flux.Statistics{}, err
}
results := flux.NewResultIteratorFromQuery(q)
defer results.Release()
// Setup headers
if w, ok := w.(http.ResponseWriter); ok {
w.Header().Set("Trailer", "Influx-Query-Statistics")
}
encoder := req.Dialect.Encoder()
n, err := encoder.Encode(w, results)
_, err = encoder.Encode(w, results)
if err != nil {
return n, err
return flux.Statistics{}, err
}
if w, ok := w.(http.ResponseWriter); ok {
data, _ := json.Marshal(results.Statistics())
w.Header().Set("Influx-Query-Statistics", string(data))
}
return n, nil
stats := results.Statistics()
return stats, nil
}
// REPLQuerier implements the repl.Querier interface while consuming a QueryService

View File

@ -16,8 +16,8 @@ type LoggingServiceBridge struct {
}
// Query executes and logs the query.
func (s *LoggingServiceBridge) Query(ctx context.Context, w io.Writer, req *ProxyRequest) (n int64, err error) {
var stats flux.Statistics
func (s *LoggingServiceBridge) Query(ctx context.Context, w io.Writer, req *ProxyRequest) (stats flux.Statistics, err error) {
var n int64
defer func() {
r := recover()
if r != nil {
@ -38,7 +38,7 @@ func (s *LoggingServiceBridge) Query(ctx context.Context, w io.Writer, req *Prox
results, err := s.QueryService.Query(ctx, &req.Request)
if err != nil {
return 0, err
return stats, err
}
// Check if this result iterator reports stats. We call this defer before cancel because
// the query needs to be finished before it will have valid statistics.
@ -50,8 +50,8 @@ func (s *LoggingServiceBridge) Query(ctx context.Context, w io.Writer, req *Prox
encoder := req.Dialect.Encoder()
n, err = encoder.Encode(w, results)
if err != nil {
return n, err
return stats, err
}
// The results iterator may have had an error independent of encoding errors.
return n, results.Err()
return stats, results.Err()
}

View File

@ -10,11 +10,11 @@ import (
// ProxyQueryService mocks the idep QueryService for testing.
type ProxyQueryService struct {
QueryF func(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error)
QueryF func(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error)
}
// Query writes the results of the query request.
func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (int64, error) {
func (s *ProxyQueryService) Query(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
return s.QueryF(ctx, w, req)
}

View File

@ -26,5 +26,5 @@ type AsyncQueryService interface {
type ProxyQueryService interface {
// Query performs the requested query and encodes the results into w.
// The number of bytes written to w is returned __independent__ of any error.
Query(ctx context.Context, w io.Writer, req *ProxyRequest) (int64, error)
Query(ctx context.Context, w io.Writer, req *ProxyRequest) (flux.Statistics, error)
}

View File

@ -13,10 +13,8 @@ import (
// NewProxyQueryService returns a proxy query service based on the given queryController
// suitable for the storage read service.
func NewProxyQueryService(queryController *pcontrol.Controller) query.ProxyQueryService {
return query.ProxyQueryServiceBridge{
QueryService: query.QueryServiceBridge{
AsyncQueryService: queryController,
},
return query.ProxyQueryServiceAsyncBridge{
AsyncQueryService: queryController,
}
}