Repository: vitessio/vitess
Pull request #11499 (2022-10-14 08:22:39)

bug fix: using self-referencing columns in HAVING should not overflow
## Description
The planner rewrites the HAVING clause in two steps, and in the second part, we were doing the the wrong rewrite, which lead to a cycle in the expression tree. When we then tried to turn the expression back to SQL, we ran into a stack overflow.

Example query that was failing:

```sql
select a+2 as a from user having a = 42
```


## Related Issue(s)

## Checklist

-   [x] "Backport me!" label has been added if this change should be backported
-   [x] Tests were added or are not required
-   [x] Documentation was added or is not required


vitess-bot[bot]: ### Review Checklist
            
Hello reviewers! :wave: Please follow this checklist when reviewing this Pull Request.

#### General
- [ ] Ensure that the Pull Request has a descriptive title.
- [ ] If this is a change that users need to know about, please apply the `release notes (needs details)` label so that merging is blocked unless the summary release notes document is included.

#### If a new flag is being introduced:
- [ ] Is it really necessary to add this flag?
- [ ] Flag names should be clear and intuitive (as far as possible)
- [ ] Help text should be descriptive.
- [ ] Flag names should use dashes (`-`) as word separators rather than underscores (`_`).

#### If a workflow is added or modified:
- [ ] Each item in `Jobs` should be named in order to mark it as `required`.
- [ ] If the workflow should be required, the maintainer team should be notified.

#### Bug fixes
- [ ] There should be at least one unit or end-to-end test.
- [ ] The Pull Request description should include a link to an issue that describes the bug.

#### Non-trivial changes
- [ ] There should be some code comments as to why things are implemented the way they are.

#### New/Existing features
- [ ] Should be documented, either by modifying the existing documentation or creating new documentation.
- [ ] New features should have a link to a feature request issue or an RFC that documents the use cases, corner cases and test cases.

#### Backward compatibility            
- [ ] Protobuf changes should be wire-compatible.
- [ ] Changes to `_vt` tables and RPCs need to be backward compatible.
- [ ] `vtctl` command output order should be stable and `awk`-able.
- [ ] RPC changes should be compatible with vitess-operator
- [ ] If a flag is removed, then it should also be removed from VTop, if used there.


Patch:
diff --git a/go/vt/vtgate/planbuilder/abstract/queryprojection.go b/go/vt/vtgate/planbuilder/abstract/queryprojection.go
index 14572b117b..f3acdd3693 100644
--- a/go/vt/vtgate/planbuilder/abstract/queryprojection.go
+++ b/go/vt/vtgate/planbuilder/abstract/queryprojection.go
@@ -314,39 +314,6 @@ func checkForInvalidAggregations(exp *sqlparser.AliasedExpr) error {
 	}, exp.Expr)
 }
 
