wip

go/test/endtoend/vtgate/queries/misc/misc_test.go

@@ -197,6 +197,7 @@ func TestSetAndGetLastInsertIDWithInsert(t *testing.T) {
 	tests := []string{
 		"insert into t1(id1, id2) values (last_insert_id(%d),%d)",
 		"insert into t1(id1, id2) values (%d, last_insert_id(%d))",
+		"insert into t1(id1, id2) select last_insert_id(%d), %d",
 	}
 
 	i := 0
@@ -228,6 +229,50 @@ func TestSetAndGetLastInsertIDWithInsert(t *testing.T) {
 	}
 }
 
+func TestSetAndGetLastInsertIDWithInsertUnsharded(t *testing.T) {
+	// in this test we can't use the mcmp, so we need to assert the returned values manually
+	mcmp, closer := start(t)
+	defer closer()
+
+	tests := []string{
+		"insert into uks.unsharded(id1, id2) values (last_insert_id(%d),12)",
+		"insert into uks.unsharded(id1, id2) values (%d, last_insert_id(34))",
+		"insert into uks.unsharded(id1, id2) select last_insert_id(%d), 453",
+	}
+
+	i := 0
+	getVal := func() int {
+		defer func() { i++ }()
+		return i
+	}
+
+	runTests := func(mcmp *utils.MySQLCompare) {
+		for _, test := range tests {
+			lastInsertID := getVal()
+			query := fmt.Sprintf(test, lastInsertID)
+			log.Errorf("test: %s", query)
+			utils.Exec(t, mcmp.VtConn, query)
+			result := utils.Exec(t, mcmp.VtConn, "select last_insert_id()")
+			uintVal, err := result.Rows[0][0].ToCastUint64()
+			require.NoError(t, err)
+			require.EqualValues(t, lastInsertID, uintVal, query)
+		}
+	}
+
+	for _, workload := range []string{"olap", "oltp"} {
+		mcmp.Run(workload, func(mcmp *utils.MySQLCompare) {
+			_, err := mcmp.VtConn.ExecuteFetch("set workload = "+workload, 1, false)
+			require.NoError(t, err)
+			runTests(mcmp)
+
+			// run the queries again, but inside a transaction this time
+			mcmp.Exec("begin")
+			runTests(mcmp)
+			mcmp.Exec("commit")
+		})
+	}
+}
+
 // TestVindexHints tests that vindex hints work as intended.
 func TestVindexHints(t *testing.T) {
 	mcmp, closer := start(t)

go/vt/vtgate/engine/insert.go

@@ -58,11 +58,9 @@ type Insert struct {
 
 	// Alias represents the row alias with columns if specified in the query.
 	Alias string
-
-	FetchLastInsertID bool
 }
 
-// newQueryInsert creates an Insert with a query string.
+// newQueryInsert creates an Insert with a query string. Used in testing.
 func newQueryInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, query string) *Insert {
 	return &Insert{
 		InsertCommon: InsertCommon{
@@ -73,7 +71,7 @@ func newQueryInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, query stri
 	}
 }
 
-// newInsert creates a new Insert.
+// newInsert creates a new Insert. Used in testing.
 func newInsert(
 	opcode InsertOpcode,
 	ignore bool,

go/vt/vtgate/engine/insert_common.go

@@ -161,7 +161,7 @@ func (ins *InsertCommon) executeUnshardedTableQuery(ctx context.Context, vcursor
 	if err != nil {
 		return nil, err
 	}
-	qr, err := execShard(ctx, loggingPrimitive, vcursor, query, bindVars, rss[0], true, !ins.PreventAutoCommit /* canAutocommit */, false)
+	qr, err := execShard(ctx, loggingPrimitive, vcursor, query, bindVars, rss[0], true, !ins.PreventAutoCommit /* canAutocommit */, ins.FetchLastInsertID)
 	if err != nil {
 		return nil, err
 	}

go/vt/vtgate/engine/insert_select.go

@@ -51,7 +51,7 @@ type (
 	}
 )
 
-// newInsertSelect creates a new InsertSelect.
+// newInsertSelect creates a new InsertSelect. Used in testing.
 func newInsertSelect(
 	ignore bool,
 	keyspace *vindexes.Keyspace,

go/vt/vtgate/planbuilder/insert.go

@@ -51,7 +51,7 @@ func gen4InsertStmtPlanner(version querypb.ExecuteOptions_PlannerVersion, insStm
 	}
 	if ks != nil {
 		if tables[0].AutoIncrement == nil && !ctx.SemTable.ForeignKeysPresent() {
-			plan := insertUnshardedShortcut(insStmt, ks, tables)
+			plan := insertUnshardedShortcut(ctx, insStmt, ks, tables)
 			setCommentDirectivesOnPlan(plan, insStmt)
 			return newPlanResult(plan, operators.QualifiedTables(ks, tables)...), nil
 		}
@@ -90,12 +90,13 @@ func errOutIfPlanCannotBeConstructed(ctx *plancontext.PlanningContext, vTbl *vin
 	return ctx.SemTable.NotUnshardedErr
 }
 
-func insertUnshardedShortcut(stmt *sqlparser.Insert, ks *vindexes.Keyspace, tables []*vindexes.Table) engine.Primitive {
+func insertUnshardedShortcut(ctx *plancontext.PlanningContext, stmt *sqlparser.Insert, ks *vindexes.Keyspace, tables []*vindexes.Table) engine.Primitive {
 	eIns := &engine.Insert{
 		InsertCommon: engine.InsertCommon{
-			Opcode:    engine.InsertUnsharded,
-			Keyspace:  ks,
-			TableName: tables[0].Name.String(),
+			Opcode:            engine.InsertUnsharded,
+			Keyspace:          ks,
+			TableName:         tables[0].Name.String(),
+			FetchLastInsertID: ctx.SemTable.ShouldFetchLastInsertID(),
 		},
 	}
 	eIns.Query = generateQuery(stmt)

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -190,6 +190,7 @@ func transformInsertionSelection(ctx *plancontext.PlanningContext, op *operators
 			ForceNonStreaming: op.ForceNonStreaming,
 			Generate:          autoIncGenerate(ins.AutoIncrement),
 			ColVindexes:       ins.ColVindexes,
+			FetchLastInsertID: ctx.SemTable.ShouldFetchLastInsertID(),
 		},
 		VindexValueOffset: ins.VindexValueOffset,
 	}
@@ -659,9 +660,8 @@ func buildInsertPrimitive(
 	}
 
 	eins := &engine.Insert{
-		InsertCommon:      ic,
-		VindexValues:      ins.VindexValues,
-		FetchLastInsertID: ctx.SemTable.ShouldFetchLastInsertID(),
+		InsertCommon: ic,
+		VindexValues: ins.VindexValues,
 	}
 
 	// we would need to generate the query on the fly. The only exception here is

go/vt/vtgate/planbuilder/testdata/dml_cases.json

@@ -2648,6 +2648,30 @@
     },
     "skip_e2e": true
   },
+  {
+    "comment": "insert using last_insert_id with argument",
+    "query": "insert into unsharded values(last_insert_id(24), 2)",
+    "plan": {
+      "QueryType": "INSERT",
+      "Original": "insert into unsharded values(last_insert_id(24), 2)",
+      "Instructions": {
+        "OperatorType": "Insert",
+        "Variant": "Unsharded",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "TargetTabletType": "PRIMARY",
+        "FetchLastInsertID": true,
+        "Query": "insert into unsharded values (last_insert_id(24), 2)",
+        "TableName": "unsharded"
+      },
+      "TablesUsed": [
+        "main.unsharded"
+      ]
+    },
+    "skip_e2e": true
+  },
   {
     "comment": "update vindex value to null with multiple primary keyspace id",
     "query": "update user set name = null where id in (1, 2, 3)",

go/vt/vtgate/planbuilder/testdata/onecase.json

@@ -1,7 +1,7 @@
 [
   {
     "comment": "Add your test case here for debugging and run go test -run=One.",
-    "query": "",
+    "query": "select id2, last_insert_id(id+1) from user",
     "plan": {
     }
   }