Delete with subquery support (#15219)

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

changelog/20.0/20.0.0/summary.md

@@ -7,6 +7,7 @@
     - [Vindex Hints](#vindex-hints)
     - [Update with Limit Support](#update-limit)
     - [Update with Multi Table Support](#multi-table-update)
+    - [Delete with Subquery Support](#delete-subquery)
   - **[Flag changes](#flag-changes)**
     - [`pprof-http` default change](#pprof-http-default)
 - **[Minor Changes](#minor-changes)**
@@ -45,6 +46,13 @@ Example: `update t1 join t2 on t1.id = t2.id join t3 on t1.col = t3.col set t1.b
 
 More details about how it works is available in [MySQL Docs](https://dev.mysql.com/doc/refman/8.0/en/update.html)
 
+#### <a id="delete-subquery"/> Delete with Subquery Support
+
+Support is added for sharded table delete with subquery
+
+Example: `delete from t1 where id in (select col from t2 where foo = 32 and bar = 43)`
+
+
 ### <a id="flag-changes"/>Flag Changes
 
 #### <a id="pprof-http-default"/> `pprof-http` Default Change

go/test/endtoend/vtgate/queries/dml/dml_test.go

@@ -204,3 +204,92 @@ func TestUpdateWithLimit(t *testing.T) {
 	require.EqualValues(t, 0, qr.RowsAffected)
 
 }
+
+// TestMultiTableUpdate executed multi-table update queries
+func TestMultiTableUpdate(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 20, "vtgate")
+
+	mcmp, closer := start(t)
+	defer closer()
+
+	// initial rows
+	mcmp.Exec("insert into order_tbl(region_id, oid, cust_no) values (1,1,4), (1,2,2), (2,3,5), (2,4,55)")
+	mcmp.Exec("insert into oevent_tbl(oid, ename) values (1,'a'), (2,'b'), (3,'a'), (4,'c')")
+
+	// check rows
+	mcmp.AssertMatches(`select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)] [INT64(2) INT64(3) INT64(5)] [INT64(2) INT64(4) INT64(55)]]`)
+	mcmp.AssertMatches(`select oid, ename from oevent_tbl order by oid`,
+		`[[INT64(1) VARCHAR("a")] [INT64(2) VARCHAR("b")] [INT64(3) VARCHAR("a")] [INT64(4) VARCHAR("c")]]`)
+
+	// multi table delete
+	qr := mcmp.Exec(`update order_tbl o join oevent_tbl ev on o.oid = ev.oid set ev.ename = 'a' where ev.oid > 3`)
+	assert.EqualValues(t, 1, qr.RowsAffected)
+
+	// check rows
+	mcmp.AssertMatches(`select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)] [INT64(2) INT64(3) INT64(5)] [INT64(2) INT64(4) INT64(55)]]`)
+	mcmp.AssertMatches(`select oid, ename from oevent_tbl order by oid`,
+		`[[INT64(1) VARCHAR("a")] [INT64(2) VARCHAR("b")] [INT64(3) VARCHAR("a")] [INT64(4) VARCHAR("a")]]`)
+
+	qr = mcmp.Exec(`update order_tbl o, oevent_tbl ev set ev.ename = 'xyz' where o.cust_no = ev.oid`)
+	assert.EqualValues(t, 2, qr.RowsAffected)
+
+	// check rows
+	mcmp.AssertMatches(`select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)] [INT64(2) INT64(3) INT64(5)] [INT64(2) INT64(4) INT64(55)]]`)
+	mcmp.AssertMatches(`select oid, ename from oevent_tbl order by oid`,
+		`[[INT64(1) VARCHAR("a")] [INT64(2) VARCHAR("xyz")] [INT64(3) VARCHAR("a")] [INT64(4) VARCHAR("xyz")]]`)
+}
+
+// TestDeleteWithSubquery executed delete queries with subqueries
+func TestDeleteWithSubquery(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 20, "vtgate")
+
+	mcmp, closer := start(t)
+	defer closer()
+
+	// initial rows
+	mcmp.Exec("insert into s_tbl(id, num) values (1,10), (2,10), (3,10), (4,20), (5,5), (6,15), (7,17), (8,80)")
+	mcmp.Exec("insert into order_tbl(region_id, oid, cust_no) values (1,1,4), (1,2,2), (2,3,5), (2,4,55)")
+
+	// check rows
+	mcmp.AssertMatches(`select id, num from s_tbl order by id`,
+		`[[INT64(1) INT64(10)] [INT64(2) INT64(10)] [INT64(3) INT64(10)] [INT64(4) INT64(20)] [INT64(5) INT64(5)] [INT64(6) INT64(15)] [INT64(7) INT64(17)] [INT64(8) INT64(80)]]`)
+	mcmp.AssertMatches(`select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)] [INT64(2) INT64(3) INT64(5)] [INT64(2) INT64(4) INT64(55)]]`)
+
+	// delete with subquery on s_tbl
+	qr := mcmp.Exec(`delete from s_tbl where id in (select oid from order_tbl)`)
+	require.EqualValues(t, 4, qr.RowsAffected)
+
+	// check rows
+	mcmp.AssertMatches(`select id, num from s_tbl order by id`,
+		`[[INT64(5) INT64(5)] [INT64(6) INT64(15)] [INT64(7) INT64(17)] [INT64(8) INT64(80)]]`)
+	mcmp.AssertMatches(`select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)] [INT64(2) INT64(3) INT64(5)] [INT64(2) INT64(4) INT64(55)]]`)
+
+	// delete with subquery on order_tbl
+	qr = mcmp.Exec(`delete from order_tbl where cust_no > (select num from s_tbl where id = 7)`)
+	require.EqualValues(t, 1, qr.RowsAffected)
+
+	// check rows
+	mcmp.AssertMatches(`select id, num from s_tbl order by id`,
+		`[[INT64(5) INT64(5)] [INT64(6) INT64(15)] [INT64(7) INT64(17)] [INT64(8) INT64(80)]]`)
+	mcmp.AssertMatches(`select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)] [INT64(2) INT64(3) INT64(5)]]`)
+
+	// delete with subquery from same table (fails on mysql) - subquery get's merged so fails for vitess
+	_, err := mcmp.ExecAllowAndCompareError(`delete from s_tbl where id in (select id from s_tbl)`)
+	require.ErrorContains(t, err, "You can't specify target table 's_tbl' for update in FROM clause (errno 1093) (sqlstate HY000)")
+
+	// delete with subquery from same table (fails on mysql) - subquery not merged so passes for vitess
+	qr = utils.Exec(t, mcmp.VtConn, `delete from order_tbl where region_id in (select cust_no from order_tbl)`)
+	require.EqualValues(t, 1, qr.RowsAffected)
+
+	// check rows
+	utils.AssertMatches(t, mcmp.VtConn, `select id, num from s_tbl order by id`,
+		`[[INT64(5) INT64(5)] [INT64(6) INT64(15)] [INT64(7) INT64(17)] [INT64(8) INT64(80)]]`)
+	utils.AssertMatches(t, mcmp.VtConn, `select region_id, oid, cust_no from order_tbl order by oid`,
+		`[[INT64(1) INT64(1) INT64(4)] [INT64(1) INT64(2) INT64(2)]]`)
+}

go/vt/vtgate/planbuilder/delete.go

@@ -144,19 +144,5 @@ func checkIfDeleteSupported(del *sqlparser.Delete, semTable *semantics.SemTable)
 		return vterrors.VT12001("multi-table DELETE statement with multi-target")
 	}
 
-	err := sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
-		switch node.(type) {
-		case *sqlparser.Subquery, *sqlparser.DerivedTable:
-			// We have a subquery, so we must fail the planning.
-			// If this subquery and the table expression were all belonging to the same unsharded keyspace,
-			// we would have already created a plan for them before doing these checks.
-			return false, vterrors.VT12001("subqueries in DML")
-		}
-		return true, nil
-	}, del)
-	if err != nil {
-		return err
-	}
-
 	return nil
 }

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -688,6 +688,8 @@ func buildUpdateLogicalPlan(
 	var vindexes []*vindexes.ColumnVindex
 	vQuery := ""
 	if len(upd.ChangedVindexValues) > 0 {
+		upd.OwnedVindexQuery.From = stmt.GetFrom()
+		upd.OwnedVindexQuery.Where = stmt.Where
 		vQuery = sqlparser.String(upd.OwnedVindexQuery)
 		vindexes = upd.Target.VTable.ColumnVindexes
 		if upd.OwnedVindexQuery.Limit != nil && len(upd.OwnedVindexQuery.OrderBy) == 0 {
@@ -709,6 +711,8 @@ func buildDeleteLogicalPlan(ctx *plancontext.PlanningContext, rb *operators.Rout
 	var vindexes []*vindexes.ColumnVindex
 	vQuery := ""
 	if del.OwnedVindexQuery != nil {
+		del.OwnedVindexQuery.From = stmt.GetFrom()
+		del.OwnedVindexQuery.Where = stmt.Where
 		vQuery = sqlparser.String(del.OwnedVindexQuery)
 		vindexes = del.Target.VTable.Owned
 	}

go/vt/vtgate/planbuilder/operators/delete.go

@@ -74,7 +74,15 @@ func (d *Delete) GetOrdering(*plancontext.PlanningContext) []OrderBy {
 func (d *Delete) ShortDescription() string {
 	ovq := ""
 	if d.OwnedVindexQuery != nil {
-		ovq = " vindexQuery:%s" + sqlparser.String(d.OwnedVindexQuery)
+		var cols, orderby, limit string
+		cols = fmt.Sprintf("COLUMNS: [%s]", sqlparser.String(d.OwnedVindexQuery.SelectExprs))
+		if len(d.OwnedVindexQuery.OrderBy) > 0 {
+			orderby = fmt.Sprintf(" ORDERBY: [%s]", sqlparser.String(d.OwnedVindexQuery.OrderBy))
+		}
+		if d.OwnedVindexQuery.Limit != nil {
+			limit = fmt.Sprintf(" LIMIT: [%s]", sqlparser.String(d.OwnedVindexQuery.Limit))
+		}
+		ovq = fmt.Sprintf(" vindexQuery(%s%s%s)", cols, orderby, limit)
 	}
 	return fmt.Sprintf("%s.%s%s", d.Target.VTable.Keyspace.Name, d.Target.VTable.Name.String(), ovq)
 }
@@ -90,8 +98,8 @@ func createOperatorFromDelete(ctx *plancontext.PlanningContext, deleteStmt *sqlp
 	}
 
 	delClone := sqlparser.CloneRefOfDelete(deleteStmt)
-	delOp := createDeleteOperator(ctx, deleteStmt)
-	op = delOp
+	var vTbl *vindexes.Table
+	op, vTbl = createDeleteOperator(ctx, deleteStmt)
 
 	if deleteStmt.Comments != nil {
 		op = &LockAndComment{
@@ -105,7 +113,7 @@ func createOperatorFromDelete(ctx *plancontext.PlanningContext, deleteStmt *sqlp
 		return op
 	}
 
-	return createFkCascadeOpForDelete(ctx, op, delClone, childFks, delOp.Target.VTable)
+	return createFkCascadeOpForDelete(ctx, op, delClone, childFks, vTbl)
 }
 
 func deleteWithInputPlanningRequired(childFks []vindexes.ChildFKInfo, deleteStmt *sqlparser.Delete) bool {
@@ -170,11 +178,12 @@ func deleteWithInputPlanningForFk(ctx *plancontext.PlanningContext, del *sqlpars
 	}
 }
 
-func createDeleteOperator(ctx *plancontext.PlanningContext, del *sqlparser.Delete) *Delete {
+func createDeleteOperator(ctx *plancontext.PlanningContext, del *sqlparser.Delete) (Operator, *vindexes.Table) {
 	op := crossJoin(ctx, del.TableExprs)
 
+	sqc := &SubQueryBuilder{}
 	if del.Where != nil {
-		op = addWherePredicates(ctx, del.Where.Expr, op)
+		op = addWherePredsToSubQueryBuilder(ctx, del.Where.Expr, op, sqc)
 	}
 
 	target := del.Targets[0]
@@ -207,9 +216,8 @@ func createDeleteOperator(ctx *plancontext.PlanningContext, del *sqlparser.Delet
 	var ovq *sqlparser.Select
 	if vTbl.Keyspace.Sharded && vTbl.Type == vindexes.TypeTable {
 		primaryVindex, _ := getVindexInformation(tblID, vTbl)
-		ate := tblInfo.GetAliasedTableExpr()
 		if len(vTbl.Owned) > 0 {
-			ovq = generateOwnedVindexQuery(ate, del, targetTbl, primaryVindex.Columns)
+			ovq = generateOwnedVindexQuery(del, targetTbl, primaryVindex.Columns)
 		}
 	}
 
@@ -222,21 +230,18 @@ func createDeleteOperator(ctx *plancontext.PlanningContext, del *sqlparser.Delet
 		},
 	}
 
-	if del.Limit == nil {
-		return delOp
-	}
-
-	addOrdering(ctx, del.OrderBy, delOp)
-
-	delOp.Source = &Limit{
-		Source: delOp.Source,
-		AST:    del.Limit,
+	if del.Limit != nil {
+		addOrdering(ctx, del.OrderBy, delOp)
+		delOp.Source = &Limit{
+			Source: delOp.Source,
+			AST:    del.Limit,
+		}
 	}
 
-	return delOp
+	return sqc.getRootOperator(delOp, nil), vTbl
 }
 
-func generateOwnedVindexQuery(tblExpr sqlparser.TableExpr, del *sqlparser.Delete, table TargetTable, ksidCols []sqlparser.IdentifierCI) *sqlparser.Select {
+func generateOwnedVindexQuery(del *sqlparser.Delete, table TargetTable, ksidCols []sqlparser.IdentifierCI) *sqlparser.Select {
 	var selExprs sqlparser.SelectExprs
 	for _, col := range ksidCols {
 		colName := makeColName(col, table, sqlparser.MultiTable(del.TableExprs))
@@ -248,11 +253,8 @@ func generateOwnedVindexQuery(tblExpr sqlparser.TableExpr, del *sqlparser.Delete
 			selExprs = append(selExprs, aeWrap(colName))
 		}
 	}
-	sqlparser.RemoveKeyspaceInTables(tblExpr)
 	return &sqlparser.Select{
 		SelectExprs: selExprs,
-		From:        del.TableExprs,
-		Where:       del.Where,
 		OrderBy:     del.OrderBy,
 		Limit:       del.Limit,
 		Lock:        sqlparser.ForUpdateLock,

go/vt/vtgate/planbuilder/operators/dml_planning.go

@@ -109,7 +109,6 @@ func buildChangedVindexesValues(
 	ctx *plancontext.PlanningContext,
 	update *sqlparser.Update,
 	table *vindexes.Table,
-	ate *sqlparser.AliasedTableExpr,
 	ksidCols []sqlparser.IdentifierCI,
 	assignments []SetExpr,
 ) (vv map[string]*engine.VindexValues, ownedVindexQuery *sqlparser.Select, subQueriesArgOnChangedVindex []string) {
@@ -145,11 +144,8 @@ func buildChangedVindexesValues(
 		return nil, nil, nil
 	}
 	// generate rest of the owned vindex query.
-	tblExpr := sqlparser.NewAliasedTableExpr(sqlparser.TableName{Name: table.Name}, ate.As.String())
 	ovq := &sqlparser.Select{
-		From:        []sqlparser.TableExpr{tblExpr},
 		SelectExprs: selExprs,
-		Where:       update.Where,
 		OrderBy:     update.OrderBy,
 		Limit:       update.Limit,
 		Lock:        sqlparser.ForUpdateLock,

go/vt/vtgate/planbuilder/operators/update.go

@@ -97,7 +97,15 @@ func (u *Update) TablesUsed() []string {
 func (u *Update) ShortDescription() string {
 	ovq := ""
 	if u.OwnedVindexQuery != nil {
-		ovq = " vindexQuery:%s" + sqlparser.String(u.OwnedVindexQuery)
+		var cols, orderby, limit string
+		cols = fmt.Sprintf("COLUMNS: [%s]", sqlparser.String(u.OwnedVindexQuery.SelectExprs))
+		if len(u.OwnedVindexQuery.OrderBy) > 0 {
+			orderby = fmt.Sprintf(" ORDERBY: [%s]", sqlparser.String(u.OwnedVindexQuery.OrderBy))
+		}
+		if u.OwnedVindexQuery.Limit != nil {
+			limit = fmt.Sprintf(" LIMIT: [%s]", sqlparser.String(u.OwnedVindexQuery.Limit))
+		}
+		ovq = fmt.Sprintf(" vindexQuery(%s%s%s)", cols, orderby, limit)
 	}
 	return fmt.Sprintf("%s.%s%s", u.Target.VTable.Keyspace.Name, u.Target.VTable.Name.String(), ovq)
 }
@@ -217,7 +225,7 @@ func createUpdateOperator(ctx *plancontext.PlanningContext, updStmt *sqlparser.U
 		Name:   name,
 	}
 
-	_, cvv, ovq, subQueriesArgOnChangedVindex := getUpdateVindexInformation(ctx, updStmt, targetTbl, tblInfo.GetAliasedTableExpr(), assignments)
+	_, cvv, ovq, subQueriesArgOnChangedVindex := getUpdateVindexInformation(ctx, updStmt, targetTbl, assignments)
 
 	updOp := &Update{
 		DMLCommon: &DMLCommon{
@@ -249,15 +257,14 @@ func getUpdateVindexInformation(
 	ctx *plancontext.PlanningContext,
 	updStmt *sqlparser.Update,
 	table TargetTable,
-	ate *sqlparser.AliasedTableExpr,
 	assignments []SetExpr,
 ) ([]*VindexPlusPredicates, map[string]*engine.VindexValues, *sqlparser.Select, []string) {
 	if !table.VTable.Keyspace.Sharded {
 		return nil, nil, nil, nil
 	}
 
 	primaryVindex, vindexAndPredicates := getVindexInformation(table.ID, table.VTable)
-	changedVindexValues, ownedVindexQuery, subQueriesArgOnChangedVindex := buildChangedVindexesValues(ctx, updStmt, table.VTable, ate, primaryVindex.Columns, assignments)
+	changedVindexValues, ownedVindexQuery, subQueriesArgOnChangedVindex := buildChangedVindexesValues(ctx, updStmt, table.VTable, primaryVindex.Columns, assignments)
 	return vindexAndPredicates, changedVindexValues, ownedVindexQuery, subQueriesArgOnChangedVindex
 }
 

go/vt/vtgate/planbuilder/plan_test.go

@@ -279,6 +279,7 @@ func TestOne(t *testing.T) {
 	lv := loadSchema(t, "vschemas/schema.json", true)
 	setFks(t, lv)
 	addPKs(t, lv, "user", []string{"user", "music"})
+	addPKs(t, lv, "main", []string{"unsharded"})
 	vschema := &vschemawrapper.VSchemaWrapper{
 		V:           lv,
 		TestBuilder: TestBuilder,