-func (qp *QueryProjection) getNonAggrExprNotMatchingGroupByExprs() sqlparser.SelectExpr {
-	for _, expr := range qp.SelectExprs {
-		if expr.Aggr {
-			continue
-		}
-		if !qp.isExprInGroupByExprs(expr) {
-			return expr.Col
-		}
-	}
-	for _, order := range qp.OrderExprs {
-		if !qp.isOrderByExprInGroupBy(order) {
-			return &sqlparser.AliasedExpr{
-				Expr: order.Inner.Expr,
-			}
-		}
-	}
-	return nil
-}
-
-func (qp *QueryProjection) isOrderByExprInGroupBy(order OrderBy) bool {
-	// ORDER BY NULL or Aggregation functions need not be present in group by
-	_, isAggregate := order.WeightStrExpr.(sqlparser.AggrFunc)
-	if sqlparser.IsNull(order.Inner.Expr) || isAggregate {
-		return true
-	}
-	for _, groupByExpr := range qp.groupByExprs {
-		if sqlparser.EqualsExpr(groupByExpr.WeightStrExpr, order.WeightStrExpr) {
-			return true
-		}
-	}
-	return false
-}
-
 func (qp *QueryProjection) isExprInGroupByExprs(expr SelectExpr) bool {
 	for _, groupByExpr := range qp.groupByExprs {
 		exp, err := expr.GetExpr()
@@ -457,10 +424,7 @@ func (qp *QueryProjection) NeedsDistinct() bool {
 func (qp *QueryProjection) AggregationExpressions() (out []Aggr, err error) {
 orderBy:
 	for _, orderExpr := range qp.OrderExprs {
-		if qp.isOrderByExprInGroupBy(orderExpr) {
-			continue orderBy
-		}
-		orderExpr := orderExpr.Inner.Expr
+		orderExpr := orderExpr.WeightStrExpr
 		for _, expr := range qp.SelectExprs {
 			col, ok := expr.Col.(*sqlparser.AliasedExpr)
 			if !ok {
diff --git a/go/vt/vtgate/planbuilder/rewrite.go b/go/vt/vtgate/planbuilder/rewrite.go
index d1b6a789a1..284ffa9c07 100644
--- a/go/vt/vtgate/planbuilder/rewrite.go
+++ b/go/vt/vtgate/planbuilder/rewrite.go
@@ -182,12 +182,8 @@ func rewriteHavingClause(node *sqlparser.Select) {
 					return false
 				}
 				originalExpr, isInMap := selectExprMap[x.Name.Lowered()]
-				if isInMap {
-					if sqlparser.ContainsAggregation(originalExpr) {
-						hasAggr = true
-					} else {
-						cursor.Replace(originalExpr)
-					}
+				if isInMap && sqlparser.ContainsAggregation(originalExpr) {
+					hasAggr = true
 				}
 				return false
 			default:


Test patch:
diff --git a/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go b/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go
index dfcfcc0c42..df95b679a1 100644
--- a/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go
+++ b/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go
@@ -33,7 +33,7 @@ func start(t *testing.T) (utils.MySQLCompare, func()) {
 	deleteAll := func() {
 		_, _ = utils.ExecAllowError(t, mcmp.VtConn, "set workload = oltp")
 
-		tables := []string{"aggr_test", "t3", "t7_xxhash", "aggr_test_dates", "t7_xxhash_idx", "t1", "t2"}
+		tables := []string{"t9", "aggr_test", "t3", "t7_xxhash", "aggr_test_dates", "t7_xxhash_idx", "t1", "t2"}
 		for _, table := range tables {
 			_, _ = mcmp.ExecAndIgnore("delete from " + table)
 		}
@@ -365,3 +365,12 @@ func TestEmptyTableAggr(t *testing.T) {
 	}
 
 }
+
+func TestOrderByCount(t *testing.T) {
+	mcmp, closer := start(t)
+	defer closer()
+
+	mcmp.Exec("insert into t9(id1, id2, id3) values(1, '1', '1'), (2, '2', '2'), (3, '2', '2'), (4, '3', '3'), (5, '3', '3'), (6, '3', '3')")
+
+	mcmp.AssertMatches("SELECT /*vt+ PLANNER=gen4 */ t9.id2 FROM t9 GROUP BY t9.id2 ORDER BY COUNT(t9.id2) DESC", `[[VARCHAR("3")] [VARCHAR("2")] [VARCHAR("1")]]`)
+}
diff --git a/go/test/endtoend/vtgate/queries/aggregation/main_test.go b/go/test/endtoend/vtgate/queries/aggregation/main_test.go
index 65cf3a0343..a859002f44 100644
--- a/go/test/endtoend/vtgate/queries/aggregation/main_test.go
+++ b/go/test/endtoend/vtgate/queries/aggregation/main_test.go
@@ -33,8 +33,8 @@ var (
 	clusterInstance *cluster.LocalProcessCluster
 	vtParams        mysql.ConnParams
 	mysqlParams     mysql.ConnParams
-	keyspaceName    = "ks_union"
-	cell            = "test_union"
+	keyspaceName    = "ks_aggr"
+	cell            = "test_aggr"
 
 	//go:embed schema.sql
 	schemaSQL string
diff --git a/go/test/endtoend/vtgate/queries/aggregation/schema.sql b/go/test/endtoend/vtgate/queries/aggregation/schema.sql
index 944c378304..a6ce793b4e 100644
--- a/go/test/endtoend/vtgate/queries/aggregation/schema.sql
+++ b/go/test/endtoend/vtgate/queries/aggregation/schema.sql
@@ -69,3 +69,4 @@ CREATE TABLE t2 (
     shardKey bigint,
     PRIMARY KEY (id)
 ) ENGINE InnoDB;
+
diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
index 1fc9db15e7..45a6cb7238 100644
--- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json
@@ -4779,5 +4779,45 @@
         "user.user"
       ]
     }
+  },
+  {
+    "comment": "Group By X Order By X",
+    "query": "SELECT user.intcol FROM user GROUP BY user.intcol ORDER BY COUNT(user.intcol)",
+    "v3-plan": "unsupported: in scatter query: complex order by expression: count(`user`.intcol)",
+    "gen4-plan": {
+      "QueryType": "SELECT",
+      "Original": "SELECT user.intcol FROM user GROUP BY user.intcol ORDER BY COUNT(user.intcol)",
+      "Instructions": {
+        "OperatorType": "Sort",
+        "Variant": "Memory",
+        "OrderBy": "1 ASC",
+        "ResultColumns": 1,
+        "Inputs": [
+          {
+            "OperatorType": "Aggregate",
+            "Variant": "Ordered",
+            "Aggregates": "sum_count(1) AS count(`user`.intcol)",
+            "GroupBy": "0",
+            "Inputs": [
+              {
+                "OperatorType": "Route",
+                "Variant": "Scatter",
+                "Keyspace": {
+                  "Name": "user",
+                  "Sharded": true
+                },
+                "FieldQuery": "select `user`.intcol, count(`user`.intcol) from `user` where 1 != 1 group by `user`.intcol",
+                "OrderBy": "0 ASC",
+                "Query": "select `user`.intcol, count(`user`.intcol) from `user` group by `user`.intcol order by `user`.intcol asc",
+                "Table": "`user`"
+              }
+            ]
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
   }
 ]
\ No newline at end of file
diff --git a/go/vt/vtgate/planbuilder/testdata/filter_cases.json b/go/vt/vtgate/planbuilder/testdata/filter_cases.json
index 76f6def665..e9ec88d47c 100644
--- a/go/vt/vtgate/planbuilder/testdata/filter_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/filter_cases.json
@@ -4128,6 +4128,7 @@
     }
   },
   {
+    "comment": "select * from samecolvin where col = :col",
     "query": "select * from samecolvin where col = :col",
     "v3-plan": {
       "QueryType": "SELECT",
@@ -6342,5 +6343,42 @@
         "user.music"
       ]
     }
+  },
+  {
+    "comment": "Self referencing columns in HAVING should work",
+    "query": "select a+2 as a from user having a = 42",
+    "v3-plan": {
+      "QueryType": "SELECT",
+      "Original": "select a+2 as a from user having a = 42",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select a + 2 as a from `user` where 1 != 1",
+        "Query": "select a + 2 as a from `user` having a = 42",
+        "Table": "`user`"
+      }
+    },
+    "gen4-plan": {
+      "QueryType": "SELECT",
+      "Original": "select a+2 as a from user having a = 42",
+      "Instructions": {
+        "OperatorType": "Route",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select a + 2 as a from `user` where 1 != 1",
+        "Query": "select a + 2 as a from `user` where a + 2 = 42",
+        "Table": "`user`"
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
   }
-]
\ No newline at end of file
+]