Improved Compatibility Around LAST_INSERT_ID

go/mysql/conn.go

@@ -399,7 +399,7 @@ func (c *Conn) readHeaderFrom(r io.Reader) (int, error) {
 		return 0, vterrors.Wrapf(err, "io.ReadFull(header size) failed")
 	}
 
-	sequence := uint8(c.header[3])
+	sequence := c.header[3]
 	if sequence != c.sequence {
 		return 0, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "invalid sequence, expected %v got %v", c.sequence, sequence)
 	}

go/mysql/query.go

@@ -404,6 +404,7 @@ func (c *Conn) ReadQueryResult(maxrows int, wantfields bool) (*sqltypes.Result,
 		return &sqltypes.Result{
 			RowsAffected:        packetOk.affectedRows,
 			InsertID:            packetOk.lastInsertID,
+			InsertIDChanged:     packetOk.lastInsertID > 0,
 			SessionStateChanges: packetOk.sessionStateData,
 			StatusFlags:         packetOk.statusFlags,
 			Info:                packetOk.info,

go/sqltypes/proto3.go

@@ -103,6 +103,7 @@ func ResultToProto3(qr *Result) *querypb.QueryResult {
 		Fields:              qr.Fields,
 		RowsAffected:        qr.RowsAffected,
 		InsertId:            qr.InsertID,
+		InsertIdChanged:     qr.InsertIDChanged,
 		Rows:                RowsToProto3(qr.Rows),
 		Info:                qr.Info,
 		SessionStateChanges: qr.SessionStateChanges,
@@ -119,6 +120,7 @@ func Proto3ToResult(qr *querypb.QueryResult) *Result {
 		Fields:              qr.Fields,
 		RowsAffected:        qr.RowsAffected,
 		InsertID:            qr.InsertId,
+		InsertIDChanged:     qr.InsertIdChanged,
 		Rows:                proto3ToRows(qr.Fields, qr.Rows),
 		Info:                qr.Info,
 		SessionStateChanges: qr.SessionStateChanges,
@@ -136,6 +138,7 @@ func CustomProto3ToResult(fields []*querypb.Field, qr *querypb.QueryResult) *Res
 		Fields:              qr.Fields,
 		RowsAffected:        qr.RowsAffected,
 		InsertID:            qr.InsertId,
+		InsertIDChanged:     qr.InsertIdChanged,
 		Rows:                proto3ToRows(fields, qr.Rows),
 		Info:                qr.Info,
 		SessionStateChanges: qr.SessionStateChanges,

go/sqltypes/result.go

@@ -31,6 +31,7 @@ type Result struct {
 	Fields              []*querypb.Field `json:"fields"`
 	RowsAffected        uint64           `json:"rows_affected"`
 	InsertID            uint64           `json:"insert_id"`
+	InsertIDChanged     bool             `json:"insert_id_changed"`
 	Rows                []Row            `json:"rows"`
 	SessionStateChanges string           `json:"session_state_changes"`
 	StatusFlags         uint16           `json:"status_flags"`
@@ -92,6 +93,7 @@ func (result *Result) Copy() *Result {
 	out := &Result{
 		RowsAffected:        result.RowsAffected,
 		InsertID:            result.InsertID,
+		InsertIDChanged:     result.InsertIDChanged,
 		SessionStateChanges: result.SessionStateChanges,
 		StatusFlags:         result.StatusFlags,
 		Info:                result.Info,
@@ -116,6 +118,7 @@ func (result *Result) ShallowCopy() *Result {
 	return &Result{
 		Fields:              result.Fields,
 		InsertID:            result.InsertID,
+		InsertIDChanged:     result.InsertIDChanged,
 		RowsAffected:        result.RowsAffected,
 		Info:                result.Info,
 		SessionStateChanges: result.SessionStateChanges,
@@ -129,6 +132,7 @@ func (result *Result) Metadata() *Result {
 	return &Result{
 		Fields:              result.Fields,
 		InsertID:            result.InsertID,
+		InsertIDChanged:     result.InsertIDChanged,
 		RowsAffected:        result.RowsAffected,
 		Info:                result.Info,
 		SessionStateChanges: result.SessionStateChanges,
@@ -153,6 +157,7 @@ func (result *Result) Truncate(l int) *Result {
 
 	out := &Result{
 		InsertID:            result.InsertID,
+		InsertIDChanged:     result.InsertIDChanged,
 		RowsAffected:        result.RowsAffected,
 		Info:                result.Info,
 		SessionStateChanges: result.SessionStateChanges,
@@ -198,6 +203,7 @@ func (result *Result) Equal(other *Result) bool {
 	return FieldsEqual(result.Fields, other.Fields) &&
 		result.RowsAffected == other.RowsAffected &&
 		result.InsertID == other.InsertID &&
+		result.InsertIDChanged == other.InsertIDChanged &&
 		slices.EqualFunc(result.Rows, other.Rows, func(a, b Row) bool {
 			return RowEqual(a, b)
 		})
@@ -324,16 +330,14 @@ func (result *Result) StripMetadata(incl querypb.ExecuteOptions_IncludedFields)
 // to another result.Note currently it doesn't handle cases like
 // if two results have different fields.We will enhance this function.
 func (result *Result) AppendResult(src *Result) {
-	if src.RowsAffected == 0 && len(src.Rows) == 0 && len(src.Fields) == 0 {
-		return
+	result.RowsAffected += src.RowsAffected
+	if src.InsertID != 0 || src.InsertIDChanged {
+		result.InsertID = src.InsertID
 	}
+	result.InsertIDChanged = result.InsertIDChanged || src.InsertIDChanged
 	if result.Fields == nil {
 		result.Fields = src.Fields
 	}
-	result.RowsAffected += src.RowsAffected
-	if src.InsertID != 0 {
-		result.InsertID = src.InsertID
-	}
 	result.Rows = append(result.Rows, src.Rows...)
 }
 

go/test/endtoend/utils/cmp.go

@@ -313,3 +313,10 @@ func (mcmp *MySQLCompare) ExecAllowError(query string) (*sqltypes.Result, error)
 	}
 	return vtQr, vtErr
 }
+
+func (mcmp *MySQLCompare) VExplain(query string) string {
+	mcmp.t.Helper()
+	vtQr, vtErr := mcmp.VtConn.ExecuteFetch("vexplain plan "+query, 1, true)
+	require.NoError(mcmp.t, vtErr, "[Vitess Error] for query: "+query)
+	return vtQr.Rows[0][0].ToString()
+}

go/test/endtoend/vtgate/queries/misc/misc_test.go

@@ -131,6 +131,62 @@ func TestCast(t *testing.T) {
 	mcmp.AssertMatches("select cast('3.2' as unsigned)", `[[UINT64(3)]]`)
 }
 
+// TestSetAndGetLastInsertID tests that the last_insert_id function works as intended when used with different arguments.
+func TestSetAndGetLastInsertID(t *testing.T) {
+	notZero := 1
+	checkQuery := func(i string, workload string, tx bool, mcmp utils.MySQLCompare) {
+		for _, val := range []int{0, notZero} {
+			query := fmt.Sprintf(i, val)
+			name := fmt.Sprintf("%s - %s", workload, query)
+			if tx {
+				name = "tx - " + name
+			}
+			mcmp.Run(name, func(mcmp *utils.MySQLCompare) {
+				mcmp.Exec(query)
+				mcmp.Exec("select last_insert_id()")
+				t := mcmp.AsT()
+				if t.Failed() {
+					t.Log(mcmp.VExplain(query))
+				}
+			})
+		}
+		// we need this value to be not zero, and then we keep changing it so different queries don't interact with each other
+		notZero++
+	}
+
+	queries := []string{
+		"select last_insert_id(%d)",
+		"select last_insert_id(%d), id1, id2 from t1 limit 1",
+		"select last_insert_id(%d), id1, id2 from t1 where 1 = 2",
+		"select 12 from t1 where last_insert_id(%d)",
+		"update t1 set id2 = last_insert_id(%d) where id1 = 1",
+		"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",
+	}
+
+	for _, workload := range []string{"olap", "oltp"} {
+		for _, tx := range []bool{true, false} {
+			mcmp, closer := start(t)
+			_, err := mcmp.VtConn.ExecuteFetch(fmt.Sprintf("set workload = %s", workload), 1000, false)
+			require.NoError(t, err)
+			if tx {
+				_, err := mcmp.VtConn.ExecuteFetch("begin", 1000, false)
+				require.NoError(t, err)
+			}
+
+			// Insert a few rows for UPDATE tests
+			mcmp.Exec("insert into t1 (id1, id2) values (1, 10)")
+
+			for _, query := range queries {
+				checkQuery(query, workload, tx, mcmp)
+			}
+			closer()
+		}
+	}
+}
+
 // TestVindexHints tests that vindex hints work as intended.
 func TestVindexHints(t *testing.T) {
 	mcmp, closer := start(t)