server: add test for /debug/sub-optimal-plan HTTP API (#14302)

pingcap · Jan 6, 2020 · 297e453 · 297e453
1 parent 39ed9f3
commit 297e453
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 18 deletions.
diff --git a/server/http_handler_test.go b/server/http_handler_test.go
@@ -22,6 +22,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"net/http/httputil"
 	"net/url"
 	"os"
 	"sort"
@@ -989,17 +990,66 @@ func (ts *HTTPHandlerTestSuite) TestDebugZip(c *C) {
 	resp, err := http.Get("http://127.0.0.1:10090/debug/zip?seconds=1")
 	c.Assert(err, IsNil)
 	c.Assert(resp.StatusCode, Equals, http.StatusOK)
-	out, err := os.Create("/tmp/tidb_debug.zip")
+	b, err := httputil.DumpResponse(resp, true)
 	c.Assert(err, IsNil)
-	_, err = io.Copy(out, resp.Body)
+	c.Assert(len(b), Greater, 0)
+	c.Assert(resp.Body.Close(), IsNil)
+}
+
+func (ts *HTTPHandlerTestSuite) TestZipInfoForSQL(c *C) {
+	ts.startServer(c)
+	defer ts.stopServer(c)
+
+	db, err := sql.Open("mysql", getDSN())
+	c.Assert(err, IsNil, Commentf("Error connecting"))
+	defer db.Close()
+	dbt := &DBTest{c, db}
+
+	dbt.mustExec("use test")
+	dbt.mustExec("create table if not exists t (a int)")
+
+	urlValues := url.Values{
+		"sql":        {"select * from t"},
+		"current_db": {"test"},
+	}
+	resp, err := http.PostForm("http://127.0.0.1:10090/debug/sub-optimal-plan", urlValues)
 	c.Assert(err, IsNil)
-	fileInfo, err := out.Stat()
+	c.Assert(resp.StatusCode, Equals, http.StatusOK)
+	b, err := httputil.DumpResponse(resp, true)
 	c.Assert(err, IsNil)
-	c.Assert(fileInfo.Size(), Greater, int64(0))
-	err = out.Close()
+	c.Assert(len(b), Greater, 0)
+	c.Assert(resp.Body.Close(), IsNil)
+
+	resp, err = http.PostForm("http://127.0.0.1:10090/debug/sub-optimal-plan?pprof_time=5&timeout=0", urlValues)
+	c.Assert(err, IsNil)
+	c.Assert(resp.StatusCode, Equals, http.StatusOK)
+	b, err = httputil.DumpResponse(resp, true)
+	c.Assert(err, IsNil)
+	c.Assert(len(b), Greater, 0)
+	c.Assert(resp.Body.Close(), IsNil)
+
+	resp, err = http.PostForm("http://127.0.0.1:10090/debug/sub-optimal-plan?pprof_time=5", urlValues)
+	c.Assert(err, IsNil)
+	c.Assert(resp.StatusCode, Equals, http.StatusOK)
+	b, err = httputil.DumpResponse(resp, true)
+	c.Assert(err, IsNil)
+	c.Assert(len(b), Greater, 0)
+	c.Assert(resp.Body.Close(), IsNil)
+
+	resp, err = http.PostForm("http://127.0.0.1:10090/debug/sub-optimal-plan?timeout=1", urlValues)
 	c.Assert(err, IsNil)
-	err = os.Remove("/tmp/tidb_debug.zip")
+	c.Assert(resp.StatusCode, Equals, http.StatusOK)
+	b, err = httputil.DumpResponse(resp, true)
+	c.Assert(err, IsNil)
+	c.Assert(len(b), Greater, 0)
+	c.Assert(resp.Body.Close(), IsNil)
+
+	urlValues.Set("current_db", "non_exists_db")
+	resp, err = http.PostForm("http://127.0.0.1:10090/debug/sub-optimal-plan", urlValues)
 	c.Assert(err, IsNil)
-	err = resp.Body.Close()
+	c.Assert(resp.StatusCode, Equals, http.StatusInternalServerError)
+	b, err = ioutil.ReadAll(resp.Body)
 	c.Assert(err, IsNil)
+	c.Assert(string(b), Equals, "use database non_exists_db failed, err: [schema:1049]Unknown database 'non_exists_db'\n")
+	c.Assert(resp.Body.Close(), IsNil)
 }
diff --git a/server/sql_info_fetcher.go b/server/sql_info_fetcher.go
@@ -15,6 +15,7 @@ package server
 
 import (
 	"archive/zip"
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -87,7 +88,7 @@ func (sh *sqlInfoFetcher) zipInfoForSQL(w http.ResponseWriter, r *http.Request)
 	timeoutString := r.FormValue("timeout")
 	curDB := strings.ToLower(r.FormValue("current_db"))
 	if curDB != "" {
-		_, err = sh.s.Execute(reqCtx, "use %v"+curDB)
+		_, err = sh.s.Execute(reqCtx, fmt.Sprintf("use %v", curDB))
 		if err != nil {
 			serveError(w, http.StatusInternalServerError, fmt.Sprintf("use database %v failed, err: %v", curDB, err))
 			return
@@ -191,7 +192,8 @@ func (sh *sqlInfoFetcher) zipInfoForSQL(w http.ResponseWriter, r *http.Request)
 		resultChan := make(chan *explainAnalyzeResult)
 		go sh.getExplainAnalyze(ctx, sql, resultChan)
 		errChan := make(chan error)
-		go sh.catchCPUProfile(reqCtx, pprofTime, zw, errChan)
+		var buf bytes.Buffer
+		go sh.catchCPUProfile(reqCtx, pprofTime, &buf, errChan)
 		select {
 		case result := <-resultChan:
 			timer.Stop()
@@ -215,7 +217,7 @@ func (sh *sqlInfoFetcher) zipInfoForSQL(w http.ResponseWriter, r *http.Request)
 		case <-timer.C:
 			cancelFunc()
 		}
-		err = <-errChan
+		err = dumpCPUProfile(errChan, &buf, zw)
 		if err != nil {
 			err = sh.writeErrFile(zw, "profile.err.txt", err)
 			terror.Log(err)
@@ -224,6 +226,22 @@ func (sh *sqlInfoFetcher) zipInfoForSQL(w http.ResponseWriter, r *http.Request)
 	}
 }
 
+func dumpCPUProfile(errChan chan error, buf *bytes.Buffer, zw *zip.Writer) error {
+	err := <-errChan
+	if err != nil {
+		return err
+	}
+	fw, err := zw.Create("profile")
+	if err != nil {
+		return err
+	}
+	_, err = fw.Write(buf.Bytes())
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (sh *sqlInfoFetcher) writeErrFile(zw *zip.Writer, name string, err error) error {
 	fw, err1 := zw.Create(name)
 	if err1 != nil {
@@ -255,14 +273,8 @@ func (sh *sqlInfoFetcher) getExplainAnalyze(ctx context.Context, sql string, res
 	resultChan <- &explainAnalyzeResult{rows: rows}
 }
 
-func (sh *sqlInfoFetcher) catchCPUProfile(ctx context.Context, sec int, zw *zip.Writer, errChan chan<- error) {
-	// dump profile
-	fw, err := zw.Create("profile")
-	if err != nil {
-		errChan <- err
-		return
-	}
-	if err := pprof.StartCPUProfile(fw); err != nil {
+func (sh *sqlInfoFetcher) catchCPUProfile(ctx context.Context, sec int, buf *bytes.Buffer, errChan chan<- error) {
+	if err := pprof.StartCPUProfile(buf); err != nil {
 		errChan <- err
 		return
 	}