Support last_insert_id() with argument in SELECT

changelog/22.0/22.0.0/summary.md

@@ -8,10 +8,12 @@
   - **[RPC Changes](#rpc-changes)**
   - **[Prefer not promoting a replica that is currently taking a backup](#reparents-prefer-not-backing-up)**
   - **[VTOrc Config File Changes](#vtorc-config-file-changes)**
+  - **[Added support for LAST_INSERT_ID(expr)](#last_insert_id)**
 - **[Minor Changes](#minor-changes)**
   - **[VTTablet Flags](#flags-vttablet)**
 
 
+
 ## <a id="major-changes"/>Major Changes</a>
 
 ### <a id="rpc-changes"/>RPC Changes</a>
@@ -59,6 +61,9 @@ The following fields can be dynamically changed -
 
 To upgrade to the newer version of the configuration file, first switch to using the flags in your current deployment before upgrading. Then you can switch to using the configuration file in the newer release.
 
+### <a id="last_insert_id"/>Added support for LAST_INSERT_ID(expr)
+
+Added support for LAST_INSERT_ID(expr) to align with MySQL behavior, enabling session-level assignment of the last insert ID via query expressions.
 
 ## <a id="minor-changes"/>Minor Changes</a>
 

go/test/endtoend/vtgate/vitess_tester/aggregation/aggregation.test

@@ -64,4 +64,10 @@ from (select id, count(*) as num_segments from t1 group by 1 order by 2 desc lim
          join t2 u on u.id = t.id;
 
 select name
-from (select name from t1 group by name having count(t1.id) > 1) t1;
+from (select name from t1 group by name having count(t1.id) > 1) t1;
+
+# this query uses last_insert_id with a column argument to show that this works well
+select id, last_insert_id(count(*)) as num_segments from t1 group by id;
+
+# checking that we stored the correct value in the last_insert_id
+select last_insert_id();

go/test/endtoend/vtgate/vitess_tester/expressions/expressions.test

@@ -28,3 +28,7 @@ SELECT (~ (1 || 0)) IS NULL;
 
 SELECT 1
 WHERE (~ (1 || 0)) IS NULL;
+
+select last_insert_id(12);
+
+select last_insert_id();

go/vt/sqlparser/cow.go → go/vt/sqlparser/copy_on_write.go

@@ -53,6 +53,35 @@ func CopyOnRewrite(
 	return out
 }
 
+// CopyAndReplaceExpr traverses a syntax tree recursively, starting with root,
+// and replaceFn is called for each node that is an Expr. If replaceFn returns
+// true, the node is replaced with the returned node.
+func CopyAndReplaceExpr(node SQLNode, replaceFn func(node Expr) (Expr, bool)) SQLNode {
+	var replace Expr
+	pre := func(node, _ SQLNode) bool {
+		expr, ok := node.(Expr)
+		if !ok {
+			return true
+		}
+		newExpr, ok := replaceFn(expr)
+		if !ok {
+			return true
+		}
+		replace = newExpr
+		return false
+	}
+
+	post := func(cursor *CopyOnWriteCursor) {
+		if replace == nil {
+			return
+		}
+		cursor.Replace(replace)
+		replace = nil
+	}
+
+	return CopyOnRewrite(node, pre, post, nil)
+}
+
 // StopTreeWalk aborts the current tree walking. No more nodes will be visited, and the rewriter will exit out early
 func (c *CopyOnWriteCursor) StopTreeWalk() {
 	c.stop = true

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -263,6 +263,10 @@ func (t *noopVCursor) SetFoundRows(u uint64) {
 	panic("implement me")
 }
 
+func (t *noopVCursor) SetLastInsertID(id uint64) {
+	panic("implement me")
+}
+
 func (t *noopVCursor) InTransactionAndIsDML() bool {
 	panic("implement me")
 }
@@ -500,6 +504,9 @@ func (f *loggingVCursor) GetSystemVariables(func(k string, v string)) {
 func (f *loggingVCursor) SetFoundRows(u uint64) {
 	panic("implement me")
 }
+func (f *loggingVCursor) SetLastInsertID(id uint64) {
+	panic("implement me")
+}
 
 func (f *loggingVCursor) InTransactionAndIsDML() bool {
 	return false

go/vt/vtgate/engine/primitive.go

@@ -182,7 +182,7 @@ type (
 		SetPriority(string)
 		SetExecQueryTimeout(timeout *int)
 		SetFoundRows(uint64)
-
+		SetLastInsertID(uint64)
 		SetDDLStrategy(string)
 		GetDDLStrategy() string
 		SetMigrationContext(string)

go/vt/vtgate/engine/save_to_session.go

@@ -0,0 +1,125 @@
+/*
+Copyright 2019 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package engine
+
+import (
+	"context"
+
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/evalengine"
+)
+
+type SaveToSession struct {
+	noTxNeeded
+
+	Source Primitive
+	Offset evalengine.Expr
+}
+
+var _ Primitive = (*SaveToSession)(nil)
+
+func (s *SaveToSession) RouteType() string {
+	return s.Source.RouteType()
+}
+
+func (s *SaveToSession) GetKeyspaceName() string {
+	return s.Source.GetKeyspaceName()
+}
+
+func (s *SaveToSession) GetTableName() string {
+	return s.Source.GetTableName()
+}
+
+func (s *SaveToSession) GetFields(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	return s.Source.GetFields(ctx, vcursor, bindVars)
+}
+
+// TryExecute on SaveToSession will execute the Source and save the last row's value of Offset into the session.
+func (s *SaveToSession) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	result, err := s.Source.TryExecute(ctx, vcursor, bindVars, wantfields)
+	if err != nil {
+		return nil, err
+	}
+
+	intEvalResult, ok, err := s.getUintFromOffset(ctx, vcursor, bindVars, result)
+	if err != nil {
+		return nil, err
+	}
+	if ok {
+		vcursor.Session().SetLastInsertID(intEvalResult)
+	}
+	return result, nil
+}
+
+func (s *SaveToSession) getUintFromOffset(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, result *sqltypes.Result) (uint64, bool, error) {
+	if len(result.Rows) == 0 {
+		return 0, false, nil
+	}
+
+	env := evalengine.NewExpressionEnv(ctx, bindVars, vcursor)
+	env.Row = result.Rows[len(result.Rows)-1] // last row
+	evalResult, err := env.Evaluate(s.Offset)
+	if err != nil {
+		return 0, false, err
+	}
+	value, err := evalResult.Value(vcursor.ConnCollation()).ToCastUint64()
+	if err != nil {
+		return 0, false, err
+	}
+	return value, true, nil
+}
+
+// TryStreamExecute on SaveToSession will execute the Source and save the last row's value of Offset into the session.
+func (s *SaveToSession) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	var value *uint64
+
+	f := func(qr *sqltypes.Result) error {
+		v, ok, err := s.getUintFromOffset(ctx, vcursor, bindVars, qr)
+		if err != nil {
+			return err
+		}
+		if ok {
+			value = &v
+		}
+		return callback(qr)
+	}
+
+	err := s.Source.TryStreamExecute(ctx, vcursor, bindVars, wantfields, f)
+	if err != nil {
+		return err
+	}
+	if value != nil {
+		vcursor.Session().SetLastInsertID(*value)
+	}
+
+	return nil
+}
+
+func (s *SaveToSession) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{s.Source}, nil
+}
+
+func (s *SaveToSession) description() PrimitiveDescription {
+	return PrimitiveDescription{
+		OperatorType: "SaveToSession",
+		Other: map[string]interface{}{
+			"Offset": sqlparser.String(s.Offset),
+		},
+	}
+}

go/vt/vtgate/executorcontext/vcursor_impl.go

@@ -1130,6 +1130,10 @@ func (vc *VCursorImpl) SetFoundRows(foundRows uint64) {
 	vc.SafeSession.SetFoundRows(foundRows)
 }
 
+func (vc *VCursorImpl) SetLastInsertID(id uint64) {
+	vc.SafeSession.LastInsertId = id
+}
+
 // SetDDLStrategy implements the SessionActions interface
 func (vc *VCursorImpl) SetDDLStrategy(strategy string) {
 	vc.SafeSession.SetDDLStrategy(strategy)

go/vt/vtgate/planbuilder/operator_transformers.go

@@ -81,6 +81,8 @@ func transformToPrimitive(ctx *plancontext.PlanningContext, op operators.Operato
 		return transformRecurseCTE(ctx, op)
 	case *operators.PercentBasedMirror:
 		return transformPercentBasedMirror(ctx, op)
+	case *operators.SaveToSession:
+		return transformSaveToSession(ctx, op)
 	}
 
 	return nil, vterrors.VT13001(fmt.Sprintf("unknown type encountered: %T (transformToPrimitive)", op))
@@ -481,6 +483,34 @@ func newSimpleProjection(cols []int, colNames []string, src engine.Primitive) en
 	}
 }
 
+func transformSaveToSession(ctx *plancontext.PlanningContext, op *operators.SaveToSession) (engine.Primitive, error) {
+	src, err := transformToPrimitive(ctx, op.Source)
+	if err != nil {
+		return nil, err
+	}
+
+	v := op.Offset
+	return createSaveToSessionOnOffset(ctx, v, src)
+}
+
+func createSaveToSessionOnOffset(ctx *plancontext.PlanningContext, v int, src engine.Primitive) (engine.Primitive, error) {
+	cfg := &evalengine.Config{
+		ResolveType: ctx.TypeForExpr,
+		Collation:   ctx.SemTable.Collation,
+		Environment: ctx.VSchema.Environment(),
+	}
+
+	offset, err := evalengine.Translate(sqlparser.NewOffset(v, nil), cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &engine.SaveToSession{
+		Source: src,
+		Offset: offset,
+	}, nil
+}
+
 func transformFilter(ctx *plancontext.PlanningContext, op *operators.Filter) (engine.Primitive, error) {
 	src, err := transformToPrimitive(ctx, op.Source)
 	if err != nil {