feat: handle last_insert_id with arguments in the evalengine

Signed-off-by: Andres Taylor <[email protected]>

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

@@ -163,6 +163,7 @@ 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",
 	}
 
 	for _, workload := range []string{"olap", "oltp"} {
@@ -175,7 +176,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 {

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -893,6 +893,9 @@ func (t *loggingVCursor) RecordMirrorStats(sourceExecTime, targetExecTime time.D
 	}
 }
 
+func (t *loggingVCursor) SetLastInsertID(id uint64) {}
+func (t *noopVCursor) SetLastInsertID(id uint64)    {}
+
 func (t *noopVCursor) VExplainLogging() {}
 func (t *noopVCursor) DisableLogging()  {}
 func (t *noopVCursor) GetVExplainLogs() []ExecuteEntry {

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

go/vt/vtgate/evalengine/compiler_asm.go

@@ -5138,3 +5138,24 @@ 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].(*evalUint64)
+		env.VCursor().SetLastInsertID(arg.u)
+		return 1
+	}, "FN LAST_INSERT_ID UINT64(SP-1)")
+}
+
+func (asm *assembler) Fn_LAST_INSERT_ID_NULL() {
+	asm.emit(func(env *ExpressionEnv) int {
+		env.VCursor().SetLastInsertID(0)
+		return 1
+	}, "FN LAST_INSERT_ID NULL")
+}
+
+func (asm *assembler) addJump(end *jump) {
+	asm.emit(func(env *ExpressionEnv) int {
+		return end.offset()
+	}, "JUMP")
+}

go/vt/vtgate/evalengine/compiler_test.go

@@ -883,6 +883,99 @@ func TestBindVarLiteral(t *testing.T) {
 	}
 }
 
