Merge pull request #8382 from planetscale/gen4-fail-more2

Gen4 fail more2

go/vt/sqlparser/ast.go

@@ -201,16 +201,17 @@ type (
 		Distinct         bool
 		StraightJoinHint bool
 		SQLCalcFoundRows bool
-		From             []TableExpr
-		Comments         Comments
-		SelectExprs      SelectExprs
-		Where            *Where
-		GroupBy          GroupBy
-		Having           *Where
-		OrderBy          OrderBy
-		Limit            *Limit
-		Lock             Lock
-		Into             *SelectInto
+		// The From field must be the first AST element of this struct so the rewriter sees it first
+		From        []TableExpr
+		Comments    Comments
+		SelectExprs SelectExprs
+		Where       *Where
+		GroupBy     GroupBy
+		Having      *Where
+		OrderBy     OrderBy
+		Limit       *Limit
+		Lock        Lock
+		Into        *SelectInto
 	}
 
 	// SelectInto is a struct that represent the INTO part of a select query

go/vt/vtgate/planbuilder/join.go

@@ -17,8 +17,6 @@ limitations under the License.
 package planbuilder
 
 import (
-	"errors"
-
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
@@ -100,7 +98,7 @@ func newJoin(lpb, rpb *primitiveBuilder, ajoin *sqlparser.JoinTableExpr, reserve
 				return err
 			}
 		case ajoin.Condition.Using != nil:
-			return errors.New("unsupported: join with USING(column_list) clause for complex queries")
+			return vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: join with USING(column_list) clause for complex queries")
 		}
 	}
 	lpb.plan = &join{

go/vt/vtgate/planbuilder/jointree.go

@@ -141,20 +141,31 @@ func (p parenTables) tableID() semantics.TableSet {
 }
 
 // visit will traverse the route tables, going inside parenTables and visiting all routeTables
-func (p parenTables) visit(f func(tbl *routeTable) error) error {
-	for _, r := range p {
-		switch r := r.(type) {
-		case *routeTable:
-			err := f(r)
-			if err != nil {
-				return err
-			}
-		case parenTables:
-			err := r.visit(f)
+func visitTables(r relation, f func(tbl *routeTable) error) error {
+	switch r := r.(type) {
+	case *routeTable:
+		err := f(r)
+		if err != nil {
+			return err
+		}
+	case parenTables:
+		for _, r := range r {
+			err := visitTables(r, f)
 			if err != nil {
 				return err
 			}
 		}
+		return nil
+	case *leJoin:
+		err := visitTables(r.lhs, f)
+		if err != nil {
+			return err
+		}
+		err = visitTables(r.rhs, f)
+		if err != nil {
+			return err
+		}
+		return nil
 	}
 	return nil
 }
@@ -355,11 +366,7 @@ func (rp *routePlan) planEqualOp(node *sqlparser.ComparisonExpr) (bool, error) {
 	return rp.haveMatchingVindex(node, column, *val, equalOrEqualUnique, justTheVindex), err
 }
 
-func (rp *routePlan) planInOp(node *sqlparser.ComparisonExpr) (bool, error) {
-	column, ok := node.Left.(*sqlparser.ColName)
-	if !ok {
-		return false, nil
-	}
+func (rp *routePlan) planSimpleInOp(node *sqlparser.ComparisonExpr, left *sqlparser.ColName) (bool, error) {
 	value, err := sqlparser.NewPlanValue(node.Right)
 	if err != nil {
 		// if we are unable to create a PlanValue, we can't use a vindex, but we don't have to fail
@@ -377,7 +384,70 @@ func (rp *routePlan) planInOp(node *sqlparser.ComparisonExpr) (bool, error) {
 		}
 	}
 	opcode := func(*vindexes.ColumnVindex) engine.RouteOpcode { return engine.SelectIN }
-	return rp.haveMatchingVindex(node, column, value, opcode, justTheVindex), err
+	return rp.haveMatchingVindex(node, left, value, opcode, justTheVindex), err
+}
+
+func (rp *routePlan) planCompositeInOp(node *sqlparser.ComparisonExpr, left sqlparser.ValTuple) (bool, error) {
+	right, rightIsValTuple := node.Right.(sqlparser.ValTuple)
+	if !rightIsValTuple {
+		return false, nil
+	}
+	return rp.planCompositeInOpRecursive(node, left, right, nil)
+}
+
+func (rp *routePlan) planCompositeInOpRecursive(node *sqlparser.ComparisonExpr, left, right sqlparser.ValTuple, coordinates []int) (bool, error) {
+	foundVindex := false
+	cindex := len(coordinates)
+	coordinates = append(coordinates, 0)
+	for i, expr := range left {
+		coordinates[cindex] = i
+		switch expr := expr.(type) {
+		case sqlparser.ValTuple:
+			ok, err := rp.planCompositeInOpRecursive(node, expr, right, coordinates)
+			if err != nil {
+				return false, err
+			}
+			return ok || foundVindex, nil
+		case *sqlparser.ColName:
+			// check if left col is a vindex
+			if !rp.hasVindex(expr) {
+				continue
+			}
+
+			rightVals := make(sqlparser.ValTuple, len(right))
+			for j, currRight := range right {
+				switch currRight := currRight.(type) {
+				case sqlparser.ValTuple:
+					val := tupleAccess(currRight, coordinates)
+					if val == nil {
+						return false, nil
+					}
+					rightVals[j] = val
+				default:
+					return false, nil
+				}
+			}
+			newPlanValues, err := makePlanValue(rightVals)
+			if newPlanValues == nil || err != nil {
+				return false, err
+			}
+
+			opcode := func(*vindexes.ColumnVindex) engine.RouteOpcode { return engine.SelectMultiEqual }
+			newVindex := rp.haveMatchingVindex(node, expr, *newPlanValues, opcode, justTheVindex)
+			foundVindex = newVindex || foundVindex
+		}
+	}
+	return foundVindex, nil
+}
+
+func (rp *routePlan) planInOp(node *sqlparser.ComparisonExpr) (bool, error) {
+	switch left := node.Left.(type) {
+	case *sqlparser.ColName:
+		return rp.planSimpleInOp(node, left)
+	case sqlparser.ValTuple:
+		return rp.planCompositeInOp(node, left)
+	}
+	return false, nil
 }
 
 func (rp *routePlan) planLikeOp(node *sqlparser.ComparisonExpr) (bool, error) {
@@ -434,6 +504,17 @@ func makePlanValue(n sqlparser.Expr) (*sqltypes.PlanValue, error) {
 	return &value, nil
 }
 
+func (rp routePlan) hasVindex(column *sqlparser.ColName) bool {
+	for _, v := range rp.vindexPreds {
+		for _, col := range v.colVindex.Columns {
+			if column.Name.Equal(col) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func (rp *routePlan) haveMatchingVindex(
 	node sqlparser.Expr,
 	column *sqlparser.ColName,
@@ -482,7 +563,12 @@ func (rp *routePlan) pickBestAvailableVindex() {
 
 // Predicates takes all known predicates for this route and ANDs them together
 func (rp *routePlan) Predicates() sqlparser.Expr {
-	return sqlparser.AndExpressions(rp.predicates...)
+	predicates := rp.predicates
+	_ = visitTables(rp.tables, func(tbl *routeTable) error {
+		predicates = append(predicates, tbl.qtable.Predicates...)
+		return nil
+	})
+	return sqlparser.AndExpressions(predicates...)
 }
 
 func (rp *routePlan) pushOutputColumns(col []*sqlparser.ColName, _ *semantics.SemTable) []int {

go/vt/vtgate/planbuilder/jointree_transformers.go

@@ -83,7 +83,10 @@ func transformRoutePlan(n *routePlan) (*route, error) {
 		switch predicate := predicate.(type) {
 		case *sqlparser.ComparisonExpr:
 			if predicate.Operator == sqlparser.InOp {
-				predicate.Right = sqlparser.ListArg(engine.ListVarName)
+				switch predicate.Left.(type) {
+				case *sqlparser.ColName:
+					predicate.Right = sqlparser.ListArg(engine.ListVarName)
+				}
 			}
 		}
 	}

go/vt/vtgate/planbuilder/plan_test.go

@@ -150,7 +150,10 @@ func init() {
 	vindexes.Register("costly", newCostlyIndex)
 }
 
-const samePlanMarker = "Gen4 plan same as above\n"
+const (
+	samePlanMarker  = "Gen4 plan same as above\n"
+	gen4ErrorPrefix = "Gen4 error: "
+)
 
 func TestPlan(t *testing.T) {
 	vschemaWrapper := &vschemaWrapper{
@@ -235,7 +238,11 @@ func TestWithDefaultKeyspaceFromFile(t *testing.T) {
 	// We are testing this separately so we can set a default keyspace
 	testOutputTempDir, err := ioutil.TempDir("", "plan_test")
 	require.NoError(t, err)
-	defer os.RemoveAll(testOutputTempDir)
+	defer func() {
+		if !t.Failed() {
+			_ = os.RemoveAll(testOutputTempDir)
+		}
+	}()
 	vschema := &vschemaWrapper{
 		v: loadSchema(t, "schema_test.json"),
 		keyspace: &vindexes.Keyspace{
@@ -607,9 +614,11 @@ func iterateExecFile(name string) (testCaseIterator chan testCase) {
 			if err != nil && err != io.EOF {
 				panic(fmt.Sprintf("error reading file %s line# %d: %s", name, lineno, err.Error()))
 			}
-			if len(binput) > 0 && string(binput) == samePlanMarker {
+			nextLine := string(binput)
+			switch {
+			case nextLine == samePlanMarker:
 				output2Planner = output
-			} else if len(binput) > 0 && (binput[0] == '"' || binput[0] == '{') {
+			case strings.HasPrefix(nextLine, "{"):
 				output2Planner = append(output2Planner, binput...)
 				for {
 					l, err := r.ReadBytes('\n')
@@ -627,8 +636,9 @@ func iterateExecFile(name string) (testCaseIterator chan testCase) {
 						break
 					}
 				}
+			case strings.HasPrefix(nextLine, gen4ErrorPrefix):
+				output2Planner = []byte(nextLine[len(gen4ErrorPrefix) : len(nextLine)-1])
 			}
-
 			testCaseIterator <- testCase{
 				file:             name,
 				lineno:           lineno,

go/vt/vtgate/planbuilder/project.go

@@ -20,7 +20,7 @@ import (
 	"errors"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
-	"vitess.io/vitess/go/vt/proto/vtrpc"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/engine"
@@ -46,7 +46,7 @@ func planProjection(pb *primitiveBuilder, in logicalPlan, expr *sqlparser.Aliase
 		} else {
 			// Pushing of non-trivial expressions not allowed for RHS of left joins.
 			if _, ok := expr.Expr.(*sqlparser.ColName); !ok && node.ejoin.Opcode == engine.LeftJoin {
-				return nil, nil, 0, errors.New("unsupported: cross-shard left join and column expressions")
+				return nil, nil, 0, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: cross-shard left join and column expressions")
 			}
 
 			newRight, col, colNumber, err := planProjection(pb, node.Right, expr, origin)
@@ -167,5 +167,5 @@ func planProjection(pb *primitiveBuilder, in logicalPlan, expr *sqlparser.Aliase
 		return node, rc, len(node.resultColumns) - 1, nil
 
 	}
-	return nil, nil, 0, vterrors.Errorf(vtrpc.Code_UNIMPLEMENTED, "[BUG] unreachable %T.projection", in)
+	return nil, nil, 0, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "[BUG] unreachable %T.projection", in)
 }

go/vt/vtgate/planbuilder/route.go

@@ -727,7 +727,12 @@ func (rb *route) computeCompositeINPlan(pb *primitiveBuilder, comparison *sqlpar
 // iterateCompositeIN recursively walks the LHS tuple of the IN clause looking
 // for column names. For those that match a vindex, it builds a multi-value plan
 // using the corresponding values in the RHS. It returns the best of the plans built.
-func (rb *route) iterateCompositeIN(pb *primitiveBuilder, comparison *sqlparser.ComparisonExpr, coordinates []int, tuple sqlparser.ValTuple) (opcode engine.RouteOpcode, vindex vindexes.SingleColumn, values sqlparser.Expr) {
+func (rb *route) iterateCompositeIN(
+	pb *primitiveBuilder,
+	comparison *sqlparser.ComparisonExpr,
+	coordinates []int,
+	tuple sqlparser.ValTuple,
+) (opcode engine.RouteOpcode, vindex vindexes.SingleColumn, values sqlparser.Expr) {
 	opcode = engine.SelectScatter
 
 	cindex := len(coordinates)

go/vt/vtgate/planbuilder/route_planning.go

@@ -46,6 +46,11 @@ func gen4Planner(_ string) func(sqlparser.Statement, *sqlparser.ReservedVars, Co
 }
 
 func newBuildSelectPlan(sel *sqlparser.Select, vschema ContextVSchema) (engine.Primitive, error) {
+
+	directives := sqlparser.ExtractCommentDirectives(sel.Comments)
+	if len(directives) > 0 {
+		return nil, semantics.Gen4NotSupportedF("comment directives")
+	}
 	keyspace, err := vschema.DefaultKeyspace()
 	if err != nil {
 		return nil, err
@@ -256,6 +261,10 @@ func planLimit(limit *sqlparser.Limit, plan logicalPlan) (logicalPlan, error) {
 
 func planHorizon(sel *sqlparser.Select, plan logicalPlan, semTable *semantics.SemTable) (logicalPlan, error) {
 	rb, ok := plan.(*route)
+	if !ok && semTable.ProjectionErr != nil {
+		return nil, semTable.ProjectionErr
+	}
+
 	if ok && rb.isSingleShard() {
 		createSingleShardRoutePlan(sel, rb)
 		return plan, nil
@@ -270,7 +279,7 @@ func planHorizon(sel *sqlparser.Select, plan logicalPlan, semTable *semantics.Se
 		return nil, err
 	}
 	for _, e := range qp.selectExprs {
-		if _, err := pushProjection(e, plan, semTable); err != nil {
+		if _, err := pushProjection(e, plan, semTable, true); err != nil {
 			return nil, err
 		}
 	}
@@ -625,7 +634,7 @@ func findColumnVindex(a *routePlan, exp sqlparser.Expr, sem *semantics.SemTable)
 
 	var singCol vindexes.SingleColumn
 
-	_ = a.tables.visit(func(table *routeTable) error {
+	_ = visitTables(a.tables, func(table *routeTable) error {
 		if leftDep.IsSolvedBy(table.qtable.TableID) {
 			for _, vindex := range table.vtable.ColumnVindexes {
 				sC, isSingle := vindex.Vindex.(vindexes.SingleColumn)