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

@@ -162,6 +162,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"} {
@@ -174,7 +175,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,11 @@ 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)")
+}

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)
@@ -194,6 +199,32 @@ 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 arg == nil || err != nil {
+		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
+	}
+
+	skip := c.compileNullCheck1(arg)
+
+	c.compileToUint64(arg, 1)
+	c.asm.Fn_LAST_INSERT_ID()
+
+	c.asm.jumpDestination(skip)
+
+	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/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
+}