Improved Compatibility Around LAST_INSERT_ID - evalengine (#17409)

Signed-off-by: Andres Taylor <[email protected]> Signed-off-by: Vicent Marti <[email protected]> Signed-off-by: Harshit Gangal <[email protected]> Signed-off-by: Florent Poinsard <[email protected]> Co-authored-by: Vicent Marti <[email protected]> Co-authored-by: Harshit Gangal <[email protected]> Co-authored-by: Florent Poinsard <[email protected]>
vitessio · Jan 8, 2025 · 6ac8998 · 6ac8998
1 parent 72d91bd
commit 6ac8998
Show file tree

Hide file tree

Showing 26 changed files with 403 additions and 45 deletions.
diff --git a/changelog/22.0/22.0.0/summary.md b/changelog/22.0/22.0.0/summary.md
@@ -10,6 +10,7 @@
   - **[VTOrc Config File Changes](#vtorc-config-file-changes)**
   - **[VTGate Config File Changes](#vtgate-config-file-changes)**
   - **[Support for More Efficient JSON Replication](#efficient-json-replication)**
+  - **[Support for LAST_INSERT_ID(x)](#last-insert-id)**
 - **[Minor Changes](#minor-changes)**
   - **[VTTablet Flags](#flags-vttablet)**
   - **[Topology read concurrency behaviour changes](#topo-read-concurrency-changes)**
@@ -80,6 +81,12 @@ In [#7345](https://github.com/vitessio/vitess/pull/17345) we added support for [
 
 If you are using MySQL 8.0 or later and using JSON columns, you can now enable this MySQL feature across your Vitess cluster(s) to lower the disk space needed for binary logs and improve the CPU and memory usage in both `mysqld` (standard intrashard MySQL replication) and `vttablet` ([VReplication](https://vitess.io/docs/reference/vreplication/vreplication/)) without losing any capabilities or features.
 
+### <a id="last-insert-id"/>Support for `LAST_INSERT_ID(x)`</a>
+
+In [#17408](https://github.com/vitessio/vitess/pull/17408) and [#17409](https://github.com/vitessio/vitess/pull/17409), we added the ability to use `LAST_INSERT_ID(x)` in Vitess directly at vtgate. This improvement allows certain queries—like `SELECT last_insert_id(123);` or `SELECT last_insert_id(count(*)) ...`—to be handled without relying on MySQL for the final value.
+
+**Limitations**:
+- When using `LAST_INSERT_ID(x)` in ordered queries (e.g., `SELECT last_insert_id(col) FROM table ORDER BY foo`), MySQL sets the session’s last-insert-id value according to the *last row returned*. Vitess does not guarantee the same behavior.
 
 ## <a id="minor-changes"/>Minor Changes</a>
 

diff --git a/go/test/endtoend/utils/cmp.go b/go/test/endtoend/utils/cmp.go
@@ -215,8 +215,8 @@ func (mcmp *MySQLCompare) Exec(query string) *sqltypes.Result {
 	return vtQr
 }
 
-// ExecVitessAndMySQL executes Vitess and MySQL with the queries provided.
-func (mcmp *MySQLCompare) ExecVitessAndMySQL(vtQ, mQ string) *sqltypes.Result {
+// ExecVitessAndMySQLDifferentQueries executes Vitess and MySQL with the queries provided.
+func (mcmp *MySQLCompare) ExecVitessAndMySQLDifferentQueries(vtQ, mQ string) *sqltypes.Result {
 	mcmp.t.Helper()
 	vtQr, err := mcmp.VtConn.ExecuteFetch(vtQ, 1000, true)
 	require.NoError(mcmp.t, err, "[Vitess Error] for query: "+vtQ)

diff --git a/go/test/endtoend/vtgate/plan_tests/plan_e2e_test.go b/go/test/endtoend/vtgate/plan_tests/plan_e2e_test.go
@@ -49,7 +49,7 @@ func TestE2ECases(t *testing.T) {
 				require.NoError(mcmp.AsT(), err)
 				sqlparser.RemoveKeyspaceIgnoreSysSchema(stmt)
 
-				mcmp.ExecVitessAndMySQL(test.Query, sqlparser.String(stmt))
+				mcmp.ExecVitessAndMySQLDifferentQueries(test.Query, sqlparser.String(stmt))
 				pd := utils.ExecTrace(mcmp.AsT(), mcmp.VtConn, test.Query)
 				verifyTestExpectations(mcmp.AsT(), pd, test)
 				if mcmp.VtConn.IsClosed() {

diff --git a/go/test/endtoend/vtgate/queries/misc/main_test.go b/go/test/endtoend/vtgate/queries/misc/main_test.go
@@ -95,7 +95,7 @@ func TestMain(m *testing.M) {
 		vtParams = clusterInstance.GetVTParams(keyspaceName)
 
 		// create mysql instance and connection parameters
-		conn, closer, err := utils.NewMySQL(clusterInstance, keyspaceName, schemaSQL)
+		conn, closer, err := utils.NewMySQL(clusterInstance, keyspaceName, schemaSQL, uschemaSQL)
 		if err != nil {
 			fmt.Println(err)
 			return 1

diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go
@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/vt/sqlparser"
 
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/stretchr/testify/assert"
@@ -163,6 +164,8 @@ func TestSetAndGetLastInsertID(t *testing.T) {
 		"update t1 set id2 = last_insert_id(%d) where id1 = 2",
 		"update t1 set id2 = 88 where id1 = last_insert_id(%d)",
 		"delete from t1 where id1 = last_insert_id(%d)",
+		"select id2, last_insert_id(count(*)) from t1 where %d group by id2",
+		"set @x = last_insert_id(%d)",
 	}
 
 	for _, workload := range []string{"olap", "oltp"} {
@@ -175,7 +178,7 @@ func TestSetAndGetLastInsertID(t *testing.T) {
 				require.NoError(t, err)
 			}
 
-			// Insert a row for UPDATE tests
+			// Insert a few rows for UPDATE tests
 			mcmp.Exec("insert into t1 (id1, id2) values (1, 10)")
 
 			for _, query := range queries {
@@ -186,6 +189,96 @@ func TestSetAndGetLastInsertID(t *testing.T) {
 	}
 }
 
+func TestSetAndGetLastInsertIDWithInsertUnsharded(t *testing.T) {
+	mcmp, closer := start(t)
+	defer closer()
+
+	tests := []string{
+		"insert into uks.unsharded(id1, id2) values (last_insert_id(%d),12)",
+		"insert into uks.unsharded(id1, id2) select last_insert_id(%d), 453",
+	}
+
+	i := 0
+	getVal := func() int {
+		defer func() { i++ }()
+		return i
+	}
+
+	runTests := func(mcmp *utils.MySQLCompare) {
+		for _, test := range tests {
+
+			lastInsertID := getVal()
+			query := fmt.Sprintf(test, lastInsertID)
+
+			stmt, err := sqlparser.NewTestParser().Parse(query)
+			require.NoError(mcmp.AsT(), err)
+			sqlparser.RemoveKeyspaceIgnoreSysSchema(stmt)
+
+			mcmp.ExecVitessAndMySQLDifferentQueries(query, sqlparser.String(stmt))
+			mcmp.Exec("select last_insert_id()")
+		}
+	}
+
+	for _, workload := range []string{"olap", "oltp"} {
+		mcmp.Run(workload, func(mcmp *utils.MySQLCompare) {
+			_, err := mcmp.VtConn.ExecuteFetch("set workload = "+workload, 1, false)
+			require.NoError(t, err)
+			runTests(mcmp)
+
+			// run the queries again, but inside a transaction this time
+			mcmp.Exec("begin")
+			runTests(mcmp)
+			mcmp.Exec("commit")
+		})
+	}
+
+	// Now test to set the last insert id to 0, see that it has changed correctly even if the value is 0
+	mcmp.ExecVitessAndMySQLDifferentQueries(
+		"insert into uks.unsharded(id1, id2) values (last_insert_id(0),12)",
+		"insert into unsharded(id1, id2) values (last_insert_id(0),12)",
+	)
+	mcmp.Exec("select last_insert_id()")
+}
+
+func TestSetAndGetLastInsertIDWithInsert(t *testing.T) {
+	mcmp, closer := start(t)
+	defer closer()
+
+	tests := []string{
+		"insert into t1(id1, id2) values (last_insert_id(%d) ,%d)",
+		"insert into t1(id1, id2) values (%d, last_insert_id(%d))",
+		"insert into t1(id1, id2) select last_insert_id(%d), %d",
+		"insert into t1(id1, id2) select last_insert_id(id1+%d), 12 from t1 where 1 > %d",
+	}
+
+	i := 0
+	getVal := func() int {
+		defer func() { i++ }()
+		return i
+	}
+
+	runTests := func(mcmp *utils.MySQLCompare) {
+		for _, test := range tests {
+			query := fmt.Sprintf(test, getVal(), getVal())
+			mcmp.Exec(query)
+			mcmp.Exec("select last_insert_id()")
+		}
+	}
+
+	for _, workload := range []string{"olap", "oltp"} {
+		mcmp.Run(workload, func(mcmp *utils.MySQLCompare) {
+			_, err := mcmp.VtConn.ExecuteFetch("set workload = "+workload, 1, false)
+			require.NoError(t, err)
+			runTests(mcmp)
+
+			// run the queries again, but inside a transaction this time
+			mcmp.Exec("begin")
+			runTests(mcmp)
+			mcmp.Exec("commit")
+		})
+	}
+}
+
 // TestVindexHints tests that vindex hints work as intended.
 func TestVindexHints(t *testing.T) {
 	mcmp, closer := start(t)

diff --git a/go/vt/sqlparser/ast_test.go b/go/vt/sqlparser/ast_test.go
@@ -917,3 +917,11 @@ func TestCloneComments(t *testing.T) {
 		assert.Equal(t, "b", val)
 	}
 }
+
+func TestRemoveKeyspace(t *testing.T) {
+	stmt, err := NewTestParser().Parse("select 1 from uks.unsharded")
+	require.NoError(t, err)
+	RemoveKeyspaceIgnoreSysSchema(stmt)
+
+	require.Equal(t, "select 1 from unsharded", String(stmt))
+}
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/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go
@@ -400,6 +400,20 @@ func (t *noopVCursor) GetDBDDLPluginName() string {
 	panic("unimplemented")
 }
 
+func (t *noopVCursor) SetLastInsertID(uint64) {}
+func (t *noopVCursor) VExplainLogging()       {}
+func (t *noopVCursor) DisableLogging()        {}
+func (t *noopVCursor) GetVExplainLogs() []ExecuteEntry {
+	return nil
+}
+func (t *noopVCursor) GetLogs() ([]ExecuteEntry, error) {
+	return nil, nil
+}
+
+// RecordMirrorStats implements VCursor.
+func (t *noopVCursor) RecordMirrorStats(sourceExecTime, targetExecTime time.Duration, targetErr error) {
+}
+
 var (
 	_ VCursor        = (*loggingVCursor)(nil)
 	_ SessionActions = (*loggingVCursor)(nil)
@@ -893,20 +907,6 @@ func (t *loggingVCursor) RecordMirrorStats(sourceExecTime, targetExecTime time.D
 	}
 }
 
-func (t *noopVCursor) VExplainLogging() {}
-func (t *noopVCursor) DisableLogging()  {}
-func (t *noopVCursor) GetVExplainLogs() []ExecuteEntry {
-	return nil
-}
-
-func (t *noopVCursor) GetLogs() ([]ExecuteEntry, error) {
-	return nil, nil
-}
-
-// RecordMirrorStats implements VCursor.
-func (t *noopVCursor) RecordMirrorStats(sourceExecTime, targetExecTime time.Duration, targetErr error) {
-}
-
 func expectResult(t *testing.T, result, want *sqltypes.Result) {
 	t.Helper()
 	fieldsResult := fmt.Sprintf("%v", result.Fields)

diff --git a/go/vt/vtgate/engine/insert.go b/go/vt/vtgate/engine/insert.go
@@ -58,11 +58,9 @@ type Insert struct {
 
 	// Alias represents the row alias with columns if specified in the query.
 	Alias string
-
-	FetchLastInsertID bool
 }
 
-// newQueryInsert creates an Insert with a query string.
+// newQueryInsert creates an Insert with a query string. Used in testing.
 func newQueryInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, query string) *Insert {
 	return &Insert{
 		InsertCommon: InsertCommon{
@@ -73,7 +71,7 @@ func newQueryInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, query stri
 	}
 }
 
-// newInsert creates a new Insert.
+// newInsert creates a new Insert. Used in testing.
 func newInsert(
 	opcode InsertOpcode,
 	ignore bool,

diff --git a/go/vt/vtgate/engine/insert_common.go b/go/vt/vtgate/engine/insert_common.go
@@ -161,7 +161,7 @@ func (ins *InsertCommon) executeUnshardedTableQuery(ctx context.Context, vcursor
 	if err != nil {
 		return nil, err
 	}
-	qr, err := execShard(ctx, loggingPrimitive, vcursor, query, bindVars, rss[0], true, !ins.PreventAutoCommit /* canAutocommit */, false)
+	qr, err := execShard(ctx, loggingPrimitive, vcursor, query, bindVars, rss[0], true, !ins.PreventAutoCommit /* canAutocommit */, ins.FetchLastInsertID)
 	if err != nil {
 		return nil, err
 	}

diff --git a/go/vt/vtgate/engine/insert_select.go b/go/vt/vtgate/engine/insert_select.go
@@ -51,7 +51,7 @@ type (
 	}
 )
 
-// newInsertSelect creates a new InsertSelect.
+// newInsertSelect creates a new InsertSelect. Used in testing.
 func newInsertSelect(
 	ignore bool,
 	keyspace *vindexes.Keyspace,

diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go
@@ -147,6 +147,8 @@ type (
 
 		// RecordMirrorStats is used to record stats about a mirror query.
 		RecordMirrorStats(time.Duration, time.Duration, error)
+
+		SetLastInsertID(uint64)
 	}
 
 	// SessionActions gives primitives ability to interact with the session state

diff --git a/go/vt/vtgate/engine/set.go b/go/vt/vtgate/engine/set.go
@@ -22,23 +22,18 @@ import (
 	"fmt"
 	"strings"
 
-	"vitess.io/vitess/go/vt/sqlparser"
-	"vitess.io/vitess/go/vt/sysvars"
-
-	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
-
-	"vitess.io/vitess/go/vt/log"
-
-	"vitess.io/vitess/go/vt/srvtopo"
-
-	"vitess.io/vitess/go/vt/vtgate/evalengine"
-
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/key"
+	"vitess.io/vitess/go/vt/log"
 	querypb "vitess.io/vitess/go/vt/proto/query"
+	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/schema"
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/srvtopo"
+	"vitess.io/vitess/go/vt/sysvars"
 	"vitess.io/vitess/go/vt/vterrors"
+	"vitess.io/vitess/go/vt/vtgate/evalengine"
 	"vitess.io/vitess/go/vt/vtgate/vindexes"
 )
 

diff --git a/go/vt/vtgate/evalengine/cached_size.go b/go/vt/vtgate/evalengine/cached_size.go
diff --git a/go/vt/vtgate/evalengine/compiler_asm.go b/go/vt/vtgate/evalengine/compiler_asm.go
@@ -5138,3 +5138,18 @@ func (asm *assembler) Introduce(offset int, t sqltypes.Type, col collations.Type
 		return 1
 	}, "INTRODUCE (SP-1)")
 }
+
+func (asm *assembler) Fn_LAST_INSERT_ID() {
+	asm.emit(func(env *ExpressionEnv) int {
+		arg := env.vm.stack[env.vm.sp-1]
+		if arg == nil {
+			env.VCursor().SetLastInsertID(0)
+		} else {
+			iarg := evalToInt64(arg)
+			uarg := env.vm.arena.newEvalUint64(uint64(iarg.i))
+			env.vm.stack[env.vm.sp-1] = uarg
+			env.VCursor().SetLastInsertID(uarg.u)
+		}
+		return 1
+	}, "FN LAST_INSERT_ID UINT64(SP-1)")
+}