diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go index d1547e0fc46..60bc97ed389 100644 --- a/go/test/endtoend/vtgate/queries/misc/misc_test.go +++ b/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) diff --git a/go/vt/vtgate/engine/insert.go b/go/vt/vtgate/engine/insert.go index 10a4048572f..5bc206f7465 100644 --- a/go/vt/vtgate/engine/insert.go +++ b/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, diff --git a/go/vt/vtgate/engine/insert_common.go b/go/vt/vtgate/engine/insert_common.go index 629d848d978..d4cae045e86 100644 --- a/go/vt/vtgate/engine/insert_common.go +++ b/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 } diff --git a/go/vt/vtgate/engine/insert_select.go b/go/vt/vtgate/engine/insert_select.go index bccee5f2cf9..af834858175 100644 --- a/go/vt/vtgate/engine/insert_select.go +++ b/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, diff --git a/go/vt/vtgate/planbuilder/insert.go b/go/vt/vtgate/planbuilder/insert.go index 80516871623..d3ad5afac72 100644 --- a/go/vt/vtgate/planbuilder/insert.go +++ b/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) diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go index bc71c7195b4..b51eac449fc 100644 --- a/go/vt/vtgate/planbuilder/operator_transformers.go +++ b/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 diff --git a/go/vt/vtgate/planbuilder/testdata/dml_cases.json b/go/vt/vtgate/planbuilder/testdata/dml_cases.json index 8893b4df0c0..47cbc066aea 100644 --- a/go/vt/vtgate/planbuilder/testdata/dml_cases.json +++ b/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)", diff --git a/go/vt/vtgate/planbuilder/testdata/onecase.json b/go/vt/vtgate/planbuilder/testdata/onecase.json index 9d653b2f6e9..93338272d7f 100644 --- a/go/vt/vtgate/planbuilder/testdata/onecase.json +++ b/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": { } }