diff --git a/go/vt/sqlparser/ast_rewriting.go b/go/vt/sqlparser/ast_rewriting.go index 9a66202e7be..1e7029e3b87 100644 --- a/go/vt/sqlparser/ast_rewriting.go +++ b/go/vt/sqlparser/ast_rewriting.go @@ -564,12 +564,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 d41e586019b..a3cb1120dff 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 d1a69fbc7a9..467421f7258 100644 --- a/go/vt/vtgate/planbuilder/testdata/select_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/select_cases.json @@ -8068,5 +8068,42 @@ "user.user" ] } + }, + { + "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" + ] + } } ]