gen4: support for union in subquery

go/vt/sqlparser/ast.go

@@ -1966,14 +1966,17 @@ type (
 
 	// ExtractedSubquery is a subquery that has been extracted from the original AST
 	// This is a struct that the parser will never produce - it's written and read by the gen4 planner
+	// CAUTION: you should only change argName and hasValuesArg through the setter methods
 	ExtractedSubquery struct {
 		Original     Expr // original expression that was replaced by this ExtractedSubquery
-		ArgName      string
-		HasValuesArg string
-		OpCode       int // this should really be engine.PulloutOpCode, but we cannot depend on engine :(
+		OpCode       int  // this should really be engine.PulloutOpCode, but we cannot depend on engine :(
 		Subquery     *Subquery
 		OtherSide    Expr // represents the side of the comparison, this field will be nil if Original is not a comparison
 		NeedsRewrite bool // tells whether we need to rewrite this subquery to Original or not
+
+		hasValuesArg string
+		argName      string
+		alternative  Expr // this is what will be used to Format this struct
 	}
 )
 

go/vt/sqlparser/ast_format.go

@@ -1726,31 +1726,5 @@ func (node *RenameTable) Format(buf *TrackedBuffer) {
 // it will be formatted as if the subquery has been extracted, and instead
 // show up like argument comparisons
 func (node *ExtractedSubquery) Format(buf *TrackedBuffer) {
-	switch original := node.Original.(type) {
-	case *ExistsExpr:
-		buf.astPrintf(node, "%v", NewArgument(node.ArgName))
-	case *ComparisonExpr:
-		// other_side = :__sq
-		cmp := &ComparisonExpr{
-			Left:     node.OtherSide,
-			Right:    NewArgument(node.ArgName),
-			Operator: original.Operator,
-		}
-		var expr Expr = cmp
-		switch original.Operator {
-		case InOp:
-			// :__sq_has_values = 1 and other_side in ::__sq
-			cmp.Right = NewListArg(node.ArgName)
-			hasValue := &ComparisonExpr{Left: NewArgument(node.HasValuesArg), Right: NewIntLiteral("1"), Operator: EqualOp}
-			expr = AndExpressions(hasValue, cmp)
-		case NotInOp:
-			// :__sq_has_values = 0 or other_side not in ::__sq
-			cmp.Right = NewListArg(node.ArgName)
-			hasValue := &ComparisonExpr{Left: NewArgument(node.HasValuesArg), Right: NewIntLiteral("0"), Operator: EqualOp}
-			expr = &OrExpr{hasValue, cmp}
-		}
-		buf.astPrintf(node, "%v", expr)
-	case *Subquery:
-		buf.astPrintf(node, "%v", NewArgument(node.ArgName))
-	}
+	node.alternative.Format(buf)
 }

go/vt/sqlparser/ast_funcs.go

@@ -1455,3 +1455,56 @@ func GetAllSelects(selStmt SelectStatement) []*Select {
 	}
 	panic("[BUG]: unknown type for SelectStatement")
 }
+
+// SetArgName sets argument name.
+func (es *ExtractedSubquery) SetArgName(n string) {
+	es.argName = n
+	es.updateAlternative()
+}
+
+// SetHasValuesArg sets has_values argument.
+func (es *ExtractedSubquery) SetHasValuesArg(n string) {
+	es.hasValuesArg = n
+	es.updateAlternative()
+}
+
+// GetArgName returns argument name.
+func (es *ExtractedSubquery) GetArgName() string {
+	return es.argName
+}
+
+// GetHasValuesArg returns has values argument.
+func (es *ExtractedSubquery) GetHasValuesArg() string {
+	return es.hasValuesArg
+
+}
+
+func (es *ExtractedSubquery) updateAlternative() {
+	switch original := es.Original.(type) {
+	case *ExistsExpr:
+		es.alternative = NewArgument(es.argName)
+	case *Subquery:
+		es.alternative = NewArgument(es.argName)
+	case *ComparisonExpr:
+		// other_side = :__sq
+		cmp := &ComparisonExpr{
+			Left:     es.OtherSide,
+			Right:    NewArgument(es.argName),
+			Operator: original.Operator,
+		}
+		var expr Expr = cmp
+		switch original.Operator {
+		case InOp:
+			// :__sq_has_values = 1 and other_side in ::__sq
+			cmp.Right = NewListArg(es.argName)
+			hasValue := &ComparisonExpr{Left: NewArgument(es.hasValuesArg), Right: NewIntLiteral("1"), Operator: EqualOp}
+			expr = AndExpressions(hasValue, cmp)
+		case NotInOp:
+			// :__sq_has_values = 0 or other_side not in ::__sq
+			cmp.Right = NewListArg(es.argName)
+			hasValue := &ComparisonExpr{Left: NewArgument(es.hasValuesArg), Right: NewIntLiteral("0"), Operator: EqualOp}
+			expr = &OrExpr{hasValue, cmp}
+		}
+		es.alternative = expr
+	}
+}

go/vt/sqlparser/precedence.go

@@ -91,7 +91,7 @@ func precedenceFor(in Expr) Precendence {
 	case *IntervalExpr:
 		return P1
 	case *ExtractedSubquery:
-		return precedenceFor(node.Original)
+		return precedenceFor(node.alternative)
 	}
 
 	return Syntactic

