Subqueries in SET condition of UPDATE statement in presence of foreig…

…n keys (#15163)

Signed-off-by: Manan Gupta <[email protected]>

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

@@ -170,6 +170,14 @@ func TestUpdateWithFK(t *testing.T) {
 	// Verify the result in u_t2 and u_t3 as well.
 	utils.AssertMatches(t, conn, `select * from u_t2 order by id`, `[[INT64(19) INT64(1234)] [INT64(342) NULL]]`)
 	utils.AssertMatches(t, conn, `select * from u_t3 order by id`, `[[INT64(1) INT64(12)] [INT64(32) INT64(13)]]`)
+
+	// Update with a subquery inside, such that the update is on a foreign key related column.
+	qr = utils.Exec(t, conn, `update u_t2 set col2 = (select col1 from u_t1 where id = 100) where id = 342`)
+	assert.EqualValues(t, 1, qr.RowsAffected)
+	// Verify the result in u_t1, u_t2 and u_t3.
+	utils.AssertMatches(t, conn, `select * from u_t1 order by id`, `[[INT64(1) INT64(13)] [INT64(10) INT64(12)] [INT64(100) INT64(13)] [INT64(1000) INT64(1234)]]`)
+	utils.AssertMatches(t, conn, `select * from u_t2 order by id`, `[[INT64(19) INT64(1234)] [INT64(342) INT64(13)]]`)
+	utils.AssertMatches(t, conn, `select * from u_t3 order by id`, `[[INT64(1) INT64(12)] [INT64(32) INT64(13)]]`)
 }
 
 // TestVstreamForFKBinLog tests that dml queries with fks are written with child row first approach in the binary logs.

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

@@ -104,8 +104,7 @@ func createOperatorFromUpdate(ctx *plancontext.PlanningContext, updStmt *sqlpars
 
 	vindexTable, routing := buildVindexTableForDML(ctx, tableInfo, qt, "update")
 
-	updClone := sqlparser.CloneRefOfUpdate(updStmt)
-	updOp := createUpdateOperator(ctx, updStmt, vindexTable, qt, routing)
+	updOp, updClone := createUpdateOperator(ctx, updStmt, vindexTable, qt, routing)
 
 	parentFks := ctx.SemTable.GetParentForeignKeysList()
 	childFks := ctx.SemTable.GetChildForeignKeysList()
@@ -127,13 +126,19 @@ func createOperatorFromUpdate(ctx *plancontext.PlanningContext, updStmt *sqlpars
 	return buildFkOperator(ctx, updOp, updClone, parentFks, childFks, vindexTable)
 }
 
-func createUpdateOperator(ctx *plancontext.PlanningContext, updStmt *sqlparser.Update, vindexTable *vindexes.Table, qt *QueryTable, routing Routing) Operator {
+func createUpdateOperator(ctx *plancontext.PlanningContext, updStmt *sqlparser.Update, vindexTable *vindexes.Table, qt *QueryTable, routing Routing) (Operator, *sqlparser.Update) {
 	sqc := &SubQueryBuilder{}
 	assignments := make([]SetExpr, len(updStmt.Exprs))
+	// updClone is used in foreign key planning to create the selection statements to be used for verification and selection.
+	// If we encounter subqueries, we want to fix the updClone to use the replaced expression, so that the pulled out subquery's
+	// result is used everywhere instead of running the subquery multiple times, which is wasteful.
+	updClone := sqlparser.CloneRefOfUpdate(updStmt)
 	for idx, updExpr := range updStmt.Exprs {
 		expr, subqs := sqc.pullOutValueSubqueries(ctx, updExpr.Expr, qt.ID, true)
 		if len(subqs) == 0 {
 			expr = updExpr.Expr
+		} else {
+			updClone.Exprs[idx].Expr = sqlparser.CloneExpr(expr)
 		}
 		proj := newProjExpr(aeWrap(expr))
 		if len(subqs) != 0 {
@@ -187,10 +192,17 @@ func createUpdateOperator(ctx *plancontext.PlanningContext, updStmt *sqlparser.U
 		}
 	}
 
-	return sqc.getRootOperator(route, decorator)
+	return sqc.getRootOperator(route, decorator), updClone
 }
 
 func buildFkOperator(ctx *plancontext.PlanningContext, updOp Operator, updClone *sqlparser.Update, parentFks []vindexes.ParentFKInfo, childFks []vindexes.ChildFKInfo, updatedTable *vindexes.Table) Operator {
+	// If the outermost operator is a subquery container, we want to do the foreign key planning inside it,
+	// because we want to Inner of the subquery to execute first and its result be used for the entire update planning.
+	subqc, isSubqc := updOp.(*SubQueryContainer)
+	if isSubqc {
+		subqc.Outer = buildFkOperator(ctx, subqc.Outer, updClone, parentFks, childFks, updatedTable)
+		return subqc
+	}
 	restrictChildFks, cascadeChildFks := splitChildFks(childFks)
 
 	op := createFKCascadeOp(ctx, updOp, updClone, cascadeChildFks, updatedTable)

go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json

@@ -3440,5 +3440,68 @@
         "unsharded_fk_allow.u_tbl3"
       ]
     }
+  },
+  {
+    "comment": "update query with an uncorrelated subquery",
+    "query": "update u_tbl4 set col41 = (select col14 from u_tbl1 where x = 2 and y = 4) where col4 = 3",
+    "plan": {
+      "QueryType": "UPDATE",
+      "Original": "update u_tbl4 set col41 = (select col14 from u_tbl1 where x = 2 and y = 4) where col4 = 3",
+      "Instructions": {
+        "OperatorType": "UncorrelatedSubquery",
+        "Variant": "PulloutValue",
+        "PulloutVars": [
+          "__sq1"
+        ],
+        "Inputs": [
+          {
+            "InputName": "SubQuery",
+            "OperatorType": "Route",
+            "Variant": "Unsharded",
+            "Keyspace": {
+              "Name": "unsharded_fk_allow",
+              "Sharded": false
+            },
+            "FieldQuery": "select col14 from u_tbl1 where 1 != 1",
+            "Query": "select col14 from u_tbl1 where x = 2 and y = 4 lock in share mode",
+            "Table": "u_tbl1"
+          },
+          {
+            "InputName": "Outer",
+            "OperatorType": "FKVerify",
+            "Inputs": [
+              {
+                "InputName": "VerifyParent-1",
+                "OperatorType": "Route",
+                "Variant": "Unsharded",
+                "Keyspace": {
+                  "Name": "unsharded_fk_allow",
+                  "Sharded": false
+                },
+                "FieldQuery": "select 1 from u_tbl4 left join u_tbl1 on u_tbl1.col14 = cast(:__sq1 as SIGNED) where 1 != 1",
+                "Query": "select 1 from u_tbl4 left join u_tbl1 on u_tbl1.col14 = cast(:__sq1 as SIGNED) where not (u_tbl4.col41) <=> (cast(:__sq1 as SIGNED)) and u_tbl4.col4 = 3 and cast(:__sq1 as SIGNED) is not null and u_tbl1.col14 is null limit 1 for share",
+                "Table": "u_tbl1, u_tbl4"
+              },
+              {
+                "InputName": "PostVerify",
+                "OperatorType": "Update",
+                "Variant": "Unsharded",
+                "Keyspace": {
+                  "Name": "unsharded_fk_allow",
+                  "Sharded": false
+                },
+                "TargetTabletType": "PRIMARY",
+                "Query": "update /*+ SET_VAR(foreign_key_checks=OFF) */ u_tbl4 set col41 = :__sq1 where col4 = 3",
+                "Table": "u_tbl4"
+              }
+            ]
+          }
+        ]
+      },
+      "TablesUsed": [
+        "unsharded_fk_allow.u_tbl1",
+        "unsharded_fk_allow.u_tbl4"
+      ]
+    }
   }
 ]