addressing review comments:

- increased test coverage
 - revert route to pre-union-shenanigan state

data/test/vtgate/select_cases.txt

@@ -544,10 +544,6 @@
   }
 }
 
-# Union all
-"select col1, col2 from user union all select col1, col2 from user_extra"
-"unsupported: scatter subquery"
-
 "select * from (select col1, col2 from unsharded where id = 1 union select col1, col2 from unsharded where id = 3) a"
 {
   "Original": "select * from (select col1, col2 from unsharded where id = 1 union select col1, col2 from unsharded where id = 3) a",
@@ -618,18 +614,6 @@
   }
 }
 
-"(select id from user union select id from music) union select 1 from dual"
-"unsupported: scatter subquery"
-
-"select 1 from music union (select id from user union all select name from unsharded)"
-"unsupported: subquery keyspace different from outer query"
-
-"select 1 from music union (select id from user union select name from unsharded)"
-"unsupported: subquery keyspace different from outer query"
-
-"select id from user union all select id from music"
-"unsupported: scatter subquery"
-
 "(select id from unsharded order by id asc limit 1) union (select id from unsharded order by id desc limit 1) order by id asc limit 1"
 {
   "Original": "(select id from unsharded order by id asc limit 1) union (select id from unsharded order by id desc limit 1) order by id asc limit 1",

data/test/vtgate/unsupported_cases.txt

@@ -235,7 +235,7 @@
 "insert into unsharded_auto values(1,1)"
 "column list required for tables with auto-inc columns"
 
-# unsharded insert, col list doesn't match values
+# unsharded insert, col list does not match values
 "insert into unsharded_auto(id, val) values(1)"
 "column list doesn't match values"
 
@@ -258,3 +258,29 @@
 # complex expression in parenthesis with order by not supported yet
 "select * from user where (id = 4 AND name ='abc') order by id"
 "unsupported: scatter and order by"
+
+# multi-shard union
+"(select id from user union select id from music) union select 1 from dual"
+"unsupported: scatter subquery"
+
+# multi-shard union
+"select 1 from music union (select id from user union all select name from unsharded)"
+"unsupported: subquery keyspace different from outer query"
+
+# multi-shard union
+"select 1 from music union (select id from user union select name from unsharded)"
+"unsupported: subquery keyspace different from outer query"
+
+# multi-shard union
+"select id from user union all select id from music"
+"unsupported: scatter subquery"
+
+# Union all
+"select col1, col2 from user union all select col1, col2 from user_extra"
+"unsupported: scatter subquery"
+
+"(select user.id, user.name from user join user_extra where user_extra.extra = 'asdf') union select 'b','c' from user"
+"unsupported construct: SELECT of UNION is non-trivial"
+
+"select 'b','c' from user union (select user.id, user.name from user join user_extra where user_extra.extra = 'asdf')"
+"unsupported construct: SELECT of UNION is non-trivial"

go/vt/sqlparser/ast.go

@@ -254,12 +254,12 @@ type ParenSelect struct {
 
 // AddOrder adds an order by element
 func (node *ParenSelect) AddOrder(order *Order) {
-	node.Select.AddOrder(order)
+	panic("unreachable")
 }
 
 // SetLimit sets the limit clause
 func (node *ParenSelect) SetLimit(limit *Limit) {
-	node.Select.SetLimit(limit)
+	panic("unreachable")
 }
 
 // Format formats the node.

go/vt/sqlparser/ast_test.go

@@ -95,6 +95,66 @@ func TestSelect(t *testing.T) {
 	}
 }
 
+func TestAddOrder(t *testing.T) {
+	src, err := Parse("select foo, bar from baz order by foo")
+	if err != nil {
+		t.Error(err)
+	}
+	order := src.(*Select).OrderBy[0]
+	dst, err := Parse("select * from t")
+	if err != nil {
+		t.Error(err)
+	}
+	dst.(*Select).AddOrder(order)
+	buf := NewTrackedBuffer(nil)
+	dst.Format(buf)
+	want := "select * from t order by foo asc"
+	if buf.String() != want {
+		t.Errorf("order: %q, want %s", buf.String(), want)
+	}
+	dst, err = Parse("select * from t union select * from s")
+	if err != nil {
+		t.Error(err)
+	}
+	dst.(*Union).AddOrder(order)
+	buf = NewTrackedBuffer(nil)
+	dst.Format(buf)
+	want = "select * from t union select * from s order by foo asc"
+	if buf.String() != want {
+		t.Errorf("order: %q, want %s", buf.String(), want)
+	}
+}
+
+func TestSetLimit(t *testing.T) {
+	src, err := Parse("select foo, bar from baz limit 4")
+	if err != nil {
+		t.Error(err)
+	}
+	limit := src.(*Select).Limit
+	dst, err := Parse("select * from t")
+	if err != nil {
+		t.Error(err)
+	}
+	dst.(*Select).SetLimit(limit)
+	buf := NewTrackedBuffer(nil)
+	dst.Format(buf)
+	want := "select * from t limit 4"
+	if buf.String() != want {
+		t.Errorf("limit: %q, want %s", buf.String(), want)
+	}
+	dst, err = Parse("select * from t union select * from s")
+	if err != nil {
+		t.Error(err)
+	}
+	dst.(*Union).SetLimit(limit)
+	buf = NewTrackedBuffer(nil)
+	dst.Format(buf)
+	want = "select * from t union select * from s limit 4"
+	if buf.String() != want {
+		t.Errorf("order: %q, want %s", buf.String(), want)
+	}
+}
+
 func TestWhere(t *testing.T) {
 	var w *Where
 	buf := NewTrackedBuffer(nil)

go/vt/vtgate/engine/route.go

@@ -243,38 +243,25 @@ func newScatterParams(ks string, bv map[string]interface{}, shards []string) *sc
 	}
 }
 
-func skewCopyQueryConstruct(queryConstruct *queryinfo.QueryConstruct, joinvars map[string]interface{}) *queryinfo.QueryConstruct {
-	qc := &queryinfo.QueryConstruct{
-		SQL:              queryConstruct.SQL,
-		Comments:         queryConstruct.Comments,
-		Keyspace:         queryConstruct.Keyspace,
-		BindVars:         make(map[string]interface{}),
-		NotInTransaction: queryConstruct.NotInTransaction,
-	}
-	for k, v := range queryConstruct.BindVars {
-		qc.BindVars[k] = v
-	}
-	for k, v := range joinvars {
-		qc.BindVars[k] = v
-	}
-	return qc
-}
-
 // Execute performs a non-streaming exec.
 func (route *Route) Execute(vcursor VCursor, queryConstruct *queryinfo.QueryConstruct, joinvars map[string]interface{}, wantfields bool) (*sqltypes.Result, error) {
-	qc := skewCopyQueryConstruct(queryConstruct, joinvars)
+	saved := copyBindVars(queryConstruct.BindVars)
+	defer func() { queryConstruct.BindVars = saved }()
+	for k, v := range joinvars {
+		queryConstruct.BindVars[k] = v
+	}
 
 	switch route.Opcode {
 	case UpdateEqual:
-		return route.execUpdateEqual(vcursor, qc)
+		return route.execUpdateEqual(vcursor, queryConstruct)
 	case DeleteEqual:
-		return route.execDeleteEqual(vcursor, qc)
+		return route.execDeleteEqual(vcursor, queryConstruct)
 	case InsertSharded:
-		return route.execInsertSharded(vcursor, qc)
+		return route.execInsertSharded(vcursor, queryConstruct)
 	case InsertUnsharded:
-		return route.execInsertUnsharded(vcursor, qc)
+		return route.execInsertUnsharded(vcursor, queryConstruct)
 	case Show:
-		return route.execShow(vcursor, qc)
+		return route.execShow(vcursor, queryConstruct)
 	}
 
 	var err error
@@ -282,13 +269,13 @@ func (route *Route) Execute(vcursor VCursor, queryConstruct *queryinfo.QueryCons
 	switch route.Opcode {
 	case SelectUnsharded, UpdateUnsharded,
 		DeleteUnsharded:
-		params, err = route.paramsUnsharded(vcursor, qc)
+		params, err = route.paramsUnsharded(vcursor, queryConstruct)
 	case SelectEqual, SelectEqualUnique:
-		params, err = route.paramsSelectEqual(vcursor, qc)
+		params, err = route.paramsSelectEqual(vcursor, queryConstruct)
 	case SelectIN:
-		params, err = route.paramsSelectIN(vcursor, qc)
+		params, err = route.paramsSelectIN(vcursor, queryConstruct)
 	case SelectScatter:
-		params, err = route.paramsSelectScatter(vcursor, qc)
+		params, err = route.paramsSelectScatter(vcursor, queryConstruct)
 	default:
 		// TODO(sougou): improve error.
 		return nil, fmt.Errorf("unsupported query route: %v", route)
@@ -297,25 +284,29 @@ func (route *Route) Execute(vcursor VCursor, queryConstruct *queryinfo.QueryCons
 		return nil, err
 	}
 
-	shardQueries := route.getShardQueries(route.Query+qc.Comments, params)
-	return vcursor.ExecuteMultiShard(params.ks, shardQueries, qc.NotInTransaction)
+	shardQueries := route.getShardQueries(route.Query+queryConstruct.Comments, params)
+	return vcursor.ExecuteMultiShard(params.ks, shardQueries, queryConstruct.NotInTransaction)
 }
 
 // StreamExecute performs a streaming exec.
 func (route *Route) StreamExecute(vcursor VCursor, queryConstruct *queryinfo.QueryConstruct, joinvars map[string]interface{}, wantfields bool, callback func(*sqltypes.Result) error) error {
-	qc := skewCopyQueryConstruct(queryConstruct, joinvars)
+	saved := copyBindVars(queryConstruct.BindVars)
+	defer func() { queryConstruct.BindVars = saved }()
+	for k, v := range joinvars {
+		queryConstruct.BindVars[k] = v
+	}
 
 	var err error
 	var params *scatterParams
 	switch route.Opcode {
 	case SelectUnsharded:
-		params, err = route.paramsUnsharded(vcursor, qc)
+		params, err = route.paramsUnsharded(vcursor, queryConstruct)
 	case SelectEqual, SelectEqualUnique:
-		params, err = route.paramsSelectEqual(vcursor, qc)
+		params, err = route.paramsSelectEqual(vcursor, queryConstruct)
 	case SelectIN:
-		params, err = route.paramsSelectIN(vcursor, qc)
+		params, err = route.paramsSelectIN(vcursor, queryConstruct)
 	case SelectScatter:
-		params, err = route.paramsSelectScatter(vcursor, qc)
+		params, err = route.paramsSelectScatter(vcursor, queryConstruct)
 	default:
 		return fmt.Errorf("query %q cannot be used for streaming", route.Query)
 	}
@@ -323,7 +314,7 @@ func (route *Route) StreamExecute(vcursor VCursor, queryConstruct *queryinfo.Que
 		return err
 	}
 	return vcursor.StreamExecuteMulti(
-		route.Query+qc.Comments,
+		route.Query+queryConstruct.Comments,
 		params.ks,
 		params.shardVars,
 		callback,
@@ -332,13 +323,25 @@ func (route *Route) StreamExecute(vcursor VCursor, queryConstruct *queryinfo.Que
 
 // GetFields fetches the field info.
 func (route *Route) GetFields(vcursor VCursor, queryConstruct *queryinfo.QueryConstruct, joinvars map[string]interface{}) (*sqltypes.Result, error) {
-	qc := skewCopyQueryConstruct(queryConstruct, joinvars)
+	saved := copyBindVars(queryConstruct.BindVars)
+	defer func() { queryConstruct.BindVars = saved }()
+	for k := range joinvars {
+		queryConstruct.BindVars[k] = nil
+	}
 	ks, shard, err := vcursor.GetAnyShard(route.Keyspace.Name)
 	if err != nil {
 		return nil, err
 	}
 
-	return vcursor.ScatterConnExecute(route.FieldQuery, qc.BindVars, ks, []string{shard}, qc.NotInTransaction)
+	return vcursor.ScatterConnExecute(route.FieldQuery, queryConstruct.BindVars, ks, []string{shard}, queryConstruct.NotInTransaction)
+}
+
+func copyBindVars(bindVars map[string]interface{}) map[string]interface{} {
+	out := make(map[string]interface{})
+	for k, v := range bindVars {
+		out[k] = v
+	}
+	return out
 }
 
 func (route *Route) paramsUnsharded(vcursor VCursor, queryConstruct *queryinfo.QueryConstruct) (*scatterParams, error) {

go/vt/vtgate/planbuilder/expr.go

@@ -93,7 +93,7 @@ func findRoute(expr sqlparser.Expr, bldr builder) (rb *route, err error) {
 		return nil, err
 	}
 	for _, subroute := range subroutes {
-		err = subqueryCanMerge(highestRoute, subroute)
+		err = routesCanMerge(highestRoute, subroute)
 		if err != nil {
 			return nil, err
 		}
@@ -104,10 +104,10 @@ func findRoute(expr sqlparser.Expr, bldr builder) (rb *route, err error) {
 	return highestRoute, nil
 }
 
-// subqueryCanMerge returns nil if the inner subquery
+// routesCanMerge returns nil if the inner subquery
 // can be merged with the specified outer route. If it
 // cannot, then it returns an appropriate error.
-func subqueryCanMerge(outer, inner *route) error {
+func routesCanMerge(outer, inner *route) error {
 	if outer.ERoute.Keyspace.Name != inner.ERoute.Keyspace.Name {
 		return errors.New("unsupported: subquery keyspace different from outer query")
 	}

go/vt/vtgate/planbuilder/union.go

@@ -71,13 +71,13 @@ func processPart(part sqlparser.SelectStatement, vschema VSchema, outer builder)
 func unionRouteMerge(union *sqlparser.Union, left, right builder, vschema VSchema) (builder, error) {
 	lroute, ok := left.(*route)
 	if !ok {
-		return nil, errors.New("unused message")
+		return nil, errors.New("unsupported construct: SELECT of UNION is non-trivial")
 	}
 	rroute, ok := right.(*route)
 	if !ok {
-		return nil, errors.New("unused message")
+		return nil, errors.New("unsupported construct: SELECT of UNION is non-trivial")
 	}
-	if err := subqueryCanMerge(lroute, rroute); err != nil {
+	if err := routesCanMerge(lroute, rroute); err != nil {
 		return nil, err
 	}
 	table := &vindexes.Table{