Add support of GROUPING SETS syntax for the group by clause (#150)

Author: git-hulk

parser/ast.go

@@ -5048,6 +5048,7 @@ func (w *PrewhereClause) Accept(visitor ASTVisitor) error {
 
 type GroupByClause struct {
 	GroupByPos    Pos
+	GroupByEnd    Pos
 	AggregateType string
 	Expr          Expr
 	WithCube      bool
@@ -5060,7 +5061,7 @@ func (g *GroupByClause) Pos() Pos {
 }
 
 func (g *GroupByClause) End() Pos {
-	return g.Expr.End()
+	return g.GroupByEnd
 }
 
 func (g *GroupByClause) String() string {
@@ -5069,7 +5070,9 @@ func (g *GroupByClause) String() string {
 	if g.AggregateType != "" {
 		builder.WriteString(g.AggregateType)
 	}
-	builder.WriteString(g.Expr.String())
+	if g.Expr != nil {
+		builder.WriteString(g.Expr.String())
+	}
 	if g.WithCube {
 		builder.WriteString(" WITH CUBE")
 	}
@@ -5085,8 +5088,10 @@ func (g *GroupByClause) String() string {
 func (g *GroupByClause) Accept(visitor ASTVisitor) error {
 	visitor.enter(g)
 	defer visitor.leave(g)
-	if err := g.Expr.Accept(visitor); err != nil {
-		return err
+	if g.Expr != nil {
+		if err := g.Expr.Accept(visitor); err != nil {
+			return err
+		}
 	}
 	return visitor.VisitGroupByExpr(g)
 }

parser/keyword.go

@@ -89,6 +89,7 @@ const (
 	KeywordGrant        = "GRANT"
 	KeywordGranularity  = "GRANULARITY"
 	KeywordGroup        = "GROUP"
+	KeywordGrouping     = "GROUPING"
 	KeywordHaving       = "HAVING"
 	KeywordHierarchical = "HIERARCHICAL"
 	KeywordHour         = "HOUR"
@@ -181,6 +182,7 @@ const (
 	KeywordSemi         = "SEMI"
 	KeywordSends        = "SENDS"
 	KeywordSet          = "SET"
+	KeywordSets         = "SETS"
 	KeywordSettings     = "SETTINGS"
 	KeywordShow         = "SHOW"
 	KeywordShutdown     = "SHUTDOWN"
@@ -317,6 +319,7 @@ var keywords = NewSet(
 	KeywordGrant,
 	KeywordGranularity,
 	KeywordGroup,
+	KeywordGrouping,
 	KeywordHaving,
 	KeywordHierarchical,
 	KeywordHour,
@@ -409,6 +412,7 @@ var keywords = NewSet(
 	KeywordSemi,
 	KeywordSends,
 	KeywordSet,
+	KeywordSets,
 	KeywordSettings,
 	KeywordShow,
 	KeywordShutdown,

parser/parser_query.go

@@ -419,17 +419,22 @@ func (p *Parser) parseGroupByClause(pos Pos) (*GroupByClause, error) {
 	var expr Expr
 	var err error
 	aggregateType := ""
-	if p.matchKeyword(KeywordCube) || p.matchKeyword(KeywordRollup) {
+	switch {
+	case p.matchKeyword(KeywordCube) || p.matchKeyword(KeywordRollup):
 		aggregateType = p.last().String
 		_ = p.lexer.consumeToken()
 		expr, err = p.parseFunctionParams(p.Pos())
-	} else {
+	case p.tryConsumeKeywords(KeywordGrouping, KeywordSets):
+		aggregateType = "GROUPING SETS"
+		expr, err = p.parseFunctionParams(p.Pos())
+	case p.tryConsumeKeywords(KeywordAll):
+		aggregateType = "ALL"
+	default:
 		expr, err = p.parseColumnExprListWithLParen(p.Pos())
 	}
 	if err != nil {
 		return nil, err
 	}
-
 	groupBy := &GroupByClause{
 		GroupByPos:    pos,
 		AggregateType: aggregateType,
@@ -449,6 +454,7 @@ func (p *Parser) parseGroupByClause(pos Pos) (*GroupByClause, error) {
 			return nil, fmt.Errorf("expected CUBE, ROLLUP or TOTALS, got %s", p.lastTokenKind())
 		}
 	}
+	groupBy.GroupByEnd = p.Pos()
 
 	return groupBy, nil
 }

parser/testdata/ddl/output/alter_table_add_projection.sql.golden.json

@@ -50,6 +50,7 @@
             },
             "GroupBy": {
               "GroupByPos": 86,
+              "GroupByEnd": 105,
               "AggregateType": "",
               "Expr": {
                 "ListPos": 95,

parser/testdata/query/format/select_with_group_by.sql

@@ -0,0 +1,26 @@
+-- Origin SQL:
+SELECT
+    datacenter,
+    distro,
+    SUM (quantity) AS qty
+FROM
+    servers
+GROUP BY
+    GROUPING SETS(
+    (datacenter,distro),
+    (datacenter),
+    (distro),
+    ()
+);
+
+SELECT
+    datacenter,
+    distro,
+    SUM (quantity) AS qty
+FROM
+    servers
+GROUP BY ALL;
+
+-- Format SQL:
+SELECT datacenter, distro, SUM(quantity) AS qty FROM servers GROUP BY GROUPING SETS((datacenter, distro), (datacenter), (distro), ());
+SELECT datacenter, distro, SUM(quantity) AS qty FROM servers GROUP BY ALL;

parser/testdata/query/output/select_simple.sql.golden.json

@@ -357,6 +357,7 @@
     },
     "GroupBy": {
       "GroupByPos": 239,
+      "GroupByEnd": 258,
       "AggregateType": "",
       "Expr": {
         "ListPos": 248,

parser/testdata/query/output/select_simple_with_group_by_with_cube_totals.sql.golden.json

@@ -78,6 +78,7 @@
     "Where": null,
     "GroupBy": {
       "GroupByPos": 37,
+      "GroupByEnd": 76,
       "AggregateType": "CUBE",
       "Expr": {
         "LeftParenPos": 50,