feat: add support for WITH ROLLUP

Signed-off-by: Andres Taylor <[email protected]>
vitessio · May 13, 2024 · 7089aed · 7089aed
1 parent eb22cfb
commit 7089aed
Show file tree

Hide file tree

Showing 35 changed files with 7,247 additions and 7,125 deletions.
diff --git a/go/test/endtoend/vtgate/queries/random/query_gen.go b/go/test/endtoend/vtgate/queries/random/query_gen.go
@@ -423,7 +423,7 @@ func (sg *selectGenerator) createGroupBy(tables []tableT) (grouping []column) {
 		if col.typ == "date" && !testFailingQueries {
 			continue
 		}
-		sg.sel.GroupBy = append(sg.sel.GroupBy, col.getASTExpr())
+		sg.sel.AddGroupBy(col.getASTExpr())
 
 		// add to select
 		if rand.IntN(2) < 1 {
@@ -478,8 +478,10 @@ func (sg *selectGenerator) createAggregations(tables []tableT) (aggregates []col
 // orders on all grouping expressions and on random SelectExprs
 func (sg *selectGenerator) createOrderBy() {
 	// always order on grouping expressions
-	for _, expr := range sg.sel.GroupBy {
-		sg.sel.OrderBy = append(sg.sel.OrderBy, sqlparser.NewOrder(expr, getRandomOrderDirection()))
+	if sg.sel.GroupBy != nil {
+		for _, expr := range sg.sel.GroupBy.Exprs {
+			sg.sel.OrderBy = append(sg.sel.OrderBy, sqlparser.NewOrder(expr, getRandomOrderDirection()))
+		}
 	}
 
 	// randomly order on SelectExprs

diff --git a/go/vt/sqlparser/ast.go b/go/vt/sqlparser/ast.go
@@ -269,7 +269,7 @@ type (
 		Comments    *ParsedComments
 		SelectExprs SelectExprs
 		Where       *Where
-		GroupBy     GroupBy
+		GroupBy     *GroupBy
 		Having      *Where
 		Windows     NamedWindows
 		OrderBy     OrderBy
@@ -3451,7 +3451,10 @@ type ConvertType struct {
 }
 
 // GroupBy represents a GROUP BY clause.
-type GroupBy []Expr
+type GroupBy struct {
+	Exprs      []Expr
+	WithRollup bool
+}
 
 // OrderBy represents an ORDER By clause.
 type OrderBy []*Order

diff --git a/go/vt/sqlparser/ast_clone.go b/go/vt/sqlparser/ast_clone.go
diff --git a/go/vt/sqlparser/ast_copy_on_rewrite.go b/go/vt/sqlparser/ast_copy_on_rewrite.go
diff --git a/go/vt/sqlparser/ast_equals.go b/go/vt/sqlparser/ast_equals.go
diff --git a/go/vt/sqlparser/ast_format.go b/go/vt/sqlparser/ast_format.go
@@ -1922,12 +1922,18 @@ func (node *When) Format(buf *TrackedBuffer) {
 }
 
 // Format formats the node.
-func (node GroupBy) Format(buf *TrackedBuffer) {
+func (node *GroupBy) Format(buf *TrackedBuffer) {
+	if node == nil {
+		return
+	}
 	prefix := " group by "
-	for _, n := range node {
+	for _, n := range node.Exprs {
 		buf.astPrintf(node, "%s%v", prefix, n)
 		prefix = ", "
 	}
+	if node.WithRollup {
+		buf.literal(" with rollup")
+	}
 }
 
 // Format formats the node.

diff --git a/go/vt/sqlparser/ast_format_fast.go b/go/vt/sqlparser/ast_format_fast.go
diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go
@@ -788,10 +788,19 @@ func NewLimitWithoutOffset(rowCount int) *Limit {
 }
 
 // NewSelect is used to create a select statement
-func NewSelect(comments Comments, exprs SelectExprs, selectOptions []string, into *SelectInto, from TableExprs, where *Where, groupBy GroupBy, having *Where, windows NamedWindows) *Select {
+func NewSelect(
+	comments Comments,
+	exprs SelectExprs,
+	selectOptions []string,
+	into *SelectInto,
+	from TableExprs,
+	where *Where,
+	groupBy *GroupBy,
+	having *Where,
+	windows NamedWindows,
+) *Select {
 	var cache *bool
 	var distinct, straightJoinHint, sqlFoundRows bool
-
 	for _, option := range selectOptions {
 		switch strings.ToLower(option) {
 		case DistinctStr:
@@ -1168,13 +1177,17 @@ func (node *Select) AddHaving(expr Expr) {
 
 // AddGroupBy adds a grouping expression, unless it's already present
 func (node *Select) AddGroupBy(expr Expr) {
-	for _, gb := range node.GroupBy {
+	if node.GroupBy == nil {
+		node.GroupBy = &GroupBy{Exprs: []Expr{expr}}
+		return
+	}
+	for _, gb := range node.GroupBy.Exprs {
 		if Equals.Expr(gb, expr) {
 			// group by columns are sets - duplicates don't add anything, so we can just skip these
 			return
 		}
 	}
-	node.GroupBy = append(node.GroupBy, expr)
+	node.GroupBy.Exprs = append(node.GroupBy.Exprs, expr)
 }
 
 // AddWhere adds the boolean expression to the