fix aggregate expressions with aliases, fixes src-d#112. (src-d#116)

Author: smola

engine_test.go

@@ -47,6 +47,11 @@ func TestEngine_Query(t *testing.T) {
 		"SELECT COUNT(*) FROM mytable;",
 		[][]interface{}{{int64(3)}},
 	)
+
+	testQuery(t, e,
+		"SELECT COUNT(*) AS c FROM mytable;",
+		[][]interface{}{{int64(3)}},
+	)
 }
 
 func testQuery(t *testing.T, e *gitql.Engine, q string, r [][]interface{}) {
@@ -64,29 +69,26 @@ func testQuery(t *testing.T, e *gitql.Engine, q string, r [][]interface{}) {
 	assert.NoError(err)
 	assert.Equal(len(r[0]), len(cols))
 
+	vals := make([]interface{}, len(cols))
+	valPtrs := make([]interface{}, len(cols))
+	for i := 0; i < len(cols); i++ {
+		valPtrs[i] = &vals[i]
+	}
+
 	i := 0
 	for {
 		if !res.Next() {
 			break
 		}
 
-		expectedRow := r[i]
-		i++
-
-		row := make([]interface{}, len(expectedRow))
-		for i := range row {
-			i64 := int64(0)
-			row[i] = &i64
-		}
+		err := res.Scan(valPtrs...)
+		assert.NoError(err)
 
-		assert.NoError(res.Scan(row...))
-		for i := range row {
-			row[i] = *(row[i].(*int64))
-		}
-
-		assert.Equal(expectedRow, row)
+		assert.Equal(r[i], vals)
+		i++
 	}
 
+	assert.NoError(res.Err())
 	assert.Equal(len(r), i)
 }
 

sql/parse/parse.go

@@ -190,6 +190,17 @@ func limitToLimit(o sqlparser.Expr, child sql.Node) (*plan.Limit, error) {
 	return plan.NewLimit(n, child), nil
 }
 
+func isAggregate(e sql.Expression) bool {
+	switch v := e.(type) {
+	case *expression.UnresolvedFunction:
+		return v.IsAggregate
+	case *expression.Alias:
+		return isAggregate(v.Child)
+	default:
+		return false
+	}
+}
+
 func selectToProjectOrGroupBy(se sqlparser.SelectExprs, g sqlparser.GroupBy, child sql.Node) (sql.Node, error) {
 	selectExprs, err := selectExprsToExpressions(se)
 	if err != nil {
@@ -199,8 +210,9 @@ func selectToProjectOrGroupBy(se sqlparser.SelectExprs, g sqlparser.GroupBy, chi
 	isAgg := len(g) > 0
 	if !isAgg {
 		for _, e := range selectExprs {
-			if u, ok := e.(*expression.UnresolvedFunction); ok {
-				isAgg = u.IsAggregate
+			if isAggregate(e) {
+				isAgg = true
+				break
 			}
 		}
 	}

sql/plan/group_by.go

@@ -183,12 +183,19 @@ func aggregate(exprs []sql.Expression, rows []sql.Row) sql.Row {
 func exprsToAggregateExprs(exprs []sql.Expression) []sql.AggregationExpression {
 	var r []sql.AggregationExpression
 	for _, e := range exprs {
-		if ae, ok := e.(sql.AggregationExpression); ok {
-			r = append(r, ae)
-		} else {
-			r = append(r, expression.NewFirst(e))
-		}
+		r = append(r, exprToAggregateExpr(e))
 	}
 
 	return r
 }
+
+func exprToAggregateExpr(e sql.Expression) sql.AggregationExpression {
+	switch v := e.(type) {
+	case sql.AggregationExpression:
+		return v
+	case *expression.Alias:
+		return exprToAggregateExpr(v.Child)
+	default:
+		return expression.NewFirst(e)
+	}
+}