Minor improvements for /query endpoint

Signed-off-by: Carlos Martín <[email protected]>

Author: carlosms

server/handler/query.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"strings"
 
 	"github.com/src-d/gitbase-playground/server/serializer"
 	"gopkg.in/bblfsh/sdk.v1/uast"
@@ -20,20 +21,22 @@ type queryRequest struct {
 
 // genericVals returns a slice of interface{}, each one a pointer to the proper
 // type for each column
-func genericVals(colTypes []*sql.ColumnType) []interface{} {
+func genericVals(colTypes []string) []interface{} {
 	columnValsPtr := make([]interface{}, len(colTypes))
 
 	for i, colType := range colTypes {
-		switch colType.DatabaseTypeName() {
+		switch colType {
 		case "BIT":
 			columnValsPtr[i] = new(sql.NullBool)
-		case "TIMESTAMP":
+		case "TIMESTAMP", "DATE", "DATETIME":
 			columnValsPtr[i] = new(mysql.NullTime)
-		case "INT":
+		case "INT", "MEDIUMINT", "BIGINT", "SMALLINT", "TINYINT":
 			columnValsPtr[i] = new(sql.NullInt64)
+		case "DOUBLE", "FLOAT":
+			columnValsPtr[i] = new(sql.NullFloat64)
 		case "JSON":
 			columnValsPtr[i] = new([]byte)
-		default: // "TEXT" and any others
+		default: // All the text and binary variations
 			columnValsPtr[i] = new(sql.NullString)
 		}
 	}
@@ -47,7 +50,6 @@ func Query(db *sql.DB) RequestProcessFunc {
 	return func(r *http.Request) (*serializer.Response, error) {
 		var queryRequest queryRequest
 		body, err := ioutil.ReadAll(r.Body)
-		defer r.Body.Close()
 		if err == nil {
 			err = json.Unmarshal(body, &queryRequest)
 		}
@@ -56,26 +58,26 @@ func Query(db *sql.DB) RequestProcessFunc {
 			return nil, err
 		}
 
-		// TODO (carlosms) this only works if the query does not end in ;
-		// and does not have a limit. It will also fail for queries like
-		// DESCRIBE TABLE
-		query := fmt.Sprintf("%s LIMIT %d", queryRequest.Query, queryRequest.Limit)
+		query := addLimit(queryRequest.Query, queryRequest.Limit)
 		rows, err := db.Query(query)
 		if err != nil {
+			if mysqlErr, ok := err.(*mysql.MySQLError); ok {
+				return nil, serializer.NewMySQLError(
+					http.StatusBadRequest,
+					mysqlErr.Number,
+					mysqlErr.Message)
+			}
+
 			return nil, serializer.NewHTTPError(http.StatusBadRequest, err.Error())
 		}
 		defer rows.Close()
 
-		columnNames, err := rows.Columns()
+		columnNames, columnTypes, err := columnsInfo(rows)
 		if err != nil {
 			return nil, err
 		}
 
-		colTypes, err := rows.ColumnTypes()
-		if err != nil {
-			return nil, err
-		}
-		columnValsPtr := genericVals(colTypes)
+		columnValsPtr := genericVals(columnTypes)
 
 		tableData := make([]map[string]interface{}, 0)
 
@@ -84,7 +86,7 @@ func Query(db *sql.DB) RequestProcessFunc {
 				return nil, err
 			}
 
-			colData := make(map[string]interface{})
+			colData := make(map[string]interface{}, len(columnTypes))
 
 			for i, val := range columnValsPtr {
 				colData[columnNames[i]] = nil
@@ -111,23 +113,21 @@ func Query(db *sql.DB) RequestProcessFunc {
 						colData[columnNames[i]] = sqlVal.String
 					}
 				case *[]byte:
-					// TODO (carlosms) this may not be an array always
-					var protobufs [][]byte
-					if err := json.Unmarshal(*val.(*[]byte), &protobufs); err != nil {
-						return nil, err
-					}
-
-					nodes := make([]*uast.Node, len(protobufs))
-
-					for i, v := range protobufs {
-						node := uast.NewNode()
-						if err = node.Unmarshal(v); err != nil {
+					// TODO (carlosms) this may not always be a JSON array, it could
+					// be a JSON object
+
+					// DatabaseTypeName JSON is used for arrays, but we don't know the
+					// type of each element. We try with uast node first and text later
+					nodes, err := unmarshallUAST(val)
+					if err == nil {
+						colData[columnNames[i]] = nodes
+					} else {
+						var strings []string
+						if err := json.Unmarshal(*val.(*[]byte), &strings); err != nil {
 							return nil, err
 						}
-						nodes[i] = node
+						colData[columnNames[i]] = strings
 					}
-
-					colData[columnNames[i]] = nodes
 				}
 			}
 
@@ -138,6 +138,58 @@ func Query(db *sql.DB) RequestProcessFunc {
 			return nil, err
 		}
 
-		return serializer.NewQueryResponse(tableData, columnNames), nil
+		return serializer.NewQueryResponse(tableData, columnNames, columnTypes), nil
+	}
+}
+
+// columnsInfo returns the column names and column types, or error
+func columnsInfo(rows *sql.Rows) ([]string, []string, error) {
+	names, err := rows.Columns()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	types, err := rows.ColumnTypes()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	typesStr := make([]string, len(types))
+	for i, colType := range types {
+		typesStr[i] = colType.DatabaseTypeName()
+	}
+
+	return names, typesStr, nil
+}
+
+// unmarshallUAST tries to cast data as [][]byte and unmarshall uast nodes
+func unmarshallUAST(data interface{}) ([]*uast.Node, error) {
+	var protobufs [][]byte
+	if err := json.Unmarshal(*data.(*[]byte), &protobufs); err != nil {
+		return nil, err
+	}
+
+	nodes := make([]*uast.Node, len(protobufs))
+
+	for i, v := range protobufs {
+		node := uast.NewNode()
+		if err := node.Unmarshal(v); err != nil {
+			return nil, err
+		}
+		nodes[i] = node
+	}
+
+	return nodes, nil
+}
+
+// addLimit adds LIMIT to the query, performing basic tests to skip it
+// for DESCRIBE TABLE, SHOW TABLES, and avoid '; limit'
+func addLimit(query string, limit int) string {
+	upper := strings.ToUpper(query)
+	if strings.Contains(upper, "DESCRIBE") || strings.Contains(upper, "SHOW") {
+		return query
 	}
+
+	trimmed := strings.TrimRight(strings.TrimSpace(query), ";")
+	return fmt.Sprintf("%s LIMIT %d", trimmed, limit)
 }

server/serializer/serializers.go

@@ -20,9 +20,10 @@ type Response struct {
 }
 
 type httpError struct {
-	Status  int    `json:"status"`
-	Title   string `json:"title"`
-	Details string `json:"details,omitempty"`
+	Status    int    `json:"status"`
+	Title     string `json:"title"`
+	Details   string `json:"details,omitempty"`
+	MySQLCode uint16 `json:"mysqlCode,omitempty"`
 }
 
 // StatusCode returns the Status of the httpError
@@ -48,6 +49,11 @@ func NewHTTPError(statusCode int, msg ...string) HTTPError {
 	return httpError{Status: statusCode, Title: strings.Join(msg, " ")}
 }
 
+// NewHTTPError returns an Error with the MySQL error code
+func NewMySQLError(statusCode int, mysqlCode uint16, msg ...string) HTTPError {
+	return httpError{Status: statusCode, MySQLCode: mysqlCode, Title: strings.Join(msg, " ")}
+}
+
 func newResponse(data interface{}, meta interface{}) *Response {
 	if data == nil {
 		return &Response{
@@ -78,9 +84,10 @@ func NewVersionResponse(version string) *Response {
 
 type queryMetaResponse struct {
 	Headers []string `json:"headers"`
+	Types   []string `json:"types"`
 }
 
 // NewQueryResponse returns a Response with table headers and row contents
-func NewQueryResponse(rows []map[string]interface{}, columnNames []string) *Response {
-	return newResponse(rows, queryMetaResponse{columnNames})
+func NewQueryResponse(rows []map[string]interface{}, columnNames, columnTypes []string) *Response {
+	return newResponse(rows, queryMetaResponse{columnNames, columnTypes})
 }