[query] Add pickled format for graphite render endpoint (#1446)

src/query/api/v1/handler/graphite/find_parser.go

@@ -76,7 +76,12 @@ func parseFindParamsToQuery(r *http.Request) (
 		return nil, xhttp.NewParseError(errors.ErrNoQueryFound, http.StatusBadRequest)
 	}
 
-	matchers := graphiteStorage.TranslateQueryToMatchers(query)
+	matchers, err := graphiteStorage.TranslateQueryToMatchers(query)
+	if err != nil {
+		return nil, xhttp.NewParseError(fmt.Errorf("invalid 'query': %s", query),
+			http.StatusBadRequest)
+	}
+
 	return &storage.FetchQuery{
 		Raw:         query,
 		TagMatchers: matchers,

src/query/api/v1/handler/graphite/pickle/opcodes.go

@@ -0,0 +1,36 @@
+// Copyright (c) 2019 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package pickle
+
+// list of opcodes required for pickling graphite query results.
+const (
+	opNone       = 0x4e
+	opMark       = 0x28
+	opStop       = 0x2e
+	opBinInt     = 0x4a
+	opBinUnicode = 0x58
+	opBinFloat   = 0x47
+	opEmptyList  = 0x5d
+	opAppends    = 0x65
+	opEmptyDict  = 0x7d
+	opSetItems   = 0x75
+	opProto      = 0x80
+)

src/query/api/v1/handler/graphite/pickle/pickle_writer.go

@@ -0,0 +1,175 @@
+// Copyright (c) 2019 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package pickle
+
+import (
+	"bufio"
+	"encoding/binary"
+	"io"
+	"math"
+)
+
+var (
+	programStart = []uint8{opProto, 0x2}
+	programEnd   = []uint8{opStop}
+	listStart    = []uint8{opEmptyList, opMark}
+	dictStart    = []uint8{opEmptyDict, opMark}
+)
+
+// A Writer is capable writing out the opcodes required by the pickle format.
+// Note that this is a very limited implementation of pickling; just enough for
+// us to implement the opcodes required by graphite /render.
+type Writer struct {
+	w   *bufio.Writer
+	buf [8]byte
+	err error
+}
+
+// NewWriter creates a new pickle writer.
+func NewWriter(w io.Writer) *Writer {
+	pw := &Writer{
+		w: bufio.NewWriter(w),
+	}
+
+	_, pw.err = pw.w.Write(programStart)
+	return pw
+}
+
+// BeginDict starts marshalling a python dict.
+func (p *Writer) BeginDict() {
+	if p.err != nil {
+		return
+	}
+
+	if _, p.err = p.w.Write(dictStart); p.err != nil {
+		return
+	}
+}
+
+// WriteDictKey writes a dictionary key.
+func (p *Writer) WriteDictKey(s string) {
+	p.WriteString(s)
+}
+
+// EndDict ends marshalling a python dict.
+func (p *Writer) EndDict() {
+	if p.err != nil {
+		return
+	}
+
+	p.err = p.w.WriteByte(opSetItems)
+}
+
+// BeginList begins writing a new python list.
+func (p *Writer) BeginList() {
+	if p.err != nil {
+		return
+	}
+
+	_, p.err = p.w.Write(listStart)
+}
+
+// EndList ends writing a python list.
+func (p *Writer) EndList() {
+	if p.err != nil {
+		return
+	}
+
+	p.w.WriteByte(opAppends)
+}
+
+// WriteNone writes a python `None`.
+func (p *Writer) WriteNone() {
+	if p.err != nil {
+		return
+	}
+
+	p.err = p.w.WriteByte(opNone)
+}
+
+// WriteFloat64 writes a float64 value. NaNs are converted in `None`.
+func (p *Writer) WriteFloat64(v float64) {
+	if math.IsNaN(v) {
+		p.WriteNone()
+		return
+	}
+
+	if p.err != nil {
+		return
+	}
+
+	if p.err = p.w.WriteByte(opBinFloat); p.err != nil {
+		return
+	}
+
+	binary.BigEndian.PutUint64(p.buf[:], math.Float64bits(v))
+	_, p.err = p.w.Write(p.buf[:])
+}
+
+// WriteString writes a python string.
+func (p *Writer) WriteString(s string) {
+	if p.err != nil {
+		return
+	}
+
+	if p.err = p.w.WriteByte(opBinUnicode); p.err != nil {
+		return
+	}
+
+	binary.LittleEndian.PutUint32(p.buf[:4], uint32(len(s)))
+	if _, p.err = p.w.Write(p.buf[:4]); p.err != nil {
+		return
+	}
+
+	_, p.err = p.w.WriteString(s)
+}
+
+// WriteInt writes an int value.
+func (p *Writer) WriteInt(n int) {
+	if p.err != nil {
+		return
+	}
+
+	if p.err = p.w.WriteByte(opBinInt); p.err != nil {
+		return
+	}
+
+	binary.LittleEndian.PutUint32(p.buf[:4], uint32(n))
+	_, p.err = p.w.Write(p.buf[:4])
+}
+
+// Close closes the writer, marking the end of the stream and flushing any
+// pending values.
+func (p *Writer) Close() error {
+	if p.err != nil {
+		return p.err
+	}
+
+	if _, p.err = p.w.Write(programEnd); p.err != nil {
+		return p.err
+	}
+
+	if p.err = p.w.Flush(); p.err != nil {
+		return p.err
+	}
+
+	return nil
+}

src/query/api/v1/handler/graphite/pickle/pickle_writer_test.go

@@ -0,0 +1,106 @@
+// Copyright (c) 2019 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package pickle
+
+import (
+	"bytes"
+	"math"
+	"testing"
+
+	"github.com/hydrogen18/stalecucumber"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWriteEmptyDict(t *testing.T) {
+	var buf bytes.Buffer
+	w := NewWriter(&buf)
+	w.BeginDict()
+	w.EndDict()
+	require.NoError(t, w.Close())
+
+	var m map[interface{}]interface{}
+	require.NoError(t, unpickle(buf.Bytes(), &m))
+	assert.Equal(t, map[interface{}]interface{}{}, m)
+}
+
+func TestWriteEmptyList(t *testing.T) {
+	var buf bytes.Buffer
+	w := NewWriter(&buf)
+	w.BeginList()
+	w.EndList()
+	require.NoError(t, w.Close())
+
+	var m []string
+	require.NoError(t, unpickle(buf.Bytes(), &m))
+	assert.Equal(t, []string{}, m)
+}
+
+func TestWriteComplex(t *testing.T) {
+	var buf bytes.Buffer
+	w := NewWriter(&buf)
+	w.BeginDict()
+	w.WriteDictKey("step")
+	w.WriteInt(3494945)
+	w.WriteDictKey("pi")
+	w.WriteFloat64(3.45E10)
+	w.WriteDictKey("none")
+	w.WriteNone()
+	w.WriteDictKey("noNumber")
+	w.WriteFloat64(math.NaN())
+	w.WriteDictKey("skey")
+	w.WriteString("hello world")
+	w.WriteDictKey("nested")
+	w.BeginDict()
+	w.WriteDictKey("fooBar")
+	w.BeginList()
+	w.WriteFloat64(349439.3494)
+	w.WriteInt(-9459450)
+	w.WriteString("A Nested String")
+	w.EndList()
+	w.EndDict()
+	w.EndDict()
+	require.NoError(t, w.Close())
+
+	s := struct {
+		Step     int
+		Pi       float64
+		NoNumber *float64
+		Skey     string
+		Nested   struct {
+			FooBar []interface{}
+		}
+	}{}
+
+	require.NoError(t, unpickle(buf.Bytes(), &s))
+	assert.Equal(t, 3494945, s.Step)
+	assert.Equal(t, 3.45E10, s.Pi)
+	assert.Nil(t, s.NoNumber)
+	assert.Equal(t, "hello world", s.Skey)
+	assert.Equal(t, []interface{}{
+		349439.3494, int64(-9459450), "A Nested String",
+	}, s.Nested.FooBar)
+}
+
+func unpickle(b []byte, data interface{}) error {
+	r := bytes.NewReader(b)
+	return stalecucumber.UnpackInto(data).From(stalecucumber.Unpickle(r))
+}

src/query/api/v1/handler/graphite/render.go

@@ -194,6 +194,6 @@ func (h *renderHandler) serveHTTP(
 		SortApplied: true,
 	}
 
-	err = WriteRenderResponse(w, response)
+	err = WriteRenderResponse(w, response, p.Format)
 	return respError{err: err, code: http.StatusOK}
 }