+type testVcursor struct {
+	lastInsertID *uint64
+	env          *vtenv.Environment
+}
+
+func (t *testVcursor) TimeZone() *time.Location {
+	return time.UTC
+}
+
+func (t *testVcursor) GetKeyspace() string {
+	return "apa"
+}
+
+func (t *testVcursor) SQLMode() string {
+	return "oltp"
+}
+
+func (t *testVcursor) Environment() *vtenv.Environment {
+	return t.env
+}
+
+func (t *testVcursor) SetLastInsertID(id uint64) {
+	t.lastInsertID = &id
+}
+
+var _ evalengine.VCursor = (*testVcursor)(nil)
+
+func TestLastInsertID(t *testing.T) {
+	var testCases = []struct {
+		expression string
+		result     uint64
+		missing    bool
+	}{
+		{
+			expression: `last_insert_id(1)`,
+			result:     1,
+		}, {
+			expression: `12`,
+			missing:    true,
+		}, {
+			expression: `last_insert_id(666)`,
+			result:     666,
+		}, {
+			expression: `last_insert_id(null)`,
+			result:     0,
+		},
+	}
+
+	venv := vtenv.NewTestEnv()
+	for _, tc := range testCases {
+		t.Run(tc.expression, func(t *testing.T) {
+			expr, err := venv.Parser().ParseExpr(tc.expression)
+			require.NoError(t, err)
+
+			cfg := &evalengine.Config{
+				Collation:         collations.CollationUtf8mb4ID,
+				NoConstantFolding: true,
+				NoCompilation:     false,
+				Environment:       venv,
+			}
+			t.Run("eval", func(t *testing.T) {
+				cfg.NoCompilation = true
+				runTest(t, expr, cfg, tc)
+			})
+			t.Run("compiled", func(t *testing.T) {
+				cfg.NoCompilation = false
+				runTest(t, expr, cfg, tc)
+			})
+		})
+	}
+}
+
+func runTest(t *testing.T, expr sqlparser.Expr, cfg *evalengine.Config, tc struct {
+	expression string
+	result     uint64
+	missing    bool
+}) {
+	converted, err := evalengine.Translate(expr, cfg)
+	require.NoError(t, err)
+
+	vc := &testVcursor{env: vtenv.NewTestEnv()}
+	env := evalengine.NewExpressionEnv(context.Background(), nil, vc)
+
+	_, err = env.Evaluate(converted)
+	require.NoError(t, err)
+	if tc.missing {
+		require.Nil(t, vc.lastInsertID)
+	} else {
+		require.NotNil(t, vc.lastInsertID)
+		require.Equal(t, tc.result, *vc.lastInsertID)
+	}
+}
+
 func TestCompilerNonConstant(t *testing.T) {
 	var testCases = []struct {
 		expression string

go/vt/vtgate/evalengine/expr_env.go

@@ -35,6 +35,7 @@ type VCursor interface {
 	GetKeyspace() string
 	SQLMode() string
 	Environment() *vtenv.Environment
+	SetLastInsertID(id uint64)
 }
 
 type (
@@ -140,6 +141,7 @@ func (e *emptyVCursor) GetKeyspace() string {
 func (e *emptyVCursor) SQLMode() string {
 	return config.DefaultSQLMode
 }
+func (e *emptyVCursor) SetLastInsertID(_ uint64) {}
 
 func NewEmptyVCursor(env *vtenv.Environment, tz *time.Location) VCursor {
 	return &emptyVCursor{env: env, tz: tz}

go/vt/vtgate/evalengine/fn_misc.go

@@ -81,6 +81,10 @@ type (
 	builtinUUIDToBin struct {
 		CallExpr
 	}
+
+	builtinLastInsertID struct {
+		CallExpr
+	}
 )
 
 var _ IR = (*builtinInetAton)(nil)
@@ -95,6 +99,7 @@ var _ IR = (*builtinBinToUUID)(nil)
 var _ IR = (*builtinIsUUID)(nil)
 var _ IR = (*builtinUUID)(nil)
 var _ IR = (*builtinUUIDToBin)(nil)
+var _ IR = (*builtinLastInsertID)(nil)
 
 func (call *builtinInetAton) eval(env *ExpressionEnv) (eval, error) {
 	arg, err := call.arg1(env)
@@ -155,6 +160,7 @@ func (call *builtinInetNtoa) compile(c *compiler) (ctype, error) {
 	c.compileToUint64(arg, 1)
 	col := typedCoercionCollation(sqltypes.VarChar, call.collate)
 	c.asm.Fn_INET_NTOA(col)
+
 	c.asm.jumpDestination(skip)
 
 	return ctype{Type: sqltypes.VarChar, Flag: flagNullable, Col: col}, nil
@@ -194,6 +200,40 @@ func (call *builtinInet6Aton) compile(c *compiler) (ctype, error) {
 	return ctype{Type: sqltypes.VarBinary, Flag: flagNullable, Col: collationBinary}, nil
 }
 
+func (call *builtinLastInsertID) eval(env *ExpressionEnv) (eval, error) {
+	arg, err := call.arg1(env)
+	if err != nil {
+		return nil, err
+	}
+	if arg == nil {
+		env.VCursor().SetLastInsertID(0)
+		return nil, err
+	}
+	insertID := uint64(evalToInt64(arg).i)
+	env.VCursor().SetLastInsertID(insertID)
+	return newEvalUint64(insertID), nil
+}
+
+func (call *builtinLastInsertID) compile(c *compiler) (ctype, error) {
+	arg, err := call.Arguments[0].compile(c)
+	if err != nil {
+		return ctype{}, err
+	}
+
+	setZero := c.compileNullCheck1(arg)
+	c.compileToUint64(arg, 1)
+	c.asm.Fn_LAST_INSERT_ID()
+	end := c.asm.jumpFrom()
+	c.asm.addJump(end)
+
+	c.asm.jumpDestination(setZero)
+	c.asm.Fn_LAST_INSERT_ID_NULL()
+
+	c.asm.jumpDestination(end)
+
+	return ctype{Type: sqltypes.Uint64, Flag: flagNullable, Col: collationNumeric}, nil
+}
+
 func printIPv6AsIPv4(addr netip.Addr) (netip.Addr, bool) {
 	b := addr.AsSlice()
 	if len(b) != 16 {

go/vt/vtgate/evalengine/integration/comparison_test.go

@@ -209,6 +209,10 @@ type vcursor struct {
 	env *vtenv.Environment
 }
 
+func (vc *vcursor) SetLastInsertID(id uint64) {}
+
+var _ evalengine.VCursor = (*vcursor)(nil)
+
 func (vc *vcursor) GetKeyspace() string {
 	return "vttest"
 }

go/vt/vtgate/evalengine/translate_builtin.go

@@ -662,6 +662,11 @@ func (ast *astCompiler) translateFuncExpr(fn *sqlparser.FuncExpr) (IR, error) {
 			return nil, argError(method)
 		}
 		return &builtinReplace{CallExpr: call, collate: ast.cfg.Collation}, nil
+	case "last_insert_id":
+		if len(args) != 1 {
+			return nil, argError(method)
+		}
+		return &builtinLastInsertID{CallExpr: call}, nil
 	default:
 		return nil, translateExprNotSupported(fn)
 	}

go/vt/vtgate/executorcontext/vcursor_impl.go

@@ -1569,3 +1569,9 @@ func (vc *VCursorImpl) GetContextWithTimeOut(ctx context.Context) (context.Conte
 func (vc *VCursorImpl) IgnoreMaxMemoryRows() bool {
 	return vc.ignoreMaxMemoryRows
 }
+
+func (vc *VCursorImpl) SetLastInsertID(id uint64) {
+	vc.SafeSession.mu.Lock()
+	defer vc.SafeSession.mu.Unlock()
+	vc.SafeSession.LastInsertId = id
+}