feat: pass on the FetchLastInsertID to the engine primitives

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

go/vt/vtgate/engine/delete.go

@@ -131,6 +131,9 @@ func (del *Delete) description() PrimitiveDescription {
 	}
 
 	addFieldsIfNotEmpty(del.DML, other)
+	if del.FetchLastInsertID {
+		other["FetchLastInsertID"] = del.FetchLastInsertID
+	}
 
 	return PrimitiveDescription{
 		OperatorType:     "Delete",

go/vt/vtgate/engine/insert.go

@@ -385,6 +385,10 @@ func (ins *Insert) description() PrimitiveDescription {
 		}
 	}
 
+	if ins.FetchLastInsertID {
+		other["FetchLastInsertID"] = true
+	}
+
 	return PrimitiveDescription{
 		OperatorType:     "Insert",
 		Keyspace:         ins.Keyspace,

go/vt/vtgate/engine/route.go

@@ -376,6 +376,9 @@ func (route *Route) description() PrimitiveDescription {
 		"Table":      route.GetTableName(),
 		"FieldQuery": route.FieldQuery,
 	}
+	if route.FetchLastInsertID {
+		other["FetchLastInsertID"] = true
+	}
 	if route.Vindex != nil {
 		other["Vindex"] = route.Vindex.String()
 	}

go/vt/vtgate/engine/update.go

@@ -214,6 +214,9 @@ func (upd *Update) description() PrimitiveDescription {
 	if len(changedVindexes) > 0 {
 		other["ChangedVindexValues"] = changedVindexes
 	}
+	if upd.FetchLastInsertID {
+		other["FetchLastInsertID"] = upd.FetchLastInsertID
+	}
 
 	return PrimitiveDescription{
 		OperatorType:     "Update",

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -536,6 +536,7 @@ func routeToEngineRoute(ctx *plancontext.PlanningContext, op *operators.Route, h
 		TableName:           strings.Join(tableNames, ", "),
 		RoutingParameters:   rp,
 		TruncateColumnCount: op.ResultColumns,
+		FetchLastInsertID:   ctx.SemTable.ShouldFetchLastInsertID(),
 	}
 	if hints != nil {
 		e.ScatterErrorsAsWarnings = hints.scatterErrorsAsWarnings
@@ -601,7 +602,7 @@ func transformRoutePlan(ctx *plancontext.PlanningContext, op *operators.Route) (
 	case *sqlparser.Delete:
 		return buildDeletePrimitive(ctx, op, dmlOp, stmt, hints)
 	case *sqlparser.Insert:
-		return buildInsertPrimitive(op, dmlOp, stmt, hints)
+		return buildInsertPrimitive(ctx, op, dmlOp, stmt, hints)
 	default:
 		return nil, vterrors.VT13001(fmt.Sprintf("dont know how to %T", stmt))
 	}
@@ -637,7 +638,10 @@ func buildRoutePrimitive(ctx *plancontext.PlanningContext, op *operators.Route,
 }
 
 func buildInsertPrimitive(
-	rb *operators.Route, op operators.Operator, stmt *sqlparser.Insert,
+	ctx *plancontext.PlanningContext,
+	rb *operators.Route,
+	op operators.Operator,
+	stmt *sqlparser.Insert,
 	hints *queryHints,
 ) (engine.Primitive, error) {
 	ins := op.(*operators.Insert)
@@ -656,8 +660,9 @@ func buildInsertPrimitive(
 	}
 
 	eins := &engine.Insert{
-		InsertCommon: ic,
-		VindexValues: ins.VindexValues,
+		InsertCommon:      ic,
+		VindexValues:      ins.VindexValues,
+		FetchLastInsertID: ctx.SemTable.ShouldFetchLastInsertID(),
 	}
 
 	// we would need to generate the query on the fly. The only exception here is
@@ -788,6 +793,7 @@ func createDMLPrimitive(ctx *plancontext.PlanningContext, rb *operators.Route, h
 		Vindexes:          colVindexes,
 		OwnedVindexQuery:  vindexQuery,
 		RoutingParameters: rp,
+		FetchLastInsertID: ctx.SemTable.ShouldFetchLastInsertID(),
 	}
 
 	if rb.Routing.OpCode() != engine.Unsharded && vindexQuery != "" {

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

@@ -2235,6 +2235,52 @@
       ]
     }
   },
+  {
+    "comment": "update with last_insert_id in SET",
+    "query": "update user_extra set val = last_insert_id(123)",
+    "plan": {
+      "QueryType": "UPDATE",
+      "Original": "update user_extra set val = last_insert_id(123)",
+      "Instructions": {
+        "OperatorType": "Update",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "TargetTabletType": "PRIMARY",
+        "FetchLastInsertID": true,
+        "Query": "update user_extra set val = last_insert_id(123)",
+        "Table": "user_extra"
+      },
+      "TablesUsed": [
+        "user.user_extra"
+      ]
+    }
+  },
+  {
+    "comment": "delete with last_insert_id in where",
+    "query": "delete from user_extra where val = last_insert_id(123)",
+    "plan": {
+      "QueryType": "DELETE",
+      "Original": "delete from user_extra where val = last_insert_id(123)",
+      "Instructions": {
+        "OperatorType": "Delete",
+        "Variant": "Scatter",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "TargetTabletType": "PRIMARY",
+        "FetchLastInsertID": true,
+        "Query": "delete from user_extra where val = last_insert_id(123)",
+        "Table": "user_extra"
+      },
+      "TablesUsed": [
+        "user.user_extra"
+      ]
+    }
+  },
   {
     "comment": "delete from with no where clause",
     "query": "delete from user_extra",

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

@@ -5599,6 +5599,7 @@
           "Name": "user",
           "Sharded": true
         },
+        "FetchLastInsertID": true,
         "FieldQuery": "select last_insert_id(id) from `user` where 1 != 1",
         "Query": "select last_insert_id(id) from `user`",
         "Table": "`user`"

go/vt/vtgate/semantics/analyzer_test.go

@@ -638,15 +638,17 @@ func TestQuerySignatureLastInsertID(t *testing.T) {
 	}, {
 		query:    "select last_insert_id(123)",
 		expected: true,
+	}, {
+		query:    "update user_extra set val = last_insert_id(123)",
+		expected: true,
 	}}
 
 	for _, tc := range queries {
 		t.Run(tc.query, func(t *testing.T) {
 			ast, err := sqlparser.NewTestParser().Parse(tc.query)
 			require.NoError(t, err)
 
-			sel := ast.(*sqlparser.Select)
-			st, err := AnalyzeStrict(sel, "dbName", fakeSchemaInfo())
+			st, err := AnalyzeStrict(ast, "dbName", fakeSchemaInfo())
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, st.QuerySignature.LastInsertIDArg)
 		})

go/vt/vtgate/semantics/semantic_table.go

@@ -1014,6 +1014,13 @@ func (st *SemTable) GetMirrorInfo() MirrorInfo {
 	return mirrorInfo(st.Tables)
 }
 
+func (st *SemTable) ShouldFetchLastInsertID() bool {
+	if st == nil {
+		return false
+	}
+	return st.QuerySignature.LastInsertIDArg
+}
+
 // mirrorInfo looks through all tables with mirror rules defined, and returns a
 // MirrorInfo containing the lowest mirror percentage found across all rules.
 //