From 4f0f8f5a53ee4fd6a2db0007a31cc9e93548a527 Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Tue, 7 Jan 2025 15:47:05 +0100 Subject: [PATCH] handle INSERT with last_insert_id() for sharded keyspaces Signed-off-by: Andres Taylor --- go/vt/vtgate/engine/insert.go | 6 +-- go/vt/vtgate/engine/insert_common.go | 2 +- go/vt/vtgate/engine/insert_select.go | 2 +- go/vt/vtgate/planbuilder/insert.go | 11 +++--- .../planbuilder/operator_transformers.go | 6 +-- .../planbuilder/testdata/aggr_cases.json | 38 +++++++++++++++++++ .../planbuilder/testdata/dml_cases.json | 24 ++++++++++++ 7 files changed, 75 insertions(+), 14 deletions(-) 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/aggr_cases.json b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json index 49a03a8f05a..1ecbf3d4ff9 100644 --- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.json @@ -6160,6 +6160,44 @@ ] } }, + { + "comment": "last_insert_id on aggregation calculated at the vtgate level", + "query": "select last_insert_id(count(*)) from user", + "plan": { + "QueryType": "SELECT", + "Original": "select last_insert_id(count(*)) from user", + "Instructions": { + "OperatorType": "Projection", + "Expressions": [ + "last_insert_id(count(*)) as last_insert_id(count(*))" + ], + "Inputs": [ + { + "OperatorType": "Aggregate", + "Variant": "Scalar", + "Aggregates": "sum_count_star(0) AS count(*)", + "Inputs": [ + { + "OperatorType": "Route", + "Variant": "Scatter", + "Keyspace": { + "Name": "user", + "Sharded": true + }, + "FetchLastInsertID": true, + "FieldQuery": "select count(*) from `user` where 1 != 1", + "Query": "select count(*) from `user`", + "Table": "`user`" + } + ] + } + ] + }, + "TablesUsed": [ + "user.user" + ] + } + }, { "comment": "aggregation on top of aggregation works fine", "query": "select distinct count(*) from user, (select distinct count(*) from user) X", 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)",