gen4 planner: allow last_insert_id with arguments (16.0)

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))

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"
+      ]
+    }
   }
 ]

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"
+      ]
+    }
   }
 ]