Gen4: subquery refactor

go/vt/sqlparser/ast.go

@@ -1963,6 +1963,18 @@ type (
 		Name ColIdent
 		Fsp  Expr // fractional seconds precision, integer from 0 to 6
 	}
+
+	// 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
+	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 :(
+		Subquery     SelectStatement
+		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
+	}
 )
 
 // iExpr ensures that only expressions nodes can be assigned to a Expr
@@ -1997,6 +2009,7 @@ func (*ConvertUsingExpr) iExpr()  {}
 func (*MatchExpr) iExpr()         {}
 func (*GroupConcatExpr) iExpr()   {}
 func (*Default) iExpr()           {}
+func (*ExtractedSubquery) iExpr() {}
 
 // Exprs represents a list of value expressions.
 // It's not a valid expression because it's not parenthesized.

go/vt/sqlparser/ast_format.go

@@ -1720,3 +1720,33 @@ func (node *RenameTable) Format(buf *TrackedBuffer) {
 		prefix = ", "
 	}
 }
+
+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 = OrExpressions(hasValue, cmp)
+		}
+		buf.astPrintf(node, "%v", expr)
+	case *Subquery:
+		buf.astPrintf(node, "%v", NewArgument(node.ArgName))
+	}
+}

go/vt/sqlparser/precedence.go

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

go/vt/sqlparser/precedence_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -65,6 +66,28 @@ func TestAndOrPrecedence(t *testing.T) {
 	}
 }
 
+func TestNotInSubqueryPrecedence(t *testing.T) {
+	tree, err := Parse("select * from a where not id in (select 42)")
+	require.NoError(t, err)
+	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.Select,
+		OtherSide:    cmp.Left,
+	}
+	not.Expr = extracted
+
+	output := readable(not)
+	assert.Equal(t, "not (:has_values1 = 1 and id in ::arg1)", output)
+}
+
 func TestPlusStarPrecedence(t *testing.T) {
 	validSQL := []struct {
 		input  string

go/vt/sqlparser/tracked_buffer.go

@@ -169,7 +169,7 @@ func (buf *TrackedBuffer) formatter(node SQLNode) {
 	}
 }
 
-//needParens says if we need a parenthesis
+// needParens says if we need a parenthesis
 // op is the operator we are printing
 // val is the value we are checking if we need parens around or not
 // left let's us know if the value is on the lhs or rhs of the operator

go/vt/vtgate/engine/pullout_subquery.go

@@ -171,9 +171,21 @@ func (ps *PulloutSubquery) execSubquery(vcursor VCursor, bindVars map[string]*qu
 }
 
 func (ps *PulloutSubquery) description() PrimitiveDescription {
+	other := map[string]interface{}{}
+	var pulloutVars []string
+	if ps.HasValues != "" {
+		pulloutVars = append(pulloutVars, ps.HasValues)
+	}
+	if ps.SubqueryResult != "" {
+		pulloutVars = append(pulloutVars, ps.SubqueryResult)
+	}
+	if len(pulloutVars) > 0 {
+		other["PulloutVars"] = pulloutVars
+	}
 	return PrimitiveDescription{
 		OperatorType: "Subquery",
 		Variant:      ps.Opcode.String(),
+		Other:        other,
 	}
 }