From 115fd864d9aefe002b7700e5028f1d607cd802e1 Mon Sep 17 00:00:00 2001
From: Florent Poinsard <35779988+frouioui@users.noreply.github.com>
Date: Tue, 19 Nov 2024 12:54:29 -0600
Subject: [PATCH] Use proper keyspace when updating the query graph of a
 reference DML (#17226)

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>
Signed-off-by: Renan Rangel <rrangel@slack-corp.com>
---
 .../queries/reference/reference_test.go       |  9 +++++++
 go/vt/vtgate/planbuilder/operators/delete.go  |  2 +-
 go/vt/vtgate/planbuilder/plan_test.go         |  1 +
 .../dml_cases_with_user_as_default.json       | 24 +++++++++++++++++++
 4 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 go/vt/vtgate/planbuilder/testdata/dml_cases_with_user_as_default.json

diff --git a/go/test/endtoend/vtgate/queries/reference/reference_test.go b/go/test/endtoend/vtgate/queries/reference/reference_test.go
index 08e9cbe13b1..66d46dfaf15 100644
--- a/go/test/endtoend/vtgate/queries/reference/reference_test.go
+++ b/go/test/endtoend/vtgate/queries/reference/reference_test.go
@@ -171,3 +171,12 @@ func TestMultiReferenceQuery(t *testing.T) {
 
 	utils.Exec(t, conn, query)
 }
+
+func TestDMLReferenceUsingShardedKS(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 22, "vtgate")
+	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
index 4d30d9b9cc1..81e36d54315 100644
--- a/go/vt/vtgate/planbuilder/operators/delete.go
+++ b/go/vt/vtgate/planbuilder/operators/delete.go
@@ -328,7 +328,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
index 9cf92a91ddf..acba2caf937 100644
--- a/go/vt/vtgate/planbuilder/plan_test.go
+++ b/go/vt/vtgate/planbuilder/plan_test.go
@@ -557,6 +557,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
new file mode 100644
index 00000000000..ff66967c2ce
--- /dev/null
+++ 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"
+      ]
+    }
+  }
+]
\ No newline at end of file