rewrite having skip calc expr

ast/rewrite.go

@@ -20,11 +20,12 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
-	"github.com/pingcap/parser/ast"
 	"reflect"
 	"regexp"
 	"strings"
 
+	"github.com/pingcap/parser/ast"
+
 	"github.com/XiaoMi/soar/common"
 
 	"github.com/kr/pretty"
@@ -811,27 +812,43 @@ func (rw *Rewrite) RewriteInsertColumns() *Rewrite {
 
 // RewriteHaving having: 对应CLA.013，使用 WHERE 过滤条件替代 HAVING
 func (rw *Rewrite) RewriteHaving() *Rewrite {
+	calcExprs := getCalcExprFromSelect(rw.Stmt) // 获取语句中所有计算列
 	err := sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
 		switch n := node.(type) {
 		case *sqlparser.Select:
 			if n.Having != nil {
-				if n.Where == nil {
-					// WHERE 条件为空直接用 HAVING 替代 WHERE 即可
-					n.Where = n.Having
-				} else {
-					// WHERE 条件不为空，需要对已有的条件进行括号保护，然后再 AND+HAVING
-					n.Where = &sqlparser.Where{
-						Expr: &sqlparser.AndExpr{
-							Left: &sqlparser.ParenExpr{
-								Expr: n.Where.Expr,
-							},
-							Right: n.Having.Expr,
-						},
+				// 对 having 子句拆分，判断每一个 having 条件是否是计算列
+				// 当该列为计算列时不替换到 where 子句
+				exprs := sqlparser.SplitAndExpression(nil, n.Having.Expr)
+				n.Having = nil
+				for _, havingExpr := range exprs {
+					if colName, isCalc := getColNameFromExpr(havingExpr); isCalc {
+						// 如果该列为计算列，则直接添加到 having 子句
+						n.AddHaving(havingExpr)
+					} else {
+						if _, ok := calcExprs[colName]; ok {
+							n.AddHaving(havingExpr)
+						} else {
+							// WHERE 条件不为空，需要对已有的条件进行括号保护，然后再 AND+HAVING
+							if n.Where != nil {
+								n.Where = &sqlparser.Where{
+									Type: sqlparser.WhereStr,
+									Expr: &sqlparser.AndExpr{
+										Left: &sqlparser.ParenExpr{
+											Expr: n.Where.Expr,
+										},
+										Right: havingExpr,
+									},
+								}
+							} else {
+								n.Where = &sqlparser.Where{
+									Type: sqlparser.WhereStr,
+									Expr: havingExpr,
+								}
+							}
+						}
 					}
 				}
-				// 别忘了重置 HAVING 和 Where.Type
-				n.Where.Type = "where"
-				n.Having = nil
 			}
 		}
 		return true, nil
@@ -841,6 +858,76 @@ func (rw *Rewrite) RewriteHaving() *Rewrite {
 	return rw
 }
 
+// getCalcExprFromSelect 获取 select 语句中的计算列
+func getCalcExprFromSelect(stmt sqlparser.Statement) map[string]string {
+	calcExpr := make(map[string]string)
+	switch v := stmt.(type) {
+	case *sqlparser.Select:
+		for _, selectExpr := range v.SelectExprs {
+			switch ex := selectExpr.(type) {
+			case *sqlparser.AliasedExpr:
+				switch x := ex.Expr.(type) {
+				case *sqlparser.ColName:
+					continue
+				default:
+					if !ex.As.IsEmpty() {
+						calcExpr[ex.As.String()] = sqlparser.String(x)
+					}
+					calcExpr[sqlparser.String(x)] = sqlparser.String(x)
+				}
+			}
+		}
+
+		if v.Having != nil {
+			exprs := sqlparser.SplitAndExpression(nil, v.Having.Expr)
+			for _, expr := range exprs {
+				colName, isAgg := getColNameFromExpr(expr)
+				if isAgg {
+					calcExpr[colName] = colName
+				}
+			}
+		}
+
+		if v.OrderBy != nil {
+			for _, order := range v.OrderBy {
+				exprs := sqlparser.SplitAndExpression(nil, order.Expr)
+				for _, expr := range exprs {
+					colName, isAgg := getColNameFromExpr(expr)
+					if isAgg {
+						calcExpr[colName] = colName
+					}
+				}
+			}
+		}
+	default:
+		return nil
+	}
+
+	return calcExpr
+}
+
+// getColNameFromExpr 从 expr 中获取列名与该列是否为计算列
+func getColNameFromExpr(expr sqlparser.SQLNode) (colName string, isCalc bool) {
+	switch v := expr.(type) {
+	case *sqlparser.ComparisonExpr:
+		return getColNameFromExpr(v.Left)
+	case *sqlparser.RangeCond:
+		return getColNameFromExpr(v.Left)
+	case *sqlparser.FuncExpr:
+		return sqlparser.String(v), true
+	case sqlparser.SelectExprs:
+		return getColNameFromExpr(v[0])
+	case *sqlparser.AliasedExpr:
+		if !v.As.IsEmpty() {
+			return getColNameFromExpr(v.Expr)
+		}
+		return sqlparser.String(v), false
+	case *sqlparser.ColName:
+		return v.Name.String(), false
+	}
+	return "", false
+}
+
 // RewriteAddOrderByNull orderbynull: 对应 CLA.008，GROUP BY 无排序要求时添加 ORDER BY NULL
 func (rw *Rewrite) RewriteAddOrderByNull() *Rewrite {
 	err := sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {

ast/rewrite_test.go

@@ -226,6 +226,18 @@ func TestRewriteHaving(t *testing.T) {
 			"input":  `SELECT state, COUNT(*) FROM Drivers WHERE col =1 or col1 =2 GROUP BY state HAVING state IN ('GA', 'TX') ORDER BY state`,
 			"output": "select state, COUNT(*) from Drivers where (col = 1 or col1 = 2) and state in ('GA', 'TX') group by state order by state asc",
 		},
+		{
+			"input":  `SELECT state, COUNT(*) as cnt FROM Drivers WHERE col = 1 or col1 = 2 GROUP BY state HAVING cnt > 1 and state IN ('GA', 'TX') order by state asc`,
+			"output": `select state, COUNT(*) as cnt from Drivers where (col = 1 or col1 = 2) and state in ('GA', 'TX') group by state having cnt > 1 order by state asc`,
+		},
+		{
+			"input":  `SELECT state, COUNT(*) FROM Drivers WHERE col = 1 or col1 = 2 GROUP BY state HAVING COUNT(*) > 1 and state IN ('GA', 'TX') order by state asc`,
+			"output": `select state, COUNT(*) from Drivers where (col = 1 or col1 = 2) and state in ('GA', 'TX') group by state having COUNT(*) > 1 order by state asc`,
+		},
+		{
+			"input":  `SELECT state, (col - 40) as col, COUNT(*) as cnt from Drivers GROUP BY state HAVING col > 1 AND cnt > 1 AND state IN ('GA', 'TX') `,
+			"output": `select state, (col - 40) as col, COUNT(*) as cnt from Drivers where state in ('GA', 'TX') group by state having col > 1 and cnt > 1`,
+		},
 	}
 	for _, sql := range testSQL {
 		rw := NewRewrite(sql["input"]).RewriteHaving()

go.mod

@@ -3,66 +3,28 @@ module github.com/XiaoMi/soar
 go 1.15
 
 require (
-	github.com/Azure/azure-storage-blob-go v0.10.0 // indirect
-	github.com/Azure/go-autorest/autorest v0.10.0 // indirect
 	github.com/CorgiMan/json2 v0.0.0-20150213135156-e72957aba209
-	github.com/aquarapid/vaultlib v0.5.1 // indirect
 	github.com/astaxie/beego v1.12.3
-	github.com/buger/jsonparser v0.0.0-20200322175846-f7e751efca13 // indirect
-	github.com/codegangsta/inject v0.0.0-20150114235600-33e0aa1cb7c0 // indirect
-	github.com/cyberdelia/go-metrics-graphite v0.0.0-20161219230853-39f87cc3b432 // indirect
 	github.com/dchest/uniuri v0.0.0-20200228104902-7aecb25e1fe5
 	github.com/gedex/inflector v0.0.0-20170307190818-16278e9db813
-	github.com/go-martini/martini v0.0.0-20170121215854-22fa46961aab // indirect
 	github.com/go-sql-driver/mysql v1.6.0
 	github.com/golang/protobuf v1.5.2 // indirect
-	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
 	github.com/gorilla/handlers v1.5.1 // indirect
 	github.com/gorilla/mux v1.8.0 // indirect
-	github.com/hashicorp/consul/api v1.5.0 // indirect
-	github.com/hashicorp/go-immutable-radix v1.1.0 // indirect
-	github.com/hashicorp/go-sockaddr v1.0.2 // indirect
-	github.com/hashicorp/go-uuid v1.0.2 // indirect
-	github.com/hashicorp/serf v0.9.2 // indirect
-	github.com/howeyc/gopass v0.0.0-20190910152052-7cb4b85ec19c // indirect
-	github.com/jeremywohl/flatten v0.0.0-20190921043622-d936035e55cf // indirect
-	github.com/klauspost/pgzip v1.2.4 // indirect
-	github.com/konsorten/go-windows-terminal-sequences v1.0.3 // indirect
 	github.com/kr/pretty v0.2.1
-	github.com/looplab/fsm v0.2.0 // indirect
-	github.com/martini-contrib/auth v0.0.0-20150219114609-fa62c19b7ae8 // indirect
-	github.com/martini-contrib/gzip v0.0.0-20151124214156-6c035326b43f // indirect
-	github.com/martini-contrib/render v0.0.0-20150707142108-ec18f8345a11 // indirect
-	github.com/mitchellh/go-ps v1.0.0 // indirect
-	github.com/mitchellh/go-testing-interface v1.14.0 // indirect
-	github.com/mitchellh/mapstructure v1.2.3 // indirect
 	github.com/montanaflynn/stats v0.6.3 // indirect
-	github.com/olekukonko/tablewriter v0.0.5-0.20200416053754-163badb3bac6 // indirect
-	github.com/oxtoacart/bpool v0.0.0-20190530202638-03653db5a59c // indirect
-	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/percona/go-mysql v0.0.0-20210427141028-73d29c6da78c
-	github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4 // indirect
 	github.com/pingcap/parser v0.0.0-20210525032559-c37778aff307
-	github.com/pingcap/pd/v4 v4.0.0-beta.1.0.20200305072537-61d9f9cc35d3 // indirect
 	github.com/pingcap/tidb v1.1.0-beta.0.20210601085537-5d7c852770eb
 	github.com/pingcap/tipb v0.0.0-20210601083426-79a378b6d1c4 // indirect
-	github.com/planetscale/pargzip v0.0.0-20201116224723-90c7fc03ea8a // indirect
-	github.com/rcrowley/go-metrics v0.0.0-20200313005456-10cdbea86bc0 // indirect
 	github.com/russross/blackfriday v1.6.0
 	github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca
-	github.com/samuel/go-zookeeper v0.0.0-20200724154423-2164a8ac840e // indirect
 	github.com/sirupsen/logrus v1.8.1 // indirect
-	github.com/sjmudd/stopwatch v0.0.0-20170613150411-f380bf8a9be1 // indirect
 	github.com/spf13/cobra v1.1.1 // indirect
-	github.com/spyzhov/ajson v0.4.2 // indirect
 	github.com/tidwall/gjson v1.7.5
-	go.uber.org/multierr v1.7.0 // indirect
 	go.uber.org/zap v1.17.0 // indirect
 	golang.org/x/lint v0.0.0-20200302205851-738671d3881b // indirect
 	golang.org/x/sys v0.0.0-20210601080250-7ecdf8ef093b // indirect
-	golang.org/x/text v0.3.6 // indirect
-	gopkg.in/gcfg.v1 v1.2.3 // indirect
-	gopkg.in/warnings.v0 v0.1.2 // indirect
 	gopkg.in/yaml.v2 v2.4.0
 	vitess.io/vitess v0.0.0-20200325000816-eda961851d63
 )