wip

Signed-off-by: Harshit Gangal <[email protected]>
vitessio · Jan 7, 2025 · 5d718f8 · 5d718f8
1 parent ff99643
commit 5d718f8
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 15 deletions.
diff --git a/go/test/endtoend/vtgate/queries/misc/misc_test.go b/go/test/endtoend/vtgate/queries/misc/misc_test.go
@@ -218,7 +218,6 @@ func TestSetAndGetLastInsertIDWithInsert(t *testing.T) {
 	mcmp.Exec("insert into t1(id1, id2) values (19,last_insert_id(0))")
 	mcmp.Exec("select last_insert_id()")
 	mcmp.Exec("commit")
-
 }
 
 // TestVindexHints tests that vindex hints work as intended.

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)",