Add support for delete planning with limits in presence of foreign ke…

…ys (#15097) Signed-off-by: Manan Gupta <[email protected]> Signed-off-by: Harshit Gangal <[email protected]> Co-authored-by: Harshit Gangal <[email protected]>
vitessio · Feb 1, 2024 · 19af51c · 19af51c
1 parent aecf793
commit 19af51c
Show file tree

Hide file tree

Showing 6 changed files with 411 additions and 68 deletions.
diff --git a/go/test/endtoend/vtgate/foreignkey/fk_fuzz_test.go b/go/test/endtoend/vtgate/foreignkey/fk_fuzz_test.go
@@ -186,8 +186,12 @@ func (fz *fuzzer) generateSingleDeleteDMLQuery() string {
 	tableId := rand.Intn(len(fkTables))
 	idValue := 1 + rand.Intn(fz.maxValForId)
 	setVarFkChecksVal := fz.getSetVarFkChecksVal()
-	query := fmt.Sprintf("delete %vfrom %v where id = %v", setVarFkChecksVal, fkTables[tableId], idValue)
-	return query
+	delWithLimit := rand.Intn(2)
+	if delWithLimit == 0 {
+		return fmt.Sprintf("delete %vfrom %v where id = %v", setVarFkChecksVal, fkTables[tableId], idValue)
+	}
+	limitCount := rand.Intn(3)
+	return fmt.Sprintf("delete %vfrom %v order by id limit %v", setVarFkChecksVal, fkTables[tableId], limitCount)
 }
 
 // generateMultiDeleteDMLQuery generates a DELETE query using 2 tables from the parameters for the fuzzer.
@@ -604,61 +608,65 @@ func TestFkFuzzTest(t *testing.T) {
 		insertShare    int
 		deleteShare    int
 		updateShare    int
-	}{{
-		name:           "Single Thread - Only Inserts",
-		concurrency:    1,
-		timeForTesting: 5 * time.Second,
-		maxValForCol:   5,
-		maxValForId:    10,
-		insertShare:    100,
-		deleteShare:    0,
-		updateShare:    0,
-	}, {
-		name:           "Single Thread - Balanced Inserts and Deletes",
-		concurrency:    1,
-		timeForTesting: 5 * time.Second,
-		maxValForCol:   5,
-		maxValForId:    10,
-		insertShare:    50,
-		deleteShare:    50,
-		updateShare:    0,
-	}, {
-		name:           "Single Thread - Balanced Inserts and Updates",
-		concurrency:    1,
-		timeForTesting: 5 * time.Second,
-		maxValForCol:   5,
-		maxValForId:    10,
-		insertShare:    50,
-		deleteShare:    0,
-		updateShare:    50,
-	}, {
-		name:           "Single Thread - Balanced Inserts, Updates and Deletes",
-		concurrency:    1,
-		timeForTesting: 5 * time.Second,
-		maxValForCol:   5,
-		maxValForId:    10,
-		insertShare:    50,
-		deleteShare:    50,
-		updateShare:    50,
-	}, {
-		name:           "Multi Thread - Only Inserts",
-		concurrency:    30,
-		timeForTesting: 5 * time.Second,
-		maxValForCol:   5,
-		maxValForId:    30,
-		insertShare:    100,
-		deleteShare:    0,
-		updateShare:    0,
-	}, {
-		name:           "Multi Thread - Balanced Inserts, Updates and Deletes",
-		concurrency:    30,
-		timeForTesting: 5 * time.Second,
-		maxValForCol:   5,
-		maxValForId:    30,
-		insertShare:    50,
-		deleteShare:    50,
-		updateShare:    50,
-	}}
+	}{
+		{
+			name:           "Single Thread - Only Inserts",
+			concurrency:    1,
+			timeForTesting: 5 * time.Second,
+			maxValForCol:   5,
+			maxValForId:    10,
+			insertShare:    100,
+			deleteShare:    0,
+			updateShare:    0,
+		}, {
+			name:           "Single Thread - Balanced Inserts and Deletes",
+			concurrency:    1,
+			timeForTesting: 5 * time.Second,
+			maxValForCol:   5,
+			maxValForId:    10,
+			insertShare:    50,
+			deleteShare:    50,
+			updateShare:    0,
+		}, {
+			name:           "Single Thread - Balanced Inserts and Updates",
+			concurrency:    1,
+			timeForTesting: 5 * time.Second,
+			maxValForCol:   5,
+			maxValForId:    10,
+			insertShare:    50,
+			deleteShare:    0,
+			updateShare:    50,
+		},
+		{
+			name:           "Single Thread - Balanced Inserts, Updates and Deletes",
+			concurrency:    1,
+			timeForTesting: 5 * time.Second,
+			maxValForCol:   5,
+			maxValForId:    10,
+			insertShare:    50,
+			deleteShare:    50,
+			updateShare:    50,
+		},
+		{
+			name:           "Multi Thread - Only Inserts",
+			concurrency:    30,
+			timeForTesting: 5 * time.Second,
+			maxValForCol:   5,
+			maxValForId:    30,
+			insertShare:    100,
+			deleteShare:    0,
+			updateShare:    0,
+		}, {
+			name:           "Multi Thread - Balanced Inserts, Updates and Deletes",
+			concurrency:    30,
+			timeForTesting: 5 * time.Second,
+			maxValForCol:   5,
+			maxValForId:    30,
+			insertShare:    50,
+			deleteShare:    50,
+			updateShare:    50,
+		},
+	}
 
 	valTrue := true
 	valFalse := false

