Update Planning for Limits in the presence of foreign keys

go/test/endtoend/vtgate/foreignkey/fk_fuzz_test.go

@@ -157,10 +157,15 @@ func (fz *fuzzer) generateUpdateDMLQuery() string {
 	idValue := 1 + rand.Intn(fz.maxValForId)
 	tableName := fkTables[tableId]
 	setVarFkChecksVal := fz.getSetVarFkChecksVal()
+	updWithLimit := rand.Intn(2)
+	limitCount := rand.Intn(3)
 	if tableName == "fk_t20" {
 		colValue := convertIntValueToString(rand.Intn(1 + fz.maxValForCol))
 		col2Value := convertIntValueToString(rand.Intn(1 + fz.maxValForCol))
-		return fmt.Sprintf("update %v%v set col = %v, col2 = %v where id = %v", setVarFkChecksVal, tableName, colValue, col2Value, idValue)
+		if updWithLimit == 0 {
+			return fmt.Sprintf("update %v%v set col = %v, col2 = %v where id = %v", setVarFkChecksVal, tableName, colValue, col2Value, idValue)
+		}
+		return fmt.Sprintf("update %v%v set col = %v, col2 = %v order by id limit %v", setVarFkChecksVal, tableName, colValue, col2Value, limitCount)
 	} else if isMultiColFkTable(tableName) {
 		if rand.Intn(2) == 0 {
 			colaValue := convertIntValueToString(rand.Intn(1 + fz.maxValForCol))
@@ -169,15 +174,24 @@ func (fz *fuzzer) generateUpdateDMLQuery() string {
 				colaValue = fz.generateExpression(rand.Intn(4)+1, "cola", "colb", "id")
 				colbValue = fz.generateExpression(rand.Intn(4)+1, "cola", "colb", "id")
 			}
-			return fmt.Sprintf("update %v%v set cola = %v, colb = %v where id = %v", setVarFkChecksVal, tableName, colaValue, colbValue, idValue)
+			if updWithLimit == 0 {
+				return fmt.Sprintf("update %v%v set cola = %v, colb = %v where id = %v", setVarFkChecksVal, tableName, colaValue, colbValue, idValue)
+			}
+			return fmt.Sprintf("update %v%v set cola = %v, colb = %v order by id limit %v", setVarFkChecksVal, tableName, colaValue, colbValue, limitCount)
 		} else {
 			colValue := fz.generateExpression(rand.Intn(4)+1, "cola", "colb", "id")
 			colToUpdate := []string{"cola", "colb"}[rand.Intn(2)]
-			return fmt.Sprintf("update %v set %v = %v where id = %v", tableName, colToUpdate, colValue, idValue)
+			if updWithLimit == 0 {
+				return fmt.Sprintf("update %v set %v = %v where id = %v", tableName, colToUpdate, colValue, idValue)
+			}
+			return fmt.Sprintf("update %v set %v = %v order by id limit %v", tableName, colToUpdate, colValue, limitCount)
 		}
 	} else {
 		colValue := fz.generateExpression(rand.Intn(4)+1, "col", "id")
-		return fmt.Sprintf("update %v%v set col = %v where id = %v", setVarFkChecksVal, tableName, colValue, idValue)
+		if updWithLimit == 0 {
+			return fmt.Sprintf("update %v%v set col = %v where id = %v", setVarFkChecksVal, tableName, colValue, idValue)
+		}
+		return fmt.Sprintf("update %v%v set col = %v order by id limit %v", setVarFkChecksVal, tableName, colValue, limitCount)
 	}
 }
 

go/test/endtoend/vtgate/foreignkey/fk_test.go

@@ -725,6 +725,51 @@ func TestFkScenarios(t *testing.T) {
 				"select * from fk_t17 order by id",
 				"select * from fk_t19 order by id",
 			},
+		}, {
+			name: "Update 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: "update fk_t15 set col = '2' 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: "Update 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: "update fk_t15 set col = '8' 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",
+			},
+		}, {
+			name: "Update with non-literal update and 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: "update fk_t15 set col = id + 3 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",
+			},
 		},
 	}
 

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

@@ -95,6 +95,15 @@
 
 func createOperatorFromUpdate(ctx *plancontext.PlanningContext, updStmt *sqlparser.Update) (op Operator) {
 	errIfUpdateNotSupported(ctx, updStmt)
+	parentFks := ctx.SemTable.GetParentForeignKeysList()
+	childFks := ctx.SemTable.GetChildForeignKeysList()
+
+	// We check if dml with input plan is required. DML with input planning is generally
+	// slower, because it does a selection and then creates a update statement wherein we have to
+	// list all the primary key values.
+	if updateWithInputPlanningRequired(childFks, parentFks, updStmt) {
+		return updateWithInputPlanningForFk(ctx, updStmt)
+	}
 
 	var updClone *sqlparser.Update
 	var vTbl *vindexes.Table
@@ -107,24 +116,70 @@
 		Lock:     sqlparser.ShareModeLock,
 	}
 
-	parentFks := ctx.SemTable.GetParentForeignKeysList()
-	childFks := ctx.SemTable.GetChildForeignKeysList()
 	if len(childFks) == 0 && len(parentFks) == 0 {
 		return op
 	}
 
-	// If the delete statement has a limit, we don't support it yet.
-	if updStmt.Limit != nil {
-		panic(vterrors.VT12001("update with limit with foreign key constraints"))
+	return buildFkOperator(ctx, op, updClone, parentFks, childFks, vTbl)
+}
+
+func updateWithInputPlanningForFk(ctx *plancontext.PlanningContext, upd *sqlparser.Update) Operator {
+	updClone := ctx.SemTable.Clone(upd).(*sqlparser.Update)
+	upd.Limit = nil
+
+	selectStmt := &sqlparser.Select{
+		From:    updClone.TableExprs,
+		Where:   updClone.Where,
+		OrderBy: updClone.OrderBy,
+		Limit:   updClone.Limit,
+		Lock:    sqlparser.ForUpdateLock,
 	}
 
-	// Now we check if any of the foreign key columns that are being udpated have dependencies on other updated columns.
-	// This is unsafe, and we currently don't support this in Vitess.
-	if err := ctx.SemTable.ErrIfFkDependentColumnUpdated(updStmt.Exprs); err != nil {
-		panic(err)
+	ate, isAliasTableExpr := upd.TableExprs[0].(*sqlparser.AliasedTableExpr)
+	if !isAliasTableExpr {
+		panic(vterrors.VT12001("update with limit with foreign key constraints using a complex table"))
+	}
+	ts := ctx.SemTable.TableSetFor(ate)
+	ti, err := ctx.SemTable.TableInfoFor(ts)
+	if err != nil {
+		panic(vterrors.VT13001(err.Error()))
 	}
+	vTbl := ti.GetVindexTable()
 
-	return buildFkOperator(ctx, op, updClone, parentFks, childFks, vTbl)
+	var leftComp sqlparser.ValTuple
+	cols := make([]*sqlparser.ColName, 0, len(vTbl.PrimaryKey))
+	for _, col := range vTbl.PrimaryKey {
+		colName := sqlparser.NewColNameWithQualifier(col.String(), vTbl.GetTableName())
+		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.DmlVals), nil)
+
+	upd.Where = sqlparser.NewWhere(sqlparser.WhereClause, compExpr)
+	return &DMLWithInput{
+		DML:    []Operator{createOperatorFromUpdate(ctx, upd)},
+		Source: createOperatorFromSelect(ctx, selectStmt),
+		cols:   [][]*sqlparser.ColName{cols},
+	}
+}
+
+func updateWithInputPlanningRequired(childFks []vindexes.ChildFKInfo, parentFks []vindexes.ParentFKInfo, updateStmt *sqlparser.Update) bool {
+	// If there are no foreign keys, we don't need to use delete with input.
+	if len(childFks) == 0 && len(parentFks) == 0 {
+		return false
+	}
+	// Limit requires dml with input.
+	if updateStmt.Limit != nil {
+		return true
+	}
+	return false
 }
 
 func errIfUpdateNotSupported(ctx *plancontext.PlanningContext, stmt *sqlparser.Update) {
@@ -150,6 +205,12 @@
 			panic(vterrors.VT12001("multi-table UPDATE statement with multi-target column update"))
 		}
 	}
+
+	// Now we check if any of the foreign key columns that are being udpated have dependencies on other updated columns.
+	// This is unsafe, and we currently don't support this in Vitess.
+	if err := ctx.SemTable.ErrIfFkDependentColumnUpdated(stmt.Exprs); err != nil {
+		panic(err)
+	}
 }
 
 func createUpdateOperator(ctx *plancontext.PlanningContext, updStmt *sqlparser.Update) (Operator, *vindexes.Table, *sqlparser.Update) {