vtexplain support select distinct and other expressions

go/vt/vtexplain/testdata/multi-output/selectsharded-output.txt

@@ -126,3 +126,43 @@ select * from name_info order by info /* select * and order by varchar column */
 1 ks_sharded/c0-: select name, info, weight_string(info) from name_info order by info asc limit 10001 /* select * and order by varchar column */
 
 ----------------------------------------------------------------------
+select distinct(name) from user where id = 1 /* select distinct */
+
+1 ks_sharded/-40: select distinct (name) from user where id = 1 limit 10001 /* select distinct */
+
+----------------------------------------------------------------------
+select distinct name from user where id = 1 /* select distinct */
+
+1 ks_sharded/-40: select distinct name from user where id = 1 limit 10001 /* select distinct */
+
+----------------------------------------------------------------------
+select id, substring(name, 1, -1) from user where id = 123 /* select substring */
+
+1 ks_sharded/-40: select id, substr(name, 1, -1) from user where id = 123 limit 10001 /* select substring */
+
+----------------------------------------------------------------------
+select id, substring_index(name, '123456', -1) from user where id = 123 /* select substring_index */
+
+1 ks_sharded/-40: select id, substring_index(name, '123456', -1) from user where id = 123 limit 10001 /* select substring_index */
+
+----------------------------------------------------------------------
+select id, case when name = 'alice' then 'ALICE' when name = 'bob' then 'BOB' end as name from user where id = 1 /* select case */
+
+1 ks_sharded/-40: select id, case when name = 'alice' then 'ALICE' when name = 'bob' then 'BOB' end as name from user where id = 1 limit 10001 /* select case */
+
+----------------------------------------------------------------------
+select id, case when name = 'alice' then 'ALICE' when name = 'bob' then 'BOB' else 'OTHER' end as name from user where id = 1 /* select case */
+
+1 ks_sharded/-40: select id, case when name = 'alice' then 'ALICE' when name = 'bob' then 'BOB' else 'OTHER' end as name from user where id = 1 limit 10001 /* select case */
+
+----------------------------------------------------------------------
+select id, case when substr(name, 1, 5) = 'alice' then 'ALICE' when name = 'bob' then 'BOB' else 'OTHER' end as name from user where id = 1 /* select case */
+
+1 ks_sharded/-40: select id, case when substr(name, 1, 5) = 'alice' then 'ALICE' when name = 'bob' then 'BOB' else 'OTHER' end as name from user where id = 1 limit 10001 /* select case */
+
+----------------------------------------------------------------------
+select id, 'abc' as test from user where id = 1 union all select id, 'def' as test from user where id = 1 union all select id, 'ghi' as test from user where id = 1 /* union all */
+
+1 ks_sharded/-40: select id, 'abc' as test from user where id = 1 union all select id, 'def' as test from user where id = 1 union all select id, 'ghi' as test from user where id = 1 limit 10001 /* union all */
+
+----------------------------------------------------------------------

go/vt/vtexplain/testdata/selectsharded-queries.sql

@@ -19,4 +19,16 @@ select name from user where id in (select id from t1) /* non-correlated subquery
 select name from user where id not in (select id from t1) /* non-correlated subquery in NOT IN clause */;
 select name from user where exists (select id from t1) /* non-correlated subquery as EXISTS */;
 
-select * from name_info order by info /* select * and order by varchar column */
+select * from name_info order by info /* select * and order by varchar column */;
+
+select distinct(name) from user where id = 1 /* select distinct */;
+select distinct name from user where id = 1 /* select distinct */;
+
+select id, substring(name, 1, -1) from user where id = 123 /* select substring */;
+select id, substring_index(name, '123456', -1) from user where id = 123 /* select substring_index */;
+
+select id, case when name = 'alice' then 'ALICE' when name = 'bob' then 'BOB' end as name from user where id = 1 /* select case */;
+select id, case when name = 'alice' then 'ALICE' when name = 'bob' then 'BOB' else 'OTHER' end as name from user where id = 1 /* select case */;
+select id, case when substr(name, 1, 5) = 'alice' then 'ALICE' when name = 'bob' then 'BOB' else 'OTHER' end as name from user where id = 1 /* select case */;
+
+select id, 'abc' as test from user where id = 1 union all select id, 'def' as test from user where id = 1 union all select id, 'ghi' as test from user where id = 1 /* union all */;

go/vt/vtexplain/vtexplain_vttablet.go

@@ -19,6 +19,7 @@ package vtexplain
 import (
 	"encoding/json"
 	"fmt"
+	"reflect"
 	"strings"
 	"sync"
 
@@ -501,7 +502,15 @@ func (t *explainTablet) HandleQuery(c *mysql.Conn, query string, callback func(*
 			return err
 		}
 
-		selStmt := stmt.(*sqlparser.Select)
+		var selStmt *sqlparser.Select
+		switch stmt.(type) {
+		case *sqlparser.Select:
+			selStmt = stmt.(*sqlparser.Select)
+		case *sqlparser.Union:
+			selStmt = stmt.(*sqlparser.Union).Right.(*sqlparser.Select)
+		default:
+			return fmt.Errorf("vtexplain: unsupported statement type +%v", reflect.TypeOf(stmt))
+		}
 
 		if len(selStmt.From) != 1 {
 			return fmt.Errorf("unsupported select with multiple from clauses")
@@ -531,40 +540,7 @@ func (t *explainTablet) HandleQuery(c *mysql.Conn, query string, callback func(*
 		for _, node := range selStmt.SelectExprs {
 			switch node := node.(type) {
 			case *sqlparser.AliasedExpr:
-				switch node := node.Expr.(type) {
-				case *sqlparser.ColName:
-					col := strings.ToLower(node.Name.String())
-					colType := colTypeMap[col]
-					if colType == querypb.Type_NULL_TYPE {
-						return fmt.Errorf("invalid column %s", col)
-					}
-					colNames = append(colNames, col)
-					colTypes = append(colTypes, colType)
-				case *sqlparser.FuncExpr:
-					// As a shortcut, functions are integral types
-					colNames = append(colNames, sqlparser.String(node))
-					colTypes = append(colTypes, querypb.Type_INT32)
-				case *sqlparser.SQLVal:
-					colNames = append(colNames, sqlparser.String(node))
-					switch node.Type {
-					case sqlparser.IntVal:
-						fallthrough
-					case sqlparser.HexNum:
-						fallthrough
-					case sqlparser.HexVal:
-						fallthrough
-					case sqlparser.BitVal:
-						colTypes = append(colTypes, querypb.Type_INT32)
-					case sqlparser.StrVal:
-						colTypes = append(colTypes, querypb.Type_VARCHAR)
-					case sqlparser.FloatVal:
-						colTypes = append(colTypes, querypb.Type_FLOAT64)
-					default:
-						return fmt.Errorf("unsupported sql value %s", sqlparser.String(node))
-					}
-				default:
-					return fmt.Errorf("unsupported select expression %s", sqlparser.String(node))
-				}
+				colNames, colTypes = inferColTypeFromExpr(node.Expr, colTypeMap, colNames, colTypes)
 			case *sqlparser.StarExpr:
 				for col, colType := range colTypeMap {
 					colNames = append(colNames, col)
@@ -615,3 +591,46 @@ func (t *explainTablet) HandleQuery(c *mysql.Conn, query string, callback func(*
 
 	return callback(result)
 }
+
+func inferColTypeFromExpr(node sqlparser.Expr, colTypeMap map[string]querypb.Type, colNames []string, colTypes []querypb.Type) ([]string, []querypb.Type) {
+	switch node := node.(type) {
+	case *sqlparser.ColName:
+		col := strings.ToLower(node.Name.String())
+		colType := colTypeMap[col]
+		if colType == querypb.Type_NULL_TYPE {
+			log.Errorf("vtexplain: invalid column %s, typeMap +%v", col, colTypeMap)
+		}
+		colNames = append(colNames, col)
+		colTypes = append(colTypes, colType)
+	case *sqlparser.FuncExpr:
+		// As a shortcut, functions are integral types
+		colNames = append(colNames, sqlparser.String(node))
+		colTypes = append(colTypes, querypb.Type_INT32)
+	case *sqlparser.SQLVal:
+		colNames = append(colNames, sqlparser.String(node))
+		switch node.Type {
+		case sqlparser.IntVal:
+			fallthrough
+		case sqlparser.HexNum:
+			fallthrough
+		case sqlparser.HexVal:
+			fallthrough
+		case sqlparser.BitVal:
+			colTypes = append(colTypes, querypb.Type_INT32)
+		case sqlparser.StrVal:
+			colTypes = append(colTypes, querypb.Type_VARCHAR)
+		case sqlparser.FloatVal:
+			colTypes = append(colTypes, querypb.Type_FLOAT64)
+		default:
+			log.Errorf("vtexplain: unsupported sql value %s", sqlparser.String(node))
+		}
+	case *sqlparser.ParenExpr:
+		colNames, colTypes = inferColTypeFromExpr(node.Expr, colTypeMap, colNames, colTypes)
+	case *sqlparser.CaseExpr:
+		colNames, colTypes = inferColTypeFromExpr(node.Whens[0].Val, colTypeMap, colNames, colTypes)
+	default:
+		log.Errorf("vtexplain: unsupported select expression type +%v node %s", reflect.TypeOf(node), sqlparser.String(node))
+	}
+
+	return colNames, colTypes
+}