diff --git a/go/test/endtoend/vtgate/foreignkey/fk_test.go b/go/test/endtoend/vtgate/foreignkey/fk_test.go
@@ -687,6 +687,36 @@ func TestFkScenarios(t *testing.T) {
 				"select * from fk_t17 order by id",
 				"select * from fk_t19 order by id",
 			},
+		}, {
+			name: "Delete with limit success",
+			dataQueries: []string{
+				"insert into fk_t15(id, col) values (1, 7), (2, 9)",
+				"insert into fk_t16(id, col) values (1, 7), (2, 9)",
+				"insert into fk_t17(id, col) values (1, 7)",
+				"insert into fk_t19(id, col) values (1, 7)",
+			},
+			dmlQuery: "delete from fk_t15 order by id limit 1",
+			assertionQueries: []string{
+				"select * from fk_t15 order by id",
+				"select * from fk_t16 order by id",
+				"select * from fk_t17 order by id",
+				"select * from fk_t19 order by id",
+			},
+		}, {
+			name: "Delete with limit 0 success",
+			dataQueries: []string{
+				"insert into fk_t15(id, col) values (1, 7), (2, 9)",
+				"insert into fk_t16(id, col) values (1, 7), (2, 9)",
+				"insert into fk_t17(id, col) values (1, 7)",
+				"insert into fk_t19(id, col) values (1, 7)",
+			},
+			dmlQuery: "delete from fk_t15 order by id limit 0",
+			assertionQueries: []string{
+				"select * from fk_t15 order by id",
+				"select * from fk_t16 order by id",
+				"select * from fk_t17 order by id",
+				"select * from fk_t19 order by id",
+			},
 		},
 	}
 

diff --git a/go/vt/vtgate/engine/delete_with_input.go b/go/vt/vtgate/engine/delete_with_input.go
@@ -62,6 +62,9 @@ func (del *DeleteWithInput) TryExecute(ctx context.Context, vcursor VCursor, bin
 	if err != nil {
 		return nil, err
 	}
+	if inputRes == nil || len(inputRes.Rows) == 0 {
+		return &sqltypes.Result{}, nil
+	}
 
 	var bv *querypb.BindVariable
 	if len(del.OutputCols) == 1 {

diff --git a/go/vt/vtgate/planbuilder/operators/delete.go b/go/vt/vtgate/planbuilder/operators/delete.go
@@ -21,6 +21,7 @@ import (
 
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vterrors"
+	"vitess.io/vitess/go/vt/vtgate/engine"
 	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
 	"vitess.io/vitess/go/vt/vtgate/semantics"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
@@ -78,8 +79,15 @@ func (d *Delete) ShortDescription() string {
 }
 
 func createOperatorFromDelete(ctx *plancontext.PlanningContext, deleteStmt *sqlparser.Delete) (op Operator) {
-	delClone := sqlparser.CloneRefOfDelete(deleteStmt)
+	childFks := ctx.SemTable.GetChildForeignKeysForTable(deleteStmt.Targets[0])
+
+	// If the delete statement has a limit and has foreign keys, we will use a DeleteWithInput
+	// operator wherein we do a selection first and use that output for the subsequent deletes.
+	if len(childFks) > 0 && deleteStmt.Limit != nil {
+		return deletePlanningForLimitFk(ctx, deleteStmt)
+	}
 
+	delClone := sqlparser.CloneRefOfDelete(deleteStmt)
 	delOp := createDeleteOperator(ctx, deleteStmt)
 	op = delOp
 
@@ -90,19 +98,57 @@ func createOperatorFromDelete(ctx *plancontext.PlanningContext, deleteStmt *sqlp
 		}
 	}
 
-	childFks := ctx.SemTable.GetChildForeignKeysForTable(deleteStmt.Targets[0])
-	// If there are no foreign key constraints, then we don't need to do anything.
+	// If there are no foreign key constraints, then we don't need to do anything special.
 	if len(childFks) == 0 {
 		return op
 	}
-	// If the delete statement has a limit, we don't support it yet.
-	if delClone.Limit != nil {
-		panic(vterrors.VT12001("foreign keys management at vitess with limit"))
-	}
 
 	return createFkCascadeOpForDelete(ctx, op, delClone, childFks, delOp.Target.VTable)
 }
 
+func deletePlanningForLimitFk(ctx *plancontext.PlanningContext, del *sqlparser.Delete) Operator {
+	delClone := ctx.SemTable.Clone(del).(*sqlparser.Delete)
+	del.Limit = nil
+	del.OrderBy = nil
+
+	selectStmt := &sqlparser.Select{
+		From:    delClone.TableExprs,
+		Where:   delClone.Where,
+		OrderBy: delClone.OrderBy,
+		Limit:   delClone.Limit,
+		Lock:    sqlparser.ForUpdateLock,
+	}
+	ts := ctx.SemTable.Targets[del.Targets[0].Name]
+	ti, err := ctx.SemTable.TableInfoFor(ts)
+	if err != nil {
+		panic(vterrors.VT13001(err.Error()))
+	}
+	vTbl := ti.GetVindexTable()
+
+	var leftComp sqlparser.ValTuple
+	cols := make([]*sqlparser.ColName, 0, len(vTbl.PrimaryKey))
+	for _, col := range vTbl.PrimaryKey {
+		colName := sqlparser.NewColName(col.String())
+		selectStmt.SelectExprs = append(selectStmt.SelectExprs, aeWrap(colName))
+		cols = append(cols, colName)
+		leftComp = append(leftComp, colName)
+		ctx.SemTable.Recursive[colName] = ts
+	}
+	// optimize for case when there is only single column on left hand side.
+	var lhs sqlparser.Expr = leftComp
+	if len(leftComp) == 1 {
+		lhs = leftComp[0]
+	}
+	compExpr := sqlparser.NewComparisonExpr(sqlparser.InOp, lhs, sqlparser.ListArg(engine.DmVals), nil)
+	del.Where = sqlparser.NewWhere(sqlparser.WhereClause, compExpr)
+
+	return &DeleteWithInput{
+		Delete: createOperatorFromDelete(ctx, del),
+		Source: createOperatorFromSelect(ctx, selectStmt),
+		cols:   cols,
+	}
+}
+
 func createDeleteOperator(ctx *plancontext.PlanningContext, del *sqlparser.Delete) *Delete {
 	op := crossJoin(ctx, del.TableExprs)