Skip to content

Commit

Permalink
add correct HTTP query handler
Browse files Browse the repository at this point in the history
  • Loading branch information
tochka committed Oct 28, 2020
1 parent 181b454 commit 7ee2d20
Show file tree
Hide file tree
Showing 3 changed files with 99 additions and 11 deletions.
38 changes: 38 additions & 0 deletions main_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
51 changes: 41 additions & 10 deletions utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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 {
Expand All @@ -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.
Expand Down
21 changes: 20 additions & 1 deletion utils_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)
Expand Down

0 comments on commit 7ee2d20

Please sign in to comment.