refactor: move DML logic to sql_builder.go (vitessio#13920)

go/vt/sqlparser/ast_funcs.go

@@ -1150,6 +1150,22 @@ func (node *Update) AddWhere(expr Expr) {
 	}
 }
 
+// AddWhere adds the boolean expression to the
+// WHERE clause as an AND condition.
+func (node *Delete) AddWhere(expr Expr) {
+	if node.Where == nil {
+		node.Where = &Where{
+			Type: WhereClause,
+			Expr: expr,
+		}
+		return
+	}
+	node.Where.Expr = &AndExpr{
+		Left:  node.Where.Expr,
+		Right: expr,
+	}
+}
+
 // AddOrder adds an order by element
 func (node *Union) AddOrder(order *Order) {
 	node.OrderBy = append(node.OrderBy, order)

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -401,20 +401,29 @@ func newRoutingParams(ctx *plancontext.PlanningContext, opCode engine.Opcode) *e
 }
 
 func transformRoutePlan(ctx *plancontext.PlanningContext, op *operators.Route) (logicalPlan, error) {
-	switch src := op.Source.(type) {
-	case *operators.Insert:
-		return transformInsertPlan(ctx, op, src)
-	case *operators.Update:
-		return transformUpdatePlan(ctx, op, src)
-	case *operators.Delete:
-		return transformDeletePlan(ctx, op, src)
-	}
-	condition := getVindexPredicate(ctx, op)
-	sel, err := operators.ToSQL(ctx, op.Source)
+	stmt, dmlOp, err := operators.ToSQL(ctx, op.Source)
 	if err != nil {
 		return nil, err
 	}
-	replaceSubQuery(ctx, sel)
+
+	replaceSubQuery(ctx, stmt)
+
+	switch stmt := stmt.(type) {
+	case sqlparser.SelectStatement:
+		return buildRouteLogicalPlan(ctx, op, stmt)
+	case *sqlparser.Update:
+		return buildUpdateLogicalPlan(ctx, op, dmlOp, stmt)
+	case *sqlparser.Delete:
+		return buildDeleteLogicalPlan(ctx, op, dmlOp, stmt)
+	case *sqlparser.Insert:
+		return buildInsertLogicalPlan(ctx, op, dmlOp, stmt)
+	default:
+		return nil, vterrors.VT13001(fmt.Sprintf("dont know how to %T", stmt))
+	}
+}
+
+func buildRouteLogicalPlan(ctx *plancontext.PlanningContext, op *operators.Route, stmt sqlparser.SelectStatement) (logicalPlan, error) {
+	condition := getVindexPredicate(ctx, op)
 	eroute, err := routeToEngineRoute(ctx, op)
 	for _, order := range op.Ordering {
 		typ, collation, _ := ctx.SemTable.TypeForExpr(order.AST)
@@ -431,17 +440,17 @@ func transformRoutePlan(ctx *plancontext.PlanningContext, op *operators.Route) (
 	}
 	return &route{
 		eroute:    eroute,
-		Select:    sel,
+		Select:    stmt,
 		tables:    operators.TableID(op),
 		condition: condition,
 	}, nil
-
 }
 
-func transformInsertPlan(ctx *plancontext.PlanningContext, op *operators.Route, ins *operators.Insert) (i *insert, err error) {
+func buildInsertLogicalPlan(ctx *plancontext.PlanningContext, rb *operators.Route, op ops.Operator, stmt *sqlparser.Insert) (logicalPlan, error) {
+	ins := op.(*operators.Insert)
 	eins := &engine.Insert{
-		Opcode:            mapToInsertOpCode(op.Routing.OpCode(), ins.Input != nil),
-		Keyspace:          op.Routing.Keyspace(),
+		Opcode:            mapToInsertOpCode(rb.Routing.OpCode(), ins.Input != nil),
+		Keyspace:          rb.Routing.Keyspace(),
 		TableName:         ins.VTable.Name.String(),
 		Ignore:            ins.Ignore,
 		ForceNonStreaming: ins.ForceNonStreaming,
@@ -450,7 +459,7 @@ func transformInsertPlan(ctx *plancontext.PlanningContext, op *operators.Route,
 		VindexValues:      ins.VindexValues,
 		VindexValueOffset: ins.VindexValueOffset,
 	}
-	i = &insert{eInsert: eins}
+	lp := &insert{eInsert: eins}
 
 	// we would need to generate the query on the fly. The only exception here is
 	// when unsharded query with autoincrement for that there is no input operator.
@@ -459,15 +468,16 @@ func transformInsertPlan(ctx *plancontext.PlanningContext, op *operators.Route,
 	}
 
 	if ins.Input == nil {
-		eins.Query = generateQuery(ins.AST)
+		eins.Query = generateQuery(stmt)
 	} else {
-		i.source, err = transformToLogicalPlan(ctx, ins.Input)
+		newSrc, err := transformToLogicalPlan(ctx, ins.Input)
 		if err != nil {
-			return
+			return nil, err
 		}
+		lp.source = newSrc
 	}
 
-	return
+	return lp, nil
 }
 
 func mapToInsertOpCode(code engine.Opcode, insertSelect bool) engine.InsertOpcode {
@@ -541,16 +551,20 @@ func dmlFormatter(buf *sqlparser.TrackedBuffer, node sqlparser.SQLNode) {
 	node.Format(buf)
 }
 
-func transformUpdatePlan(ctx *plancontext.PlanningContext, op *operators.Route, upd *operators.Update) (logicalPlan, error) {
-	ast := upd.AST
-	replaceSubQuery(ctx, ast)
+func buildUpdateLogicalPlan(
+	ctx *plancontext.PlanningContext,
+	op *operators.Route,
+	dmlOp ops.Operator,
+	stmt *sqlparser.Update,
+) (logicalPlan, error) {
+	upd := dmlOp.(*operators.Update)
 	rp := newRoutingParams(ctx, op.Routing.OpCode())
 	err := op.Routing.UpdateRoutingParams(ctx, rp)
 	if err != nil {
 		return nil, err
 	}
 	edml := &engine.DML{
-		Query:             generateQuery(ast),
+		Query:             generateQuery(stmt),
 		TableNames:        []string{upd.VTable.Name.String()},
 		Vindexes:          upd.VTable.ColumnVindexes,
 		OwnedVindexQuery:  upd.OwnedVindexQuery,
@@ -567,11 +581,15 @@ func transformUpdatePlan(ctx *plancontext.PlanningContext, op *operators.Route,
 	return &primitiveWrapper{prim: e}, nil
 }
 
-func transformDeletePlan(ctx *plancontext.PlanningContext, op *operators.Route, del *operators.Delete) (logicalPlan, error) {
-	ast := del.AST
-	replaceSubQuery(ctx, ast)
-	rp := newRoutingParams(ctx, op.Routing.OpCode())
-	err := op.Routing.UpdateRoutingParams(ctx, rp)
+func buildDeleteLogicalPlan(
+	ctx *plancontext.PlanningContext,
+	rb *operators.Route,
+	dmlOp ops.Operator,
+	ast *sqlparser.Delete,
+) (logicalPlan, error) {
+	del := dmlOp.(*operators.Delete)
+	rp := newRoutingParams(ctx, rb.Routing.OpCode())
+	err := rb.Routing.UpdateRoutingParams(ctx, rp)
 	if err != nil {
 		return nil, err
 	}
@@ -583,7 +601,7 @@ func transformDeletePlan(ctx *plancontext.PlanningContext, op *operators.Route,
 		RoutingParameters: rp,
 	}
 
-	transformDMLPlan(del.VTable, edml, op.Routing, del.OwnedVindexQuery != "")
+	transformDMLPlan(del.VTable, edml, rb.Routing, del.OwnedVindexQuery != "")
 
 	e := &engine.Delete{
 		DML: edml,