From 6fd09ebafcf5284f9f716bf62957009b58f2fa6b Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Wed, 18 Dec 2024 12:36:30 +0100 Subject: [PATCH] feat: handle last_insert_id with arguments in the evalengine Signed-off-by: Andres Taylor --- .../endtoend/vtgate/queries/misc/misc_test.go | 3 +- go/vt/vtgate/engine/fake_vcursor_test.go | 3 ++ go/vt/vtgate/engine/primitive.go | 2 ++ go/vt/vtgate/evalengine/compiler_asm.go | 8 +++++ go/vt/vtgate/evalengine/expr_env.go | 2 ++ go/vt/vtgate/evalengine/fn_misc.go | 31 +++++++++++++++++++ go/vt/vtgate/evalengine/translate_builtin.go | 5 +++ go/vt/vtgate/executorcontext/vcursor_impl.go | 6 ++++ 8 files changed, 59 insertions(+), 1 deletion(-) diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go index 05b96c33604..55f2398c75f 100644 --- a/go/test/endtoend/vtgate/queries/misc/misc_test.go +++ b/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 { diff --git a/go/vt/vtgate/engine/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go index f27ca380876..72422193ee8 100644 --- a/go/vt/vtgate/engine/fake_vcursor_test.go +++ b/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 { diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go index e6fa102581e..7734dd81a6b 100644 --- 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/evalengine/compiler_asm.go b/go/vt/vtgate/evalengine/compiler_asm.go index dfb1a30bffc..b717b0750f6 100644 --- a/go/vt/vtgate/evalengine/compiler_asm.go +++ b/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)") +} diff --git a/go/vt/vtgate/evalengine/expr_env.go b/go/vt/vtgate/evalengine/expr_env.go index 38a65f9b4e0..4a7f9849ab0 100644 --- a/go/vt/vtgate/evalengine/expr_env.go +++ b/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} diff --git a/go/vt/vtgate/evalengine/fn_misc.go b/go/vt/vtgate/evalengine/fn_misc.go index 8813b62f823..e323d9a7dad 100644 --- a/go/vt/vtgate/evalengine/fn_misc.go +++ b/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 { diff --git a/go/vt/vtgate/evalengine/translate_builtin.go b/go/vt/vtgate/evalengine/translate_builtin.go index 476ee32483b..1f8bd7798aa 100644 --- a/go/vt/vtgate/evalengine/translate_builtin.go +++ b/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) } diff --git a/go/vt/vtgate/executorcontext/vcursor_impl.go b/go/vt/vtgate/executorcontext/vcursor_impl.go index 1896b3f267a..1890bc72202 100644 --- a/go/vt/vtgate/executorcontext/vcursor_impl.go +++ b/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 +}