diff --git a/go/vt/sqlparser/ast_rewriting.go b/go/vt/sqlparser/ast_rewriting.go index 743b01a92aa..2b6683a15d4 100644 --- a/go/vt/sqlparser/ast_rewriting.go +++ b/go/vt/sqlparser/ast_rewriting.go @@ -567,12 +567,19 @@ var funcRewrites = map[string]string{ } func (er *astRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) { - bindVar, found := funcRewrites[node.Name.Lowered()] + lowered := node.Name.Lowered() + if lowered == "last_insert_id" && len(node.Exprs) > 0 { + // if we are dealing with is LAST_INSERT_ID() with an argument, we don't need to rewrite it. + // with an argument, this is an identity function that will update the session state and + // sets the correct fields in the OK TCP packet that we send back + return + } + bindVar, found := funcRewrites[lowered] if !found || (bindVar == DBVarName && !er.shouldRewriteDatabaseFunc) { return } if len(node.Exprs) > 0 { - er.err = vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "Argument to %s() not supported", node.Name.Lowered()) + er.err = vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "Argument to %s() not supported", lowered) return } cursor.Replace(bindVarExpression(bindVar)) diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases.json b/go/vt/vtgate/planbuilder/testdata/dml_cases.json index b25ec342b46..e156460a232 100644 --- a/go/vt/vtgate/planbuilder/testdata/dml_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/dml_cases.json @@ -6244,5 +6244,28 @@ "user.ref" ] } + }, + { + "comment": "update using last_insert_id with an argument", + "query": "update main.m1 set foo = last_insert_id(foo+1) where id = 12345", + "plan": { + "QueryType": "UPDATE", + "Original": "update main.m1 set foo = last_insert_id(foo+1) where id = 12345", + "Instructions": { + "OperatorType": "Update", + "Variant": "Unsharded", + "Keyspace": { + "Name": "main", + "Sharded": false + }, + "TargetTabletType": "PRIMARY", + "MultiShardAutocommit": false, + "Query": "update m1 set foo = last_insert_id(foo + 1) where id = 12345", + "Table": "m1" + }, + "TablesUsed": [ + "main.m1" + ] + } } ] diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.json b/go/vt/vtgate/planbuilder/testdata/select_cases.json index 48d11d0e816..0b97f1b0aca 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json @@ -7944,6 +7944,43 @@ ] } }, + { + "comment": "allow last_insert_id with argument", + "query": "select last_insert_id(id) from user", + "v3-plan": { + "QueryType": "SELECT", + "Original": "select last_insert_id(id) from user", + "Instructions": { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select last_insert_id(id) from `user` where 1 != 1", + "Query": "select last_insert_id(id) from `user`", + "Table": "`user`" + } + }, + "gen4-plan": { + "QueryType": "SELECT", + "Original": "select last_insert_id(id) from user", + "Instructions": { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FieldQuery": "select last_insert_id(id) from `user` where 1 != 1", + "Query": "select last_insert_id(id) from `user`", + "Table": "`user`" + }, + "TablesUsed": [ + "user.user" + ] + } + }, { "comment": "merge subquery using MAX and join into single route", "query": "select 1 from user join music_extra on user.id = music_extra.user_id where music_extra.music_id = (select max(music_id) from music_extra where user_id = user.id)",