add correct HTTP query handler

Author: tochka

main_test.go

@@ -102,6 +102,44 @@ func TestServe(t *testing.T) {
 			},
 			startTLS,
 		},
+		{
+			"https cache with mix query source",
+			"testdata/https.cache.yml",
+			func(t *testing.T) {
+				// do request which response must be cached
+				queryURLParam := "SELECT * FROM system.numbers"
+				queryBody := "LIMIT 10"
+				expectedQuery := queryURLParam + "\n" + queryBody
+				buf := bytes.NewBufferString(queryBody)
+				req, err := http.NewRequest("GET", "https://127.0.0.1:8443?query="+url.QueryEscape(queryURLParam), buf)
+				checkErr(t, err)
+				req.SetBasicAuth("default", "qwerty")
+				resp, err := tlsClient.Do(req)
+				checkErr(t, err)
+				if resp.StatusCode != http.StatusOK {
+					t.Fatalf("unexpected status code: %d; expected: %d", resp.StatusCode, http.StatusOK)
+				}
+				resp.Body.Close()
+
+				// check cached response
+				key := &cache.Key{
+					Query:          []byte(expectedQuery),
+					AcceptEncoding: "gzip",
+				}
+				path := fmt.Sprintf("%s/cache/%s", testDir, key.String())
+				if _, err := os.Stat(path); err != nil {
+					t.Fatalf("err while getting file %q info: %s", path, err)
+				}
+				rw := httptest.NewRecorder()
+				cc := proxy.caches["https_cache"]
+				if err := cc.WriteTo(rw, key); err != nil {
+					t.Fatalf("unexpected error while writing reposnse from cache: %s", err)
+				}
+				expected := "Ok.\n"
+				checkResponse(t, rw.Body, expected)
+			},
+			startTLS,
+		},
 		{
 			"bad https cache",
 			"testdata/https.cache.yml",

utils.go

@@ -47,8 +47,10 @@ func getAuth(req *http.Request) (string, string) {
 //
 // getQuerySnippet must be called only for error reporting.
 func getQuerySnippet(req *http.Request) string {
-	if req.Method == http.MethodGet && req.URL.Query().Get("query") != "" {
-		return req.URL.Query().Get("query")
+	var query string
+
+	if req.URL.Query().Get("query") != "" {
+		query = req.URL.Query().Get("query")
 	}
 
 	crc, ok := req.Body.(*cachedReadCloser)
@@ -66,26 +68,44 @@ func getQuerySnippet(req *http.Request) string {
 
 	u := getDecompressor(req)
 	if u == nil {
-		return data
+		if len(query) != 0 && len(data) != 0 {
+			query += "\n"
+		}
+
+		return query + data
 	}
 	bs := bytes.NewBufferString(data)
 	b, err := u.decompress(bs)
 	if err == nil {
-		return string(b)
+		if len(query) != 0 && len(b) != 0 {
+			query += "\n"
+		}
+
+		return query + string(b)
 	}
 	// It is better to return partially decompressed data instead of an empty string.
 	if len(b) > 0 {
-		return string(b)
+		if len(query) != 0 && len(b) != 0 {
+			query += "\n"
+		}
+
+		return query + string(b)
 	}
+
 	// The data failed to be decompressed. Return compressed data
 	// instead of an empty string.
-	return data
+	if len(query) != 0 && len(data) != 0 {
+		query += "\n"
+	}
+
+	return query + data
 }
 
 // getFullQuery returns full query from req.
 func getFullQuery(req *http.Request) ([]byte, error) {
-	if req.Method == http.MethodGet && req.URL.Query().Get("query") != "" {
-		return []byte(req.URL.Query().Get("query")), nil
+	var result bytes.Buffer
+	if req.URL.Query().Get("query") != "" {
+		result.WriteString(req.URL.Query().Get("query"))
 	}
 	data, err := ioutil.ReadAll(req.Body)
 	if err != nil {
@@ -95,14 +115,25 @@ func getFullQuery(req *http.Request) ([]byte, error) {
 	req.Body = ioutil.NopCloser(bytes.NewBuffer(data))
 	u := getDecompressor(req)
 	if u == nil {
-		return data, nil
+		if result.Len() != 0 && len(data) != 0 {
+			result.WriteByte('\n')
+		}
+		result.Write(data)
+
+		return result.Bytes(), nil
 	}
 	br := bytes.NewReader(data)
 	b, err := u.decompress(br)
 	if err != nil {
 		return nil, fmt.Errorf("cannot uncompress query: %s", err)
 	}
-	return b, nil
+
+	if result.Len() != 0 && len(b) != 0 {
+		result.WriteByte('\n')
+	}
+	result.Write(b)
+
+	return result.Bytes(), nil
 }
 
 // canCacheQuery returns true if q can be cached.

utils_test.go

@@ -68,7 +68,7 @@ func testCanCacheQuery(t *testing.T, q string, expected bool) {
 }
 
 func TestGetQuerySnippetGET(t *testing.T) {
-	req, err := http.NewRequest("GET", "", nil)
+	req, err := http.NewRequest("GET", "", bytes.NewBuffer(nil))
 	checkErr(t, err)
 	params := make(url.Values)
 	q := "SELECT column FROM table"
@@ -91,6 +91,25 @@ func TestGetQuerySnippetGETBody(t *testing.T) {
 	}
 }
 
+func TestGetQuerySnippetGETBothQueryAndBody(t *testing.T) {
+	queryPart := "SELECT column"
+	bodyPart := "FROM table"
+	expectedQuery := "SELECT column\nFROM table"
+
+	body := bytes.NewBufferString(bodyPart)
+	req, err := http.NewRequest("GET", "", body)
+	checkErr(t, err)
+
+	params := make(url.Values)
+	params.Set("query", queryPart)
+	req.URL.RawQuery = params.Encode()
+
+	query := getQuerySnippet(req)
+	if query != expectedQuery {
+		t.Fatalf("got: %q; expected: %q", query, expectedQuery)
+	}
+}
+
 func TestGetQuerySnippetPOST(t *testing.T) {
 	q := "SELECT column FROM table"
 	body := bytes.NewBufferString(q)