Use proper keyspace when updating the query graph of a reference DML (#…

…17226)

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>

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")
+}

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")

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() {

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"
+      ]
+    }
+  }
+]