Merge pull request #264 from carlosms/fix-241

Use mysql KILL to cancel queries

Author: carlosms

server/handler/common_unit_test.go

@@ -39,5 +39,9 @@ func (suite *HandlerUnitSuite) SetupTest() {
 }
 
 func (suite *HandlerUnitSuite) TearDownTest() {
-	suite.db.Close()
+	defer suite.db.Close()
+
+	if err := suite.mock.ExpectationsWereMet(); err != nil {
+		suite.FailNowf("there were unfulfilled expectations:", err.Error())
+	}
 }

server/handler/query.go

@@ -66,7 +66,31 @@ func Query(db service.SQLDB) RequestProcessFunc {
 		}
 
 		query, limitSet := addLimit(queryRequest.Query, queryRequest.Limit)
-		rows, err := db.QueryContext(r.Context(), query)
+
+		// go-sql-driver/mysql QueryContext stops waiting for the query results on
+		// context cancel, but it does not actually cancel the query on the server
+
+		c := make(chan error, 1)
+
+		var rows *sql.Rows
+		go func() {
+			rows, err = db.QueryContext(r.Context(), query)
+			c <- err
+		}()
+
+		// It may happen that the QueryContext returns with an error because of
+		// context cancellation. In this case, the select may enter on the second
+		// case. We check if the context was cancelled with Err() instead of Done()
+		select {
+		case <-r.Context().Done():
+		case err = <-c:
+		}
+
+		if r.Context().Err() != nil {
+			killQuery(db, query)
+			return nil, dbError(r.Context().Err())
+		}
+
 		if err != nil {
 			return nil, dbError(err)
 		}
@@ -103,6 +127,47 @@ func Query(db service.SQLDB) RequestProcessFunc {
 	}
 }
 
+func killQuery(db service.SQLDB, query string) {
+	pRows, pErr := db.Query("SHOW FULL PROCESSLIST")
+	if pErr != nil {
+		// TODO (carlosms) log error when we migrate to go-log
+		return
+	}
+	defer pRows.Close()
+
+	found := false
+	var foundID int
+
+	for pRows.Next() {
+		var id int
+		var info sql.NullString
+		var rb sql.RawBytes
+		// The columns are:
+		// Id, User, Host, db, Command, Time, State, Info
+		// gitbase returns the query on "Info".
+		if err := pRows.Scan(&id, &rb, &rb, &rb, &rb, &rb, &rb, &info); err != nil {
+			// TODO (carlosms) log error when we migrate to go-log
+			return
+		}
+
+		if info.Valid && info.String == query {
+			if found {
+				// Found more than one match for current query, we cannot know which
+				// one is ours. Skip the cancellation
+				// TODO (carlosms) log error when we migrate to go-log
+				return
+			}
+
+			found = true
+			foundID = id
+		}
+	}
+
+	if found {
+		db.Exec(fmt.Sprintf("KILL %d", foundID))
+	}
+}
+
 // columnsInfo returns the column names and column types, or error
 func columnsInfo(rows *sql.Rows) ([]string, []string, error) {
 	names, err := rows.Columns()

server/handler/query_test.go

@@ -83,8 +83,6 @@ func (suite *QuerySuite) TestBadRequest() {
 		`{"query": "select * from repositories", "limit": "string"}`,
 	}
 
-	suite.mock.ExpectQuery(".*").WillReturnError(fmt.Errorf("forced err"))
-
 	for _, tc := range testCases {
 		suite.T().Run(tc, func(t *testing.T) {
 			a := assert.New(t)
@@ -201,8 +199,16 @@ func (suite *QuerySuite) TestQueryAbort() {
 	// Ideally we would test that the sql query context is canceled, but
 	// go-sqlmock does not have something like ExpectContextCancellation
 
-	mockRows := sqlmock.NewRows([]string{"a", "b", "c", "d"})
-	suite.mock.ExpectQuery(".*").WillDelayFor(2 * time.Second).WillReturnRows(mockRows)
+	mockRows := sqlmock.NewRows([]string{"a", "b", "c", "d"}).AddRow(1, "one", 1.5, 100)
+	suite.mock.ExpectQuery(`select \* from repositories`).WillDelayFor(2 * time.Second).WillReturnRows(mockRows)
+
+	mockProcessRows := sqlmock.NewRows(
+		[]string{"Id", "User", "Host", "db", "Command", "Time", "State", "Info"}).
+		AddRow(1234, nil, "localhost:3306", nil, "query", 2, "SquashedTable(refs, commit_files, files)(1/5)", "select * from files").
+		AddRow(1288, nil, "localhost:3306", nil, "query", 2, "SquashedTable(refs, commit_files, files)(1/5)", "select * from repositories")
+	suite.mock.ExpectQuery("SHOW FULL PROCESSLIST").WillReturnRows(mockProcessRows)
+
+	suite.mock.ExpectExec("KILL 1288")
 
 	json := `{"query": "select * from repositories"}`
 	req, _ := http.NewRequest("POST", "/query", strings.NewReader(json))
@@ -227,6 +233,11 @@ func (suite *QuerySuite) TestQueryAbort() {
 		handler.ServeHTTP(res, req)
 	}()
 
+	// Without this wait the Request is cancelled before the handler has time to
+	// start the query. Which also works fine, but we want to test a cancellation
+	// for a query that is in progress
+	time.Sleep(200 * time.Millisecond)
+
 	cancel()
 
 	wg.Wait()

server/service/common.go

@@ -12,4 +12,5 @@ type SQLDB interface {
 	QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error)
 	QueryRow(query string, args ...interface{}) *sql.Row
 	Ping() error
+	Exec(query string, args ...interface{}) (sql.Result, error)
 }

server/testing/common.go

@@ -33,6 +33,11 @@ func (db *MockDB) QueryRow(query string, args ...interface{}) *sql.Row {
 	return nil
 }
 
+// Exec executes a query without returning any rows
+func (db *MockDB) Exec(query string, args ...interface{}) (sql.Result, error) {
+	return nil, nil
+}
+
 // As returned by gitbase v0.17.0-rc.4, SELECT UAST('console.log("test")', 'JavaScript') AS uast
 const (
 	UASTMarshaled     = "\x00\x00\x02\x16\n\x04File\x1a\xfc\x03\n\aProgram\x12\x17\n\finternalRole\x12\aprogram\x12\x14\n\nsourceType\x12\x06module\x1a\xb0\x03\n\x13ExpressionStatement\x12\x14\n\finternalRole\x12\x04body\x1a\xf1\x02\n\x0eCallExpression\x12\x1a\n\finternalRole\x12\nexpression\x1a\xdc\x01\n\x10MemberExpression\x12\x16\n\finternalRole\x12\x06callee\x12\x11\n\bcomputed\x12\x05false\x1aC\n\nIdentifier\x12\x16\n\finternalRole\x12\x06object\x12\x0f\n\x04Name\x12\aconsole*\x04\x10\x01\x18\x012\x06\b\a\x10\x01\x18\b\x1aC\n\nIdentifier\x12\v\n\x04Name\x12\x03log\x12\x18\n\finternalRole\x12\bproperty*\x06\b\b\x10\x01\x18\t2\x06\b\v\x10\x01\x18\f*\x04\x10\x01\x18\x012\x06\b\v\x10\x01\x18\f:\x05\x02\x12\x01TU\x1aR\n\x06String\x12\r\n\x05Value\x12\x04test\x12\n\n\x06Format\x12\x00\x12\x19\n\finternalRole\x12\targuments*\x06\b\f\x10\x01\x18\r2\x06\b\x12\x10\x01\x18\x13:\x02T1*\x04\x10\x01\x18\x012\x06\b\x13\x10\x01\x18\x14:\x02\x12T*\x04\x10\x01\x18\x012\x06\b\x13\x10\x01\x18\x14:\x01\x13*\x04\x10\x01\x18\x012\x06\b\x13\x10\x01\x18\x14:\x019*\x04\x10\x01\x18\x012\x06\b\x13\x10\x01\x18\x14:\x01\""