Add name alias support for select statement and order by can use name alias field

Author: blacktear23

kvcmds/cmd_query.go

@@ -59,7 +59,7 @@ func (c QueryCmd) Handler() func(ctx context.Context) {
 				return err
 			}
 			ret := [][]string{
-				plan.FieldNames(),
+				plan.FieldNameList(),
 			}
 			for {
 				cols, err := plan.Next()

query/ast.go

@@ -136,11 +136,12 @@ type Expression interface {
 }
 
 type SelectStmt struct {
-	AllFields bool
-	Fields    []Expression
-	Where     *WhereStmt
-	Order     *OrderStmt
-	Limit     *LimitStmt
+	AllFields  bool
+	FieldNames []string
+	Fields     []Expression
+	Where      *WhereStmt
+	Order      *OrderStmt
+	Limit      *LimitStmt
 }
 
 type WhereStmt struct {

query/lexer.go

@@ -28,6 +28,7 @@ const (
 	DESC     TokenType = 17
 	TRUE     TokenType = 18
 	FALSE    TokenType = 19
+	AS       TokenType = 20
 )
 
 var (
@@ -147,6 +148,23 @@ func (l *Lexer) Split() []*Token {
 			} else {
 				curr += string(char)
 			}
+		case '`':
+			if !strStart {
+				strStart = true
+				strStartChar = char
+				tokStartPos = i
+			} else if strStartChar == char {
+				strStart = false
+				token := &Token{
+					Tp:   NAME,
+					Data: curr,
+					Pos:  tokStartPos,
+				}
+				ret = append(ret, token)
+				curr = ""
+			} else {
+				curr += string(char)
+			}
 		case '~', '^', '=', '!', '*', '+', '-', '/', '>', '<':
 			if strStart {
 				curr += string(char)
@@ -342,6 +360,9 @@ func buildToken(curr string, pos int) *Token {
 	case "false":
 		token.Tp = FALSE
 		return token
+	case "as":
+		token.Tp = AS
+		return token
 	default:
 		if isNumber(curr) {
 			token.Tp = NUMBER

query/optimizer.go

@@ -42,10 +42,11 @@ func (o *Optimizer) BuildPlan(t Txn) (*ProjectionPlan, error) {
 			return nil, err
 		}
 		return &ProjectionPlan{
-			Txn:       t,
-			ChildPlan: fp,
-			AllFields: o.stmt.AllFields,
-			Fields:    o.stmt.Fields,
+			Txn:        t,
+			ChildPlan:  fp,
+			AllFields:  o.stmt.AllFields,
+			FieldNames: o.stmt.FieldNames,
+			Fields:     o.stmt.Fields,
 		}, nil
 	}
 
@@ -64,10 +65,11 @@ func (o *Optimizer) BuildPlan(t Txn) (*ProjectionPlan, error) {
 	}
 
 	return &ProjectionPlan{
-		Txn:       t,
-		ChildPlan: fp,
-		AllFields: o.stmt.AllFields,
-		Fields:    o.stmt.Fields,
+		Txn:        t,
+		ChildPlan:  fp,
+		AllFields:  o.stmt.AllFields,
+		FieldNames: o.stmt.FieldNames,
+		Fields:     o.stmt.Fields,
 	}, nil
 }
 

query/parser.go

@@ -6,9 +6,10 @@ import (
 )
 
 var (
-	ErrSyntaxStartWhere    = errors.New("Syntax Error: not start with `where`")
-	ErrSyntaxEmptyFields   = errors.New("Syntax Error: empty select fields")
-	ErrSyntaxInvalidFields = errors.New("Syntax Error: invalid fields")
+	ErrSyntaxStartWhere       = errors.New("Syntax Error: not start with `where`")
+	ErrSyntaxEmptyFields      = errors.New("Syntax Error: empty select fields")
+	ErrSyntaxInvalidFields    = errors.New("Syntax Error: invalid fields")
+	ErrSyntaxInvalidFieldName = errors.New("Syntax Error: invalid field name")
 )
 
 const MaxNestLevel = 1e5
@@ -266,9 +267,10 @@ func (p *Parser) parseOperand() (Expression, error) {
 
 func (p *Parser) parseSelect() (*SelectStmt, error) {
 	var (
-		fields    = []Expression{}
-		allFields = false
-		err       error
+		fields     = []Expression{}
+		fieldNames = []string{}
+		allFields  = false
+		err        error
 	)
 	err = p.expect(&Token{Tp: SELECT, Data: "select"})
 	if err != nil {
@@ -291,7 +293,17 @@ func (p *Parser) parseSelect() (*SelectStmt, error) {
 		if err != nil {
 			return nil, err
 		}
+		fieldName := field.String()
+		if p.tok != nil && p.tok.Tp == AS {
+			p.next()
+			if p.tok.Tp != NAME {
+				return nil, ErrSyntaxInvalidFieldName
+			}
+			fieldName = p.tok.Data
+		}
+
 		fields = append(fields, field)
+		fieldNames = append(fieldNames, fieldName)
 		if p.tok != nil && p.tok.Tp == WHERE {
 			break
 		}
@@ -307,8 +319,9 @@ func (p *Parser) parseSelect() (*SelectStmt, error) {
 	}
 
 	return &SelectStmt{
-		Fields:    fields,
-		AllFields: allFields,
+		Fields:     fields,
+		FieldNames: fieldNames,
+		AllFields:  allFields,
 	}, nil
 }
 
@@ -353,7 +366,7 @@ func (p *Parser) parseLimit() (*LimitStmt, error) {
 	return ret, nil
 }
 
-func (p *Parser) parseOrderBy() (*OrderStmt, error) {
+func (p *Parser) parseOrderBy(selStmt *SelectStmt) (*OrderStmt, error) {
 	var (
 		err         error
 		shouldBreak bool         = false
@@ -377,6 +390,37 @@ func (p *Parser) parseOrderBy() (*OrderStmt, error) {
 		switch field.ReturnType() {
 		case TSTR, TNUMBER, TBOOL:
 			break
+		case TIDENT:
+			// Check can convert to name expression
+			nexpr, ok := field.(*NameExpr)
+			if !ok {
+				return nil, errors.New("Syntax error: invalid field name")
+			}
+			// Check has select statement
+			if selStmt == nil {
+				return nil, fmt.Errorf("Syntax error: unknown field %s in select statement", nexpr.Data)
+			}
+			fieldName := nexpr.Data
+			foundIdx := -1
+			// Found field in select statement
+			for i, fname := range selStmt.FieldNames {
+				if fname == fieldName {
+					foundIdx = i
+					break
+				}
+			}
+			// Not found return error
+			if foundIdx < 0 {
+				return nil, fmt.Errorf("Syntax error: unknown field %s in select statement", nexpr.Data)
+			}
+			// Found then check return type is valid
+			field = selStmt.Fields[foundIdx]
+			switch field.ReturnType() {
+			case TSTR, TNUMBER, TBOOL:
+				break
+			default:
+				return nil, errors.New("Syntax error: order by field is wrong type.")
+			}
 		default:
 			return nil, errors.New("Syntax error: order by field is wrong type.")
 		}
@@ -453,7 +497,7 @@ func (p *Parser) Parse() (*SelectStmt, error) {
 			if orderStmt != nil {
 				return nil, errors.New("Syntax error duplicate order by expression")
 			}
-			orderStmt, err = p.parseOrderBy()
+			orderStmt, err = p.parseOrderBy(selectStmt)
 			if err != nil {
 				return nil, err
 			}

query/plan.go

@@ -276,21 +276,18 @@ func (p *LimitPlan) Explain() []string {
 }
 
 type ProjectionPlan struct {
-	Txn       Txn
-	ChildPlan Plan
-	AllFields bool
-	Fields    []Expression
+	Txn        Txn
+	ChildPlan  Plan
+	AllFields  bool
+	FieldNames []string
+	Fields     []Expression
 }
 
-func (p *ProjectionPlan) FieldNames() []string {
+func (p *ProjectionPlan) FieldNameList() []string {
 	if p.AllFields {
 		return []string{"Key", "Value"}
 	}
-	ret := []string{}
-	for _, field := range p.Fields {
-		ret = append(ret, field.String())
-	}
-	return ret
+	return p.FieldNames
 }
 
 func (p *ProjectionPlan) Next() ([]Column, error) {

utils/utils.go

@@ -33,8 +33,9 @@ func PrintTable(data [][]string) {
 
 func PrintTableNoWrap(data [][]string) {
 	table := tablewriter.NewWriter(os.Stdout)
-	table.SetHeader(data[0])
 	table.SetAutoWrapText(false)
+	table.SetAutoFormatHeaders(false)
+	table.SetHeader(data[0])
 	table.SetBorders(tablewriter.Border{Left: true, Top: true, Right: true, Bottom: true})
 	table.SetCenterSeparator("|")
 	table.AppendBulk(data[1:])