Fix owned vindex query to keep alias in table expression

go/test/endtoend/vtgate/lookup_test.go

@@ -470,7 +470,7 @@ func TestConsistentLookupUpdate(t *testing.T) {
 	3 3 3
 	3 4 4
 	*/
-	utils.Exec(t, conn, "update t4 set id2 = '42' where id1 = 1")
+	utils.Exec(t, conn, "update t4 a set a.id2 = '42' where a.id1 = 1")
 	qr = utils.Exec(t, conn, "select id1, id2 from t4 order by id1")
 	if got, want := fmt.Sprintf("%v", qr.Rows), `[[INT64(1) VARCHAR("42")] [INT64(2) VARCHAR("2")] [INT64(3) VARCHAR("3")] [INT64(4) VARCHAR("3")]]`; got != want {
 		t.Errorf("select:\n%v want\n%v", got, want)

go/vt/sqlparser/sql.y

@@ -777,9 +777,9 @@ update_statement:
   }
 
 delete_statement:
-  with_clause_opt DELETE comment_opt ignore_opt FROM table_name opt_partition_clause where_expression_opt order_by_opt limit_opt
+  with_clause_opt DELETE comment_opt ignore_opt FROM table_name as_opt_id opt_partition_clause where_expression_opt order_by_opt limit_opt
   {
-    $$ = &Delete{With: $1, Comments: Comments($3), Ignore: $4, TableExprs:  TableExprs{&AliasedTableExpr{Expr:$6}}, Partitions: $7, Where: NewWhere(WhereClause, $8), OrderBy: $9, Limit: $10}
+    $$ = &Delete{With: $1, Comments: Comments($3), Ignore: $4, TableExprs: TableExprs{&AliasedTableExpr{Expr:$6, As: $7}}, Partitions: $8, Where: NewWhere(WhereClause, $9), OrderBy: $10, Limit: $11}
   }
 | with_clause_opt DELETE comment_opt ignore_opt FROM table_name_list USING table_references where_expression_opt
   {

go/vt/vtgate/planbuilder/delete.go

@@ -57,7 +57,12 @@ func buildDeletePlan(stmt sqlparser.Statement, reservedVars *sqlparser.ReservedV
 	}
 
 	if len(edel.Table.Owned) > 0 {
-		edel.OwnedVindexQuery = generateDMLSubquery(del.Where, del.OrderBy, del.Limit, edel.Table, ksidCol)
+		aTblExpr, ok := del.TableExprs[0].(*sqlparser.AliasedTableExpr)
+		if !ok {
+			return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: delete on complex table expression")
+		}
+		tblExpr := &sqlparser.AliasedTableExpr{Expr: sqlparser.TableName{Name: edel.Table.Name}, As: aTblExpr.As}
+		edel.OwnedVindexQuery = generateDMLSubquery(tblExpr, del.Where, del.OrderBy, del.Limit, edel.Table, ksidCol)
 		edel.KsidVindex = ksidVindex
 	}
 

go/vt/vtgate/planbuilder/dml.go

@@ -179,15 +179,15 @@ func buildDMLPlan(vschema ContextVSchema, dmlType string, stmt sqlparser.Stateme
 	return edml, ksidVindex, ksidCol, nil
 }
 
-func generateDMLSubquery(where *sqlparser.Where, orderBy sqlparser.OrderBy, limit *sqlparser.Limit, table *vindexes.Table, ksidCol string) string {
+func generateDMLSubquery(tblExpr sqlparser.TableExpr, where *sqlparser.Where, orderBy sqlparser.OrderBy, limit *sqlparser.Limit, table *vindexes.Table, ksidCol string) string {
 	buf := sqlparser.NewTrackedBuffer(nil)
 	buf.Myprintf("select %s", ksidCol)
 	for _, cv := range table.Owned {
 		for _, column := range cv.Columns {
 			buf.Myprintf(", %v", column)
 		}
 	}
-	buf.Myprintf(" from %v%v%v%v for update", table.Name, where, orderBy, limit)
+	buf.Myprintf(" from %v%v%v%v for update", tblExpr, where, orderBy, limit)
 	return buf.String()
 }
 

go/vt/vtgate/planbuilder/testdata/dml_cases.txt

@@ -519,7 +519,7 @@ Gen4 plan same as above
     "TargetTabletType": "PRIMARY",
     "KsidVindex": "user_index",
     "MultiShardAutocommit": false,
-    "OwnedVindexQuery": "select Id, `Name`, Costly from `user` where id = 1 for update",
+    "OwnedVindexQuery": "select Id, `Name`, Costly from `user` as route1 where id = 1 for update",
     "Query": "delete from `user` as route1 where id = 1",
     "Table": "user",
     "Values": [
@@ -2585,3 +2585,50 @@ Gen4 plan same as above
   }
 }
 Gen4 plan same as above
+
+#update with alias table
+"update user u set u.name = 'john' where u.col > 20"
+{
+  "QueryType": "UPDATE",
+  "Original": "update user u set u.name = 'john' where u.col \u003e 20",
+  "Instructions": {
+    "OperatorType": "Update",
+    "Variant": "Scatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "TargetTabletType": "PRIMARY",
+    "ChangedVindexValues": [
+      "name_user_map:3"
+    ],
+    "KsidVindex": "user_index",
+    "MultiShardAutocommit": false,
+    "OwnedVindexQuery": "select Id, `Name`, Costly, u.`name` = 'john' from `user` as u where u.col \u003e 20 for update",
+    "Query": "update `user` as u set u.`name` = 'john' where u.col \u003e 20",
+    "Table": "user"
+  }
+}
+Gen4 plan same as above
+
+#delete with alias table
+"delete from user u where u.col > 20"
+{
+  "QueryType": "DELETE",
+  "Original": "delete from user u where u.col \u003e 20",
+  "Instructions": {
+    "OperatorType": "Delete",
+    "Variant": "Scatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "TargetTabletType": "PRIMARY",
+    "KsidVindex": "user_index",
+    "MultiShardAutocommit": false,
+    "OwnedVindexQuery": "select Id, `Name`, Costly from `user` as u where u.col \u003e 20 for update",
+    "Query": "delete from `user` as u where u.col \u003e 20",
+    "Table": "user"
+  }
+}
+Gen4 plan same as above

go/vt/vtgate/planbuilder/update.go

@@ -116,7 +116,12 @@ func buildChangedVindexesValues(update *sqlparser.Update, table *vindexes.Table,
 		return nil, "", nil
 	}
 	// generate rest of the owned vindex query.
-	buf.Myprintf(" from %v%v%v%v for update", table.Name, update.Where, update.OrderBy, update.Limit)
+	aTblExpr, ok := update.TableExprs[0].(*sqlparser.AliasedTableExpr)
+	if !ok {
+		return nil, "", vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: update on complex table expression")
+	}
+	tblExpr := &sqlparser.AliasedTableExpr{Expr: sqlparser.TableName{Name: table.Name}, As: aTblExpr.As}
+	buf.Myprintf(" from %v%v%v%v for update", tblExpr, update.Where, update.OrderBy, update.Limit)
 	return changedVindexes, buf.String(), nil
 }