Multi Table Delete Support: join with reference table (#14784)

Signed-off-by: Harshit Gangal <[email protected]>

go/tools/asthelpergen/asthelpergen.go

@@ -29,7 +29,6 @@ import (
 	"golang.org/x/tools/go/packages"
 
 	"vitess.io/vitess/go/textutil"
-
 	"vitess.io/vitess/go/tools/codegen"
 )
 

go/vt/sqlparser/ast.go

@@ -365,8 +365,8 @@ type (
 		With       *With
 		Ignore     Ignore
 		Comments   *ParsedComments
-		Targets    TableNames
 		TableExprs TableExprs
+		Targets    TableNames
 		Partitions Partitions
 		Where      *Where
 		OrderBy    OrderBy

go/vt/sqlparser/ast_format.go

@@ -172,7 +172,7 @@ func (node *Delete) Format(buf *TrackedBuffer) {
 	if node.Ignore {
 		buf.literal("ignore ")
 	}
-	if node.Targets != nil {
+	if node.Targets != nil && !node.isSingleAliasExpr() {
 		buf.astPrintf(node, "%v ", node.Targets)
 	}
 	buf.astPrintf(node, "from %v%v%v%v%v", node.TableExprs, node.Partitions, node.Where, node.OrderBy, node.Limit)

go/vt/sqlparser/ast_funcs.go

@@ -24,13 +24,11 @@ import (
 	"strconv"
 	"strings"
 
-	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
-	"vitess.io/vitess/go/vt/vterrors"
-
-	"vitess.io/vitess/go/vt/log"
-
 	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/log"
 	querypb "vitess.io/vitess/go/vt/proto/query"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/vterrors"
 )
 
 // Walk calls postVisit on every node.
@@ -2156,25 +2154,31 @@ func (s SelectExprs) AllAggregation() bool {
 	return true
 }
 
-// RemoveKeyspaceFromColName removes the Qualifier.Qualifier on all ColNames in the expression tree
-func RemoveKeyspaceFromColName(expr Expr) {
-	RemoveKeyspace(expr)
-}
-
 // RemoveKeyspace removes the Qualifier.Qualifier on all ColNames in the AST
 func RemoveKeyspace(in SQLNode) {
 	// Walk will only return an error if we return an error from the inner func. safe to ignore here
 	_ = Walk(func(node SQLNode) (kontinue bool, err error) {
-		switch col := node.(type) {
-		case *ColName:
-			if col.Qualifier.Qualifier.NotEmpty() {
-				col.Qualifier.Qualifier = NewIdentifierCS("")
-			}
+		if col, ok := node.(*ColName); ok && col.Qualifier.Qualifier.NotEmpty() {
+			col.Qualifier.Qualifier = NewIdentifierCS("")
 		}
+
 		return true, nil
 	}, in)
 }
 
+// RemoveKeyspaceInTables removes the Qualifier on all TableNames in the AST
+func RemoveKeyspaceInTables(in SQLNode) {
+	// Walk will only return an error if we return an error from the inner func. safe to ignore here
+	Rewrite(in, nil, func(cursor *Cursor) bool {
+		if tbl, ok := cursor.Node().(TableName); ok && tbl.Qualifier.NotEmpty() {
+			tbl.Qualifier = NewIdentifierCS("")
+			cursor.Replace(tbl)
+		}
+
+		return true
+	})
+}
+
 func convertStringToInt(integer string) int {
 	val, _ := strconv.Atoi(integer)
 	return val
@@ -2536,3 +2540,14 @@ func IsLiteral(expr Expr) bool {
 func (ct *ColumnType) Invisible() bool {
 	return ct.Options.Invisible != nil && *ct.Options.Invisible
 }
+
+func (node *Delete) isSingleAliasExpr() bool {
+	if len(node.Targets) > 1 {
+		return false
+	}
+	if len(node.TableExprs) != 1 {
+		return false
+	}
+	_, isAliasExpr := node.TableExprs[0].(*AliasedTableExpr)
+	return isAliasExpr
+}

go/vt/sqlparser/parse_test.go

@@ -1360,7 +1360,7 @@ var (
 		input: "delete /* limit */ from a limit b",
 	}, {
 		input:  "delete /* alias where */ t.* from a as t where t.id = 2",
-		output: "delete /* alias where */ t from a as t where t.id = 2",
+		output: "delete /* alias where */ from a as t where t.id = 2",
 	}, {
 		input:  "delete t.* from t, t1",
 		output: "delete t from t, t1",

go/vt/vtgate/planbuilder/delete.go

@@ -144,16 +144,7 @@ func checkIfDeleteSupported(del *sqlparser.Delete, semTable *semantics.SemTable)
 		return semTable.NotUnshardedErr
 	}
 
-	// Delete is only supported for a single TableExpr which is supposed to be an aliased expression
-	multiShardErr := vterrors.VT12001("multi-shard or vindex write statement")
-	if len(del.TableExprs) != 1 {
-		return multiShardErr
-	}
-	_, isAliasedExpr := del.TableExprs[0].(*sqlparser.AliasedTableExpr)
-	if !isAliasedExpr {
-		return multiShardErr
-	}
-
+	// Delete is only supported for single Target.
 	if len(del.Targets) > 1 {
 		return vterrors.VT12001("multi-table DELETE statement in a sharded keyspace")
 	}

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -527,7 +527,7 @@ func transformRoutePlan(ctx *plancontext.PlanningContext, op *operators.Route) (
 	case *sqlparser.Update:
 		return buildUpdateLogicalPlan(ctx, op, dmlOp, stmt, hints)
 	case *sqlparser.Delete:
-		return buildDeleteLogicalPlan(ctx, op, dmlOp, hints)
+		return buildDeleteLogicalPlan(ctx, op, dmlOp, stmt, hints)
 	case *sqlparser.Insert:
 		return buildInsertLogicalPlan(op, dmlOp, stmt, hints)
 	default:
@@ -689,24 +689,20 @@ func buildUpdateLogicalPlan(
 	return &primitiveWrapper{prim: e}, nil
 }
 
-func buildDeleteLogicalPlan(
-	ctx *plancontext.PlanningContext,
-	rb *operators.Route,
-	dmlOp operators.Operator,
-	hints *queryHints,
-) (logicalPlan, error) {
+func buildDeleteLogicalPlan(ctx *plancontext.PlanningContext, rb *operators.Route, dmlOp operators.Operator, stmt *sqlparser.Delete, hints *queryHints) (logicalPlan, error) {
 	del := dmlOp.(*operators.Delete)
 	rp := newRoutingParams(ctx, rb.Routing.OpCode())
 	rb.Routing.UpdateRoutingParams(ctx, rp)
+	vtable := del.Target.VTable
 	edml := &engine.DML{
-		Query:             generateQuery(del.AST),
-		TableNames:        []string{del.VTable.Name.String()},
-		Vindexes:          del.VTable.Owned,
+		Query:             generateQuery(stmt),
+		TableNames:        []string{vtable.Name.String()},
+		Vindexes:          vtable.Owned,
 		OwnedVindexQuery:  del.OwnedVindexQuery,
 		RoutingParameters: rp,
 	}
 
-	transformDMLPlan(del.VTable, edml, rb.Routing, del.OwnedVindexQuery != "")
+	transformDMLPlan(vtable, edml, rb.Routing, del.OwnedVindexQuery != "")
 
 	e := &engine.Delete{
 		DML: edml,

go/vt/vtgate/planbuilder/operators/SQL_builder.go

@@ -309,7 +309,7 @@ func (ts *tableSorter) Swap(i, j int) {
 func removeKeyspaceFromSelectExpr(expr sqlparser.SelectExpr) {
 	switch expr := expr.(type) {
 	case *sqlparser.AliasedExpr:
-		sqlparser.RemoveKeyspaceFromColName(expr.Expr)
+		sqlparser.RemoveKeyspace(expr.Expr)
 	case *sqlparser.StarExpr:
 		expr.TableName.Qualifier = sqlparser.NewIdentifierCS("")
 	}
@@ -376,14 +376,36 @@ func buildQuery(op Operator, qb *queryBuilder) {
 	case *Update:
 		buildUpdate(op, qb)
 	case *Delete:
-		buildDML(op, qb)
+		buildDelete(op, qb)
 	case *Insert:
 		buildDML(op, qb)
 	default:
 		panic(vterrors.VT13001(fmt.Sprintf("unknown operator to convert to SQL: %T", op)))
 	}
 }
 
+func buildDelete(op *Delete, qb *queryBuilder) {
+	buildQuery(op.Source, qb)
+	// currently the qb builds a select query underneath.
+	// Will take the `From` and `Where` from this select
+	// and create a delete statement.
+	// TODO: change it to directly produce `delete` statement.
+	sel, ok := qb.stmt.(*sqlparser.Select)
+	if !ok {
+		panic(vterrors.VT13001("expected a select here"))
+	}
+
+	qb.dmlOperator = op
+	qb.stmt = &sqlparser.Delete{
+		Ignore:     sqlparser.Ignore(op.Ignore),
+		Targets:    sqlparser.TableNames{op.Target.Name},
+		TableExprs: sel.From,
+		Where:      sel.Where,
+		OrderBy:    op.OrderBy,
+		Limit:      op.Limit,
+	}
+}
+
 func buildUpdate(op *Update, qb *queryBuilder) {
 	tblName := sqlparser.NewTableName(op.QTable.Table.Name.String())
 	aTblExpr := &sqlparser.AliasedTableExpr{

go/vt/vtgate/planbuilder/operators/ast_to_op.go

@@ -72,7 +72,7 @@ func addWherePredicates(ctx *plancontext.PlanningContext, expr sqlparser.Expr, o
 	outerID := TableID(op)
 	exprs := sqlparser.SplitAndExpression(nil, expr)
 	for _, expr := range exprs {
-		sqlparser.RemoveKeyspaceFromColName(expr)
+		sqlparser.RemoveKeyspace(expr)
 		subq := sqc.handleSubquery(ctx, expr, outerID)
 		if subq != nil {
 			continue