feat: Add support for Insert with row alias (#15790)

Signed-off-by: Andres Taylor <[email protected]> Signed-off-by: Harshit Gangal <[email protected]> Co-authored-by: Andres Taylor <[email protected]>
vitessio · Apr 29, 2024 · ca2659d · ca2659d
1 parent 3eaa99d
commit ca2659d
Show file tree

Hide file tree

Showing 26 changed files with 453 additions and 81 deletions.
diff --git a/changelog/20.0/20.0.0/summary.md b/changelog/20.0/20.0.0/summary.md
@@ -22,6 +22,7 @@
     - [Delete with Subquery Support](#delete-subquery)
     - [Delete with Multi Target Support](#delete-multi-target)
     - [User Defined Functions Support](#udf-support)
+    - [Insert Row Alias Support](#insert-row-alias-support)
   - **[Query Timeout](#query-timeout)**
   - **[Flag changes](#flag-changes)**
     - [`pprof-http` default change](#pprof-http-default)
@@ -198,6 +199,16 @@ Without this flag, VTGate will not be aware that there might be aggregating user
 
 More details about how to load UDFs is available in [MySQL Docs](https://dev.mysql.com/doc/extending-mysql/8.0/en/adding-loadable-function.html)
 
+#### <a id="insert-row-alias-support"/> Insert Row Alias Support
+
+Support is added to have row alias in Insert statement to be used with `on duplicate key update`.
+
+Example:
+- `insert into user(id, name, email) valies (100, 'Alice', '[email protected]') as new on duplicate key update name = new.name, email = new.email`
+- `insert into user(id, name, email) valies (100, 'Alice', '[email protected]') as new(m, n, p) on duplicate key update name = n, email = p`
+
+More details about how it works is available in [MySQL Docs](https://dev.mysql.com/doc/refman/8.0/en/insert-on-duplicate.html)
+
 ### <a id="query-timeout"/>Query Timeout
 On a query timeout, Vitess closed the connection using the `kill connection` statement. This leads to connection churn 
 which is not desirable in some cases. To avoid this, Vitess now uses the `kill query` statement to cancel the query. 

diff --git a/go/mysql/sqlerror/constants.go b/go/mysql/sqlerror/constants.go
@@ -235,6 +235,7 @@ const (
 	ERUnknownTimeZone              = ErrorCode(1298)
 	ERInvalidCharacterString       = ErrorCode(1300)
 	ERQueryInterrupted             = ErrorCode(1317)
+	ERViewWrongList                = ErrorCode(1353)
 	ERTruncatedWrongValueForField  = ErrorCode(1366)
 	ERIllegalValueForType          = ErrorCode(1367)
 	ERDataTooLong                  = ErrorCode(1406)

diff --git a/go/mysql/sqlerror/sql_error.go b/go/mysql/sqlerror/sql_error.go
@@ -216,6 +216,7 @@ var stateToMysqlCode = map[vterrors.State]mysqlCode{
 	vterrors.OperandColumns:               {num: EROperandColumns, state: SSWrongNumberOfColumns},
 	vterrors.WrongValueCountOnRow:         {num: ERWrongValueCountOnRow, state: SSWrongValueCountOnRow},
 	vterrors.WrongArguments:               {num: ERWrongArguments, state: SSUnknownSQLState},
+	vterrors.ViewWrongList:                {num: ERViewWrongList, state: SSUnknownSQLState},
 	vterrors.UnknownStmtHandler:           {num: ERUnknownStmtHandler, state: SSUnknownSQLState},
 	vterrors.KeyDoesNotExist:              {num: ERKeyDoesNotExist, state: SSClientError},
 	vterrors.UnknownTimeZone:              {num: ERUnknownTimeZone, state: SSUnknownSQLState},

diff --git a/go/test/endtoend/vtgate/queries/dml/insert_test.go b/go/test/endtoend/vtgate/queries/dml/insert_test.go
@@ -462,3 +462,27 @@ func TestMixedCases(t *testing.T) {
 	// final check count on the lookup vindex table.
 	utils.AssertMatches(t, mcmp.VtConn, "select count(*) from lkp_mixed_idx", "[[INT64(12)]]")
 }
+
+// TestInsertAlias test the alias feature in insert statement.
+func TestInsertAlias(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 20, "vtgate")
+	utils.SkipIfBinaryIsBelowVersion(t, 20, "vttablet")
+
+	mcmp, closer := start(t)
+	defer closer()
+
+	// initial record
+	mcmp.Exec("insert into user_tbl(id, region_id, name) values (1, 1,'foo'),(2, 2,'bar'),(3, 3,'baz'),(4, 4,'buzz')")
+
+	qr := mcmp.Exec("insert into user_tbl(id, region_id, name) values (2, 2, 'foo') as new on duplicate key update name = new.name")
+	assert.EqualValues(t, 2, qr.RowsAffected)
+
+	// this validates the record.
+	mcmp.Exec("select id, region_id, name from user_tbl order by id")
+
+	qr = mcmp.Exec("insert into user_tbl(id, region_id, name) values (3, 3, 'foo') as new(m, n, p) on duplicate key update name = p")
+	assert.EqualValues(t, 2, qr.RowsAffected)
+
+	// this validates the record.
+	mcmp.Exec("select id, region_id, name from user_tbl order by id")
+}
diff --git a/go/vt/sqlparser/ast.go b/go/vt/sqlparser/ast.go
@@ -341,8 +341,8 @@ type (
 		Partitions Partitions
 		Columns    Columns
 		Rows       InsertRows
-		OnDup      OnDup
 		RowAlias   *RowAlias
+		OnDup      OnDup
 	}
 
 	// Ignore represents whether ignore was specified or not

diff --git a/go/vt/sqlparser/ast_clone.go b/go/vt/sqlparser/ast_clone.go
diff --git a/go/vt/sqlparser/ast_copy_on_rewrite.go b/go/vt/sqlparser/ast_copy_on_rewrite.go
diff --git a/go/vt/sqlparser/ast_equals.go b/go/vt/sqlparser/ast_equals.go
diff --git a/go/vt/sqlparser/ast_rewrite.go b/go/vt/sqlparser/ast_rewrite.go
diff --git a/go/vt/sqlparser/ast_visit.go b/go/vt/sqlparser/ast_visit.go
diff --git a/go/vt/sqlparser/cached_size.go b/go/vt/sqlparser/cached_size.go
diff --git a/go/vt/vterrors/code.go b/go/vt/vterrors/code.go
@@ -58,6 +58,7 @@ var (
 	VT03030 = errorWithState("VT03030", vtrpcpb.Code_INVALID_ARGUMENT, WrongValueCountOnRow, "lookup column count does not match value count with the row (columns, count): (%v, %d)", "The number of columns you want to insert do not match the number of columns of your SELECT query.")
 	VT03031 = errorWithoutState("VT03031", vtrpcpb.Code_INVALID_ARGUMENT, "EXPLAIN is only supported for single keyspace", "EXPLAIN has to be sent down as a single query to the underlying MySQL, and this is not possible if it uses tables from multiple keyspaces")
 	VT03032 = errorWithState("VT03032", vtrpcpb.Code_INVALID_ARGUMENT, NonUpdateableTable, "the target table %s of the UPDATE is not updatable", "You cannot update a table that is not a real MySQL table.")
+	VT03033 = errorWithState("VT03033", vtrpcpb.Code_INVALID_ARGUMENT, ViewWrongList, "In definition of view, derived table or common table expression, SELECT list and column names list have different column counts", "The table column list and derived column list have different column counts.")
 
 	VT05001 = errorWithState("VT05001", vtrpcpb.Code_NOT_FOUND, DbDropExists, "cannot drop database '%s'; database does not exists", "The given database does not exist; Vitess cannot drop it.")
 	VT05002 = errorWithState("VT05002", vtrpcpb.Code_NOT_FOUND, BadDb, "cannot alter database '%s'; unknown database", "The given database does not exist; Vitess cannot alter it.")
@@ -146,6 +147,7 @@ var (
 		VT03030,
 		VT03031,
 		VT03032,
+		VT03033,
 		VT05001,
 		VT05002,
 		VT05003,

diff --git a/go/vt/vterrors/state.go b/go/vt/vterrors/state.go
@@ -49,6 +49,7 @@ const (
 	WrongArguments
 	BadNullError
 	InvalidGroupFuncUse
+	ViewWrongList
 
 	// failed precondition
 	NoDB

diff --git a/go/vt/vtgate/engine/cached_size.go b/go/vt/vtgate/engine/cached_size.go
diff --git a/go/vt/vtgate/engine/insert.go b/go/vt/vtgate/engine/insert.go
@@ -55,6 +55,9 @@ type Insert struct {
 
 	// Mid is the row values for the sharded insert plans.
 	Mid sqlparser.Values
+
+	// Alias represents the row alias with columns if specified in the query.
+	Alias string
 }
 
 // newQueryInsert creates an Insert with a query string.
@@ -287,7 +290,7 @@ func (ins *Insert) getInsertShardedQueries(
 				}
 			}
 		}
-		rewritten := ins.Prefix + strings.Join(mids, ",") + sqlparser.String(ins.Suffix)
+		rewritten := ins.Prefix + strings.Join(mids, ",") + ins.Alias + sqlparser.String(ins.Suffix)
 		queries[i] = &querypb.BoundQuery{
 			Sql:           rewritten,
 			BindVariables: shardBindVars,
@@ -363,6 +366,19 @@ func (ins *Insert) description() PrimitiveDescription {
 		other["VindexValues"] = valuesOffsets
 	}
 
+	// This is a check to ensure we send the correct query to the database.
+	// "ActualQuery" should not be part of the plan output, if it does, it means the query was not rewritten correctly.
+	if ins.Mid != nil {
+		var mids []string
+		for _, n := range ins.Mid {
+			mids = append(mids, sqlparser.String(n))
+		}
+		shardedQuery := ins.Prefix + strings.Join(mids, ", ") + ins.Alias + sqlparser.String(ins.Suffix)
+		if shardedQuery != ins.Query {
+			other["ActualQuery"] = shardedQuery
+		}
+	}
+
 	return PrimitiveDescription{
 		OperatorType:     "Insert",
 		Keyspace:         ins.Keyspace,

diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go
@@ -619,6 +619,9 @@ func buildInsertLogicalPlan(
 	// when unsharded query with autoincrement for that there is no input operator.
 	if eins.Opcode != engine.InsertUnsharded {
 		eins.Prefix, eins.Mid, eins.Suffix = generateInsertShardedQuery(ins.AST)
+		if ins.AST.RowAlias != nil {
+			eins.Alias = sqlparser.String(ins.AST.RowAlias)
+		}
 	}
 
 	eins.Query = generateQuery(stmt)
@@ -660,7 +663,7 @@ func generateInsertShardedQuery(ins *sqlparser.Insert) (prefix string, mids sqlp
 	prefixBuf := sqlparser.NewTrackedBuffer(dmlFormatter)
 	prefixBuf.Myprintf(prefixFormat,
 		ins.Comments, ins.Ignore.ToString(),
-		ins.Table, ins.Columns)
+		ins.Table, ins.Columns, ins.RowAlias)
 	prefix = prefixBuf.String()
 
 	suffix = sqlparser.CopyOnRewrite(ins.OnDup, nil, func(cursor *sqlparser.CopyOnWriteCursor) {

diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases.json b/go/vt/vtgate/planbuilder/testdata/dml_cases.json
@@ -6752,5 +6752,109 @@
         "user.user"
       ]
     }
+  },
+  {
+    "comment": "RowAlias in INSERT",
+    "query": "INSERT INTO authoritative (user_id,col1,col2) VALUES (1,'2',3),(4,'5',6) AS new ON DUPLICATE KEY UPDATE col2 = new.user_id+new.col1",
+    "plan": {
+      "QueryType": "INSERT",
+      "Original": "INSERT INTO authoritative (user_id,col1,col2) VALUES (1,'2',3),(4,'5',6) AS new ON DUPLICATE KEY UPDATE col2 = new.user_id+new.col1",
+      "Instructions": {
+        "OperatorType": "Insert",
+        "Variant": "Sharded",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "TargetTabletType": "PRIMARY",
+        "InsertIgnore": true,
+        "Query": "insert into authoritative(user_id, col1, col2) values (:_user_id_0, '2', 3), (:_user_id_1, '5', 6) as new on duplicate key update col2 = new.user_id + new.col1",
+        "TableName": "authoritative",
+        "VindexValues": {
+          "user_index": "1, 4"
+        }
+      },
+      "TablesUsed": [
+        "user.authoritative"
+      ]
+    }
+  },
+  {
+    "comment": "RowAlias with explicit columns in INSERT",
+    "query": "INSERT INTO authoritative (user_id,col1,col2) VALUES (1,'2',3),(4,'5',6) AS new(a,b,c) ON DUPLICATE KEY UPDATE col1 = a+c",
+    "plan": {
+      "QueryType": "INSERT",
+      "Original": "INSERT INTO authoritative (user_id,col1,col2) VALUES (1,'2',3),(4,'5',6) AS new(a,b,c) ON DUPLICATE KEY UPDATE col1 = a+c",
+      "Instructions": {
+        "OperatorType": "Insert",
+        "Variant": "Sharded",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "TargetTabletType": "PRIMARY",
+        "InsertIgnore": true,
+        "Query": "insert into authoritative(user_id, col1, col2) values (:_user_id_0, '2', 3), (:_user_id_1, '5', 6) as new (a, b, c) on duplicate key update col1 = a + c",
+        "TableName": "authoritative",
+        "VindexValues": {
+          "user_index": "1, 4"
+        }
+      },
+      "TablesUsed": [
+        "user.authoritative"
+      ]
+    }
+  },
+  {
+    "comment": "RowAlias in INSERT (no column list)",
+    "query": "INSERT INTO authoritative VALUES (1,'2',3),(4,'5',6) AS new ON DUPLICATE KEY UPDATE col2 = new.user_id+new.col1",
+    "plan": {
+      "QueryType": "INSERT",
+      "Original": "INSERT INTO authoritative VALUES (1,'2',3),(4,'5',6) AS new ON DUPLICATE KEY UPDATE col2 = new.user_id+new.col1",
+      "Instructions": {
+        "OperatorType": "Insert",
+        "Variant": "Sharded",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "TargetTabletType": "PRIMARY",
+        "InsertIgnore": true,
+        "Query": "insert into authoritative(user_id, col1, col2) values (:_user_id_0, '2', 3), (:_user_id_1, '5', 6) as new on duplicate key update col2 = new.user_id + new.col1",
+        "TableName": "authoritative",
+        "VindexValues": {
+          "user_index": "1, 4"
+        }
+      },
+      "TablesUsed": [
+        "user.authoritative"
+      ]
+    }
+  },
+  {
+    "comment": "RowAlias with explicit columns in INSERT (no column list)",
+    "query": "INSERT INTO authoritative VALUES (1,'2',3),(4,'5',6) AS new(a,b,c) ON DUPLICATE KEY UPDATE col1 = a+c",
+    "plan": {
+      "QueryType": "INSERT",
+      "Original": "INSERT INTO authoritative VALUES (1,'2',3),(4,'5',6) AS new(a,b,c) ON DUPLICATE KEY UPDATE col1 = a+c",
+      "Instructions": {
+        "OperatorType": "Insert",
+        "Variant": "Sharded",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "TargetTabletType": "PRIMARY",
+        "InsertIgnore": true,
+        "Query": "insert into authoritative(user_id, col1, col2) values (:_user_id_0, '2', 3), (:_user_id_1, '5', 6) as new (a, b, c) on duplicate key update col1 = a + c",
+        "TableName": "authoritative",
+        "VindexValues": {
+          "user_index": "1, 4"
+        }
+      },
+      "TablesUsed": [
+        "user.authoritative"
+      ]
+    }
   }
 ]