From b6c6b3ff3580082ef51e2c87ce4585d16b519a34 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Fri, 20 Dec 2024 16:17:48 +0100
Subject: [PATCH] make sure evalengine expressions are not simplified

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 .../endtoend/vtgate/queries/misc/misc_test.go | 34 +++++++++++++++++++
 go/vt/vtgate/evalengine/fn_misc.go            |  4 +++
 .../planbuilder/testdata/set_cases.json       | 19 +++++++++++
 3 files changed, 57 insertions(+)

diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go
index e8701fe88bd..666a087c049 100644
--- a/go/test/endtoend/vtgate/queries/misc/misc_test.go
+++ b/go/test/endtoend/vtgate/queries/misc/misc_test.go
@@ -164,6 +164,7 @@ func TestSetAndGetLastInsertID(t *testing.T) {
 		"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",
+		"set @x = last_insert_id(%d)",
 	}
 
 	for _, workload := range []string{"olap", "oltp"} {
@@ -187,6 +188,39 @@ func TestSetAndGetLastInsertID(t *testing.T) {
 	}
 }
 
+func TestSetAndGetLastInsertIDWithInsert(t *testing.T) {
+	mcmp, closer := start(t)
+	defer closer()
+
+	mcmp.Exec("insert into t1(id1, id2) values (last_insert_id(12),0)")
+	mcmp.Exec("select last_insert_id()")
+	mcmp.Exec("insert into t1(id1, id2) values (13,last_insert_id(0))")
+	mcmp.Exec("select last_insert_id()")
+
+	mcmp.Exec("begin")
+	mcmp.Exec("insert into t1(id1, id2) values (last_insert_id(14),0)")
+	mcmp.Exec("select last_insert_id()")
+	mcmp.Exec("insert into t1(id1, id2) values (15,last_insert_id(0))")
+	mcmp.Exec("select last_insert_id()")
+	mcmp.Exec("commit")
+
+	_, err := mcmp.VtConn.ExecuteFetch("set workload = olap", 1, false)
+	require.NoError(t, err)
+
+	mcmp.Exec("insert into t1(id1, id2) values (last_insert_id(16),0)")
+	mcmp.Exec("select last_insert_id()")
+	mcmp.Exec("insert into t1(id1, id2) values (17,last_insert_id(0))")
+	mcmp.Exec("select last_insert_id()")
+
+	mcmp.Exec("begin")
+	mcmp.Exec("insert into t1(id1, id2) values (last_insert_id(18),0)")
+	mcmp.Exec("select last_insert_id()")
+	mcmp.Exec("insert into t1(id1, id2) values (19,last_insert_id(0))")
+	mcmp.Exec("select last_insert_id()")
+	mcmp.Exec("commit")
+
+}
+
 // TestVindexHints tests that vindex hints work as intended.
 func TestVindexHints(t *testing.T) {
 	mcmp, closer := start(t)
diff --git a/go/vt/vtgate/evalengine/fn_misc.go b/go/vt/vtgate/evalengine/fn_misc.go
index 6f7d27c1101..8d6db24a2dc 100644
--- a/go/vt/vtgate/evalengine/fn_misc.go
+++ b/go/vt/vtgate/evalengine/fn_misc.go
@@ -223,6 +223,10 @@ func (call *builtinLastInsertID) compile(c *compiler) (ctype, error) {
 	return ctype{Type: sqltypes.Uint64, Flag: arg.Flag & flagNullable, Col: collationNumeric}, nil
 }
 
+func (call *builtinLastInsertID) constant() bool {
+	return false // we don't want this function to be simplified away
+}
+
 func printIPv6AsIPv4(addr netip.Addr) (netip.Addr, bool) {
 	b := addr.AsSlice()
 	if len(b) != 16 {
diff --git a/go/vt/vtgate/planbuilder/testdata/set_cases.json b/go/vt/vtgate/planbuilder/testdata/set_cases.json
index 58cb2fffa75..5528c0e32fe 100644
--- a/go/vt/vtgate/planbuilder/testdata/set_cases.json
+++ b/go/vt/vtgate/planbuilder/testdata/set_cases.json
@@ -605,5 +605,24 @@
         ]
       }
     }
+  },
+  {
+    "QueryType": "SET",
+    "Original": "set @foo = last_insert_id(1)",
+    "Instructions": {
+      "OperatorType": "Set",
+      "Ops": [
+        {
+          "Type": "UserDefinedVariable",
+          "Name": "foo",
+          "Expr": "last_insert_id(1)"
+        }
+      ],
+      "Inputs": [
+        {
+          "OperatorType": "SingleRow"
+        }
+      ]
+    }
   }
 ]