[release-20.0] Use proper keyspace when updating the query graph of a…

… reference DML (#17226) (#17257) Signed-off-by: Florent Poinsard <[email protected]> Signed-off-by: Florent Poinsard <[email protected]> Co-authored-by: vitess-bot[bot] <108069721+vitess-bot[bot]@users.noreply.github.com> Co-authored-by: Florent Poinsard <[email protected]> Co-authored-by: Florent Poinsard <[email protected]>
vitessio · Nov 19, 2024 · ae9cc19 · ae9cc19
1 parent 495a44a
commit ae9cc19
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 1 deletion.
diff --git a/go/test/endtoend/vtgate/queries/reference/reference_test.go b/go/test/endtoend/vtgate/queries/reference/reference_test.go
@@ -171,3 +171,17 @@ func TestMultiReferenceQuery(t *testing.T) {
 
 	utils.Exec(t, conn, query)
 }
+
+func TestDMLReferenceUsingShardedKS(t *testing.T) {
+	version, err := cluster.GetMajorVersion("vtgate")
+	require.NoError(t, err)
+	if version != 20 {
+		t.Skip()
+	}
+
+	conn, closer := start(t)
+	defer closer()
+
+	utils.Exec(t, conn, "use sks")
+	utils.Exec(t, conn, "update zip_detail set zip_id = 1 where id = 1")
+}
diff --git a/go/vt/vtgate/planbuilder/operators/delete.go b/go/vt/vtgate/planbuilder/operators/delete.go
@@ -342,7 +342,7 @@ func updateQueryGraphWithSource(ctx *plancontext.PlanningContext, input Operator
 			if tbl.ID != tblID {
 				continue
 			}
-			tbl.Alias = sqlparser.NewAliasedTableExpr(sqlparser.NewTableName(vTbl.Name.String()), tbl.Alias.As.String())
+			tbl.Alias = sqlparser.NewAliasedTableExpr(sqlparser.NewTableNameWithQualifier(vTbl.Name.String(), vTbl.Keyspace.Name), tbl.Alias.As.String())
 			tbl.Table, _ = tbl.Alias.TableName()
 		}
 		return op, Rewrote("change query table point to source table")

diff --git a/go/vt/vtgate/planbuilder/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go
@@ -561,6 +561,7 @@ func (s *planTestSuite) TestWithUserDefaultKeyspaceFromFileSharded() {
 	}
 
 	s.testFile("select_cases_with_user_as_default.json", vschema, false)
+	s.testFile("dml_cases_with_user_as_default.json", vschema, false)
 }
 
 func (s *planTestSuite) TestWithSystemSchemaAsDefaultKeyspace() {

diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases_with_user_as_default.json b/go/vt/vtgate/planbuilder/testdata/dml_cases_with_user_as_default.json
@@ -0,0 +1,24 @@
+[
+  {
+    "comment": "Update reference table from sharded keyspace to unsharded keyspace",
+    "query": "update ambiguous_ref_with_source set done = true where id = 1;",
+    "plan": {
+      "QueryType": "UPDATE",
+      "Original": "update ambiguous_ref_with_source set done = true where id = 1;",
+      "Instructions": {
+        "OperatorType": "Update",
+        "Variant": "Unsharded",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "TargetTabletType": "PRIMARY",
+        "Query": "update ambiguous_ref_with_source set done = true where id = 1",
+        "Table": "ambiguous_ref_with_source"
+      },
+      "TablesUsed": [
+        "main.ambiguous_ref_with_source"
+      ]
+    }
+  }
+]