Support last_insert_id() with argument

changelog/22.0/22.0.0/summary.md

@@ -5,6 +5,8 @@
 - **[Major Changes](#major-changes)**
   - **[RPC Changes](#rpc-changes)**
   - **[Prefer not promoting a replica that is currently taking a backup](#reparents-prefer-not-backing-up)**
+  - **[Added support for LAST_INSERT_ID(expr)](#last_insert_id)**
+
 
 
 ## <a id="major-changes"/>Major Changes</a>
@@ -25,4 +27,8 @@ For planned reparents, hosts taking backups with a backup engine other than `bui
 valid candidates. This means they will never get promoted - not even if there's no other candidates.
 
 Note that behavior for `builtin` backups remains unchanged: a replica that is currently taking a `builtin` backup will
-never be promoted, neither by planned nor by emergency reparents.
+never be promoted, neither by planned nor by emergency reparents.
+
+### <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.

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/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/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 {

go/vt/vtgate/planbuilder/operators/query_planning.go

@@ -829,11 +829,34 @@ func addTruncationOrProjectionToReturnOutput(ctx *plancontext.PlanningContext, s
 
 	cols := output.GetSelectExprs(ctx)
 	sizeCorrect := len(selExprs) == len(cols) || tryTruncateColumnsAt(output, len(selExprs))
-	if sizeCorrect && colNamesAlign(selExprs, cols) {
+	if !sizeCorrect || !colNamesAlign(selExprs, cols) {
+		output = createSimpleProjection(ctx, selExprs, output)
+	}
+
+	if !ctx.SemTable.QuerySignature.LastInsertIDArg {
 		return output
 	}
 
-	return createSimpleProjection(ctx, selExprs, output)
+	var offset int
+	for i, expr := range selExprs {
+		ae, ok := expr.(*sqlparser.AliasedExpr)
+		if !ok {
+			panic(vterrors.VT09015())
+		}
+		fnc, ok := ae.Expr.(*sqlparser.FuncExpr)
+		if !ok || !fnc.Name.EqualString("last_insert_id") {
+			continue
+		}
+		offset = i
+		break
+	}
+
+	return &SaveToSession{
+		unaryOperator: unaryOperator{
+			Source: output,
+		},
+		Offset: offset,
+	}
 }
 
 func colNamesAlign(expected, actual sqlparser.SelectExprs) bool {

go/vt/vtgate/planbuilder/operators/save_to_session.go

@@ -0,0 +1,71 @@
+/*
+Copyright 2022 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 operators
+
+import (
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+)
+
+// SaveToSession is used to save a value to the session.
+// At the moment it's only used for last_insert_id, but it could be used for other things in the future.
+type SaveToSession struct {
+	unaryOperator
+	Offset int
+}
+
+var _ Operator = (*SaveToSession)(nil)
+
+func (s *SaveToSession) Clone(inputs []Operator) Operator {
+	k := *s
+	k.Source = inputs[0]
+	return &k
+}
+
+func (s *SaveToSession) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Expr) Operator {
+	src := s.Source.AddPredicate(ctx, expr)
+	s.Source = src
+	return s
+}
+
+func (s *SaveToSession) AddColumn(ctx *plancontext.PlanningContext, reuseExisting bool, addToGroupBy bool, expr *sqlparser.AliasedExpr) int {
+	return s.Source.AddColumn(ctx, reuseExisting, addToGroupBy, expr)
+}
+
+func (s *SaveToSession) AddWSColumn(ctx *plancontext.PlanningContext, offset int, underRoute bool) int {
+	return s.Source.AddWSColumn(ctx, offset, underRoute)
+}
+
+func (s *SaveToSession) FindCol(ctx *plancontext.PlanningContext, expr sqlparser.Expr, underRoute bool) int {
+	return s.Source.FindCol(ctx, expr, underRoute)
+}
+
+func (s *SaveToSession) GetColumns(ctx *plancontext.PlanningContext) []*sqlparser.AliasedExpr {
+	return s.Source.GetColumns(ctx)
+}
+
+func (s *SaveToSession) GetSelectExprs(ctx *plancontext.PlanningContext) sqlparser.SelectExprs {
+	return s.Source.GetSelectExprs(ctx)
+}
+
+func (s *SaveToSession) ShortDescription() string {
+	return ""
+}
+
+func (s *SaveToSession) GetOrdering(ctx *plancontext.PlanningContext) []OrderBy {
+	return s.Source.GetOrdering(ctx)
+}