handle INSERT with last_insert_id() for sharded keyspaces

Signed-off-by: Andres Taylor <[email protected]>

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/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",

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