From 27a78beb23d6d8976c338c149ffcea0df64d4bfd Mon Sep 17 00:00:00 2001
From: Manan Gupta <manan@planetscale.com>
Date: Thu, 19 Oct 2023 19:52:37 +0530
Subject: [PATCH] feat: copy the order by from the update query

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 go/vt/vtgate/planbuilder/operators/ast_to_op.go | 2 ++
 go/vt/vtgate/planbuilder/operators/delete.go    | 2 +-
 go/vt/vtgate/planbuilder/operators/update.go    | 4 +++-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/go/vt/vtgate/planbuilder/operators/ast_to_op.go b/go/vt/vtgate/planbuilder/operators/ast_to_op.go
index 6fef308c592..8d00e29918a 100644
--- a/go/vt/vtgate/planbuilder/operators/ast_to_op.go
+++ b/go/vt/vtgate/planbuilder/operators/ast_to_op.go
@@ -398,6 +398,7 @@ func createSelectionOp(
 	selectExprs []sqlparser.SelectExpr,
 	tableExprs sqlparser.TableExprs,
 	where *sqlparser.Where,
+	orderBy sqlparser.OrderBy,
 	limit *sqlparser.Limit,
 	lock sqlparser.Lock,
 ) (ops.Operator, error) {
@@ -405,6 +406,7 @@ func createSelectionOp(
 		SelectExprs: selectExprs,
 		From:        tableExprs,
 		Where:       where,
+		OrderBy:     orderBy,
 		Limit:       limit,
 		Lock:        lock,
 	}
diff --git a/go/vt/vtgate/planbuilder/operators/delete.go b/go/vt/vtgate/planbuilder/operators/delete.go
index b3a24524018..dd37ed5db01 100644
--- a/go/vt/vtgate/planbuilder/operators/delete.go
+++ b/go/vt/vtgate/planbuilder/operators/delete.go
@@ -190,7 +190,7 @@ func createFkCascadeOpForDelete(ctx *plancontext.PlanningContext, parentOp ops.O
 		}
 		fkChildren = append(fkChildren, fkChild)
 	}
-	selectionOp, err := createSelectionOp(ctx, selectExprs, delStmt.TableExprs, delStmt.Where, nil, sqlparser.ForUpdateLock)
+	selectionOp, err := createSelectionOp(ctx, selectExprs, delStmt.TableExprs, delStmt.Where, nil, nil, sqlparser.ForUpdateLock)
 	if err != nil {
 		return nil, err
 	}
diff --git a/go/vt/vtgate/planbuilder/operators/update.go b/go/vt/vtgate/planbuilder/operators/update.go
index 722808c9e3b..4b1f36d2e43 100644
--- a/go/vt/vtgate/planbuilder/operators/update.go
+++ b/go/vt/vtgate/planbuilder/operators/update.go
@@ -298,7 +298,7 @@ func createFKCascadeOp(ctx *plancontext.PlanningContext, parentOp ops.Operator,
 		fkChildren = append(fkChildren, fkChild)
 	}
 
-	selectionOp, err := createSelectionOp(ctx, selectExprs, updStmt.TableExprs, updStmt.Where, nil, sqlparser.ForUpdateLock)
+	selectionOp, err := createSelectionOp(ctx, selectExprs, updStmt.TableExprs, updStmt.Where, updStmt.OrderBy, nil, sqlparser.ForUpdateLock)
 	if err != nil {
 		return nil, err
 	}
@@ -616,6 +616,7 @@ func createFkVerifyOpForParentFKForUpdate(ctx *plancontext.PlanningContext, updS
 				sqlparser.NewJoinCondition(joinCond, nil)),
 		},
 		sqlparser.NewWhere(sqlparser.WhereClause, whereCond),
+		nil,
 		sqlparser.NewLimitWithoutOffset(1),
 		sqlparser.ShareModeLock)
 }
@@ -685,6 +686,7 @@ func createFkVerifyOpForChildFKForUpdate(ctx *plancontext.PlanningContext, updSt
 				sqlparser.NewJoinCondition(joinCond, nil)),
 		},
 		sqlparser.NewWhere(sqlparser.WhereClause, whereCond),
+		nil,
 		sqlparser.NewLimitWithoutOffset(1),
 		sqlparser.ShareModeLock)
 }