go/vt/sqlparser/precedence_test.go

@@ -72,22 +72,43 @@ func TestNotInSubqueryPrecedence(t *testing.T) {
 	not := tree.(*Select).Where.Expr.(*NotExpr)
 	cmp := not.Expr.(*ComparisonExpr)
 	subq := cmp.Right.(*Subquery)
-	fmt.Println(subq)
 
 	extracted := &ExtractedSubquery{
-		Original:     cmp,
-		ArgName:      "arg1",
-		HasValuesArg: "has_values1",
-		OpCode:       1,
-		Subquery:     subq,
-		OtherSide:    cmp.Left,
+		Original:  cmp,
+		OpCode:    1,
+		Subquery:  subq,
+		OtherSide: cmp.Left,
 	}
-	not.Expr = extracted
+	extracted.SetArgName("arg1")
+	extracted.SetHasValuesArg("has_values1")
 
+	not.Expr = extracted
 	output := readable(not)
 	assert.Equal(t, "not (:has_values1 = 1 and id in ::arg1)", output)
 }
 
+func TestSubqueryPrecedence(t *testing.T) {
+	tree, err := Parse("select * from a where id in (select 42) and false")
+	require.NoError(t, err)
+	where := tree.(*Select).Where
+	andExpr := where.Expr.(*AndExpr)
+	cmp := andExpr.Left.(*ComparisonExpr)
+	subq := cmp.Right.(*Subquery)
+
+	extracted := &ExtractedSubquery{
+		Original:  andExpr.Left,
+		OpCode:    1,
+		Subquery:  subq,
+		OtherSide: cmp.Left,
+	}
+	extracted.SetArgName("arg1")
+	extracted.SetHasValuesArg("has_values1")
+
+	andExpr.Left = extracted
+	output := readable(extracted)
+	assert.Equal(t, ":has_values1 = 1 and id in ::arg1", output)
+}
+
 func TestPlusStarPrecedence(t *testing.T) {
 	validSQL := []struct {
 		input  string
@@ -193,7 +214,7 @@ func TestRandom(t *testing.T) {
 	// The purpose of this test is to find discrepancies between Format and parsing. If for example our precedence rules are not consistent between the two, this test should find it.
 	// The idea is to generate random queries, and pass them through the parser and then the unparser, and one more time. The result of the first unparse should be the same as the second result.
 	seed := time.Now().UnixNano()
-	fmt.Println(fmt.Sprintf("seed is %d", seed)) //nolint
+	fmt.Println(fmt.Sprintf("seed is %d", seed)) // nolint
 	g := newGenerator(seed, 5)
 	endBy := time.Now().Add(1 * time.Second)
 

go/vt/vtgate/planbuilder/abstract/concatenate.go

@@ -50,7 +50,7 @@ func (c *Concatenate) PushPredicate(sqlparser.Expr, *semantics.SemTable) error {
 
 // UnsolvedPredicates implements the Operator interface
 func (c *Concatenate) UnsolvedPredicates(*semantics.SemTable) []sqlparser.Expr {
-	panic("implement me")
+	return nil
 }
 
 // CheckValid implements the Operator interface

go/vt/vtgate/planbuilder/abstract/operator.go

@@ -190,17 +190,13 @@ func createOperatorFromSelect(sel *sqlparser.Select, semTable *semantics.SemTabl
 	if len(semTable.SubqueryMap[sel]) > 0 {
 		resultantOp = &SubQuery{}
 		for _, sq := range semTable.SubqueryMap[sel] {
-			subquerySelectStatement, isSel := sq.Subquery.Select.(*sqlparser.Select)
-			if !isSel {
-				return nil, semantics.Gen4NotSupportedF("UNION in subquery")
-			}
-			opInner, err := createOperatorFromSelect(subquerySelectStatement, semTable)
+			opInner, err := CreateOperatorFromAST(sq.Subquery.Select, semTable)
 			if err != nil {
 				return nil, err
 			}
 			resultantOp.Inner = append(resultantOp.Inner, &SubQueryInner{
-				Inner:             opInner,
 				ExtractedSubquery: sq,
+				Inner:             opInner,
 			})
 		}
 	}

go/vt/vtgate/planbuilder/abstract/operator_test.go

@@ -133,8 +133,8 @@ func testString(op Operator) string {
 		var inners []string
 		for _, sqOp := range op.Inner {
 			subquery := fmt.Sprintf("{\n\tType: %s", engine.PulloutOpcode(sqOp.ExtractedSubquery.OpCode).String())
-			if sqOp.ExtractedSubquery.ArgName != "" {
-				subquery += fmt.Sprintf("\n\tArgName: %s", sqOp.ExtractedSubquery.ArgName)
+			if sqOp.ExtractedSubquery.GetArgName() != "" {
+				subquery += fmt.Sprintf("\n\tArgName: %s", sqOp.ExtractedSubquery.GetArgName())
 			}
 			subquery += fmt.Sprintf("\n\tQuery: %s\n}", indent(testString(sqOp.Inner)))
 			subquery = indent(subquery)