Feat/multi update

pkg/datasource/sql/undo/builder/basic_undo_log_builder.go

@@ -100,6 +100,7 @@ func (b *BasicUndoLogBuilder) buildSelectArgs(stmt *ast.SelectStmt, args []drive
 		selectArgsIndexs = make([]int32, 0)
 		selectArgs       = make([]driver.Value, 0)
 	)
+
 	b.traversalArgs(stmt.Where, &selectArgsIndexs)
 	if stmt.OrderBy != nil {
 		for _, item := range stmt.OrderBy.Items {
@@ -151,11 +152,11 @@ func (b *BasicUndoLogBuilder) traversalArgs(node ast.Node, argsIndex *[]int32) {
 	}
 }
 
-func (u *BasicUndoLogBuilder) buildRecordImages(rowsi driver.Rows, tableMetaData types.TableMeta) (*types.RecordImage, error) {
+func (b *BasicUndoLogBuilder) buildRecordImages(rowsi driver.Rows, tableMetaData types.TableMeta) (*types.RecordImage, error) {
 	// select column names
 	columnNames := rowsi.Columns()
 	rowImages := make([]types.RowImage, 0)
-	ss := u.getScanSlice(columnNames, tableMetaData)
+	ss := b.getScanSlice(columnNames, tableMetaData)
 
 	for {
 		err := rowsi.Next(ss)

pkg/datasource/sql/undo/builder/mysql_multi_undo_log_builder.go

@@ -104,8 +104,7 @@ func (u *MySQLMultiUndoLogBuilder) AfterImage(ctx context.Context, execCtx *type
 
 		switch parseContext.ExecutorType {
 		case types.UpdateExecutor:
-			// todo change to use MultiUpdateExecutor
-			tmpImages, err = GetMySQLUpdateUndoLogBuilder().AfterImage(ctx, execCtx, []*types.RecordImage{u.beforeImages[i]})
+			tmpImages, err = GetMySQLMultiUpdateUndoLogBuilder().AfterImage(ctx, execCtx, []*types.RecordImage{u.beforeImages[i]})
 			break
 		case types.DeleteExecutor:
 			// todo use MultiDeleteExecutor

pkg/datasource/sql/undo/builder/mysql_multi_update_undo_log_builder.go

@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 builder
+
+import (
+	"context"
+	"database/sql/driver"
+	"fmt"
+	"strings"
+
+	"github.com/arana-db/parser"
+	"github.com/arana-db/parser/ast"
+	"github.com/arana-db/parser/format"
+	"github.com/pkg/errors"
+
+	"github.com/seata/seata-go/pkg/datasource/sql/types"
+	"github.com/seata/seata-go/pkg/datasource/sql/undo"
+	"github.com/seata/seata-go/pkg/util/bytes"
+	"github.com/seata/seata-go/pkg/util/log"
+)
+
+func init() {
+	undo.RegistrUndoLogBuilder(types.UpdateExecutor, GetMySQLMultiUpdateUndoLogBuilder)
+}
+
+type updateVisitor struct {
+	stmt *ast.UpdateStmt
+}
+
+func (m *updateVisitor) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
+	return n, true
+}
+
+func (m *updateVisitor) Leave(n ast.Node) (node ast.Node, ok bool) {
+	node = n
+	return node, true
+}
+
+type MySQLMultiUpdateUndoLogBuilder struct {
+	BasicUndoLogBuilder
+}
+
+func GetMySQLMultiUpdateUndoLogBuilder() undo.UndoLogBuilder {
+	return &MySQLMultiUpdateUndoLogBuilder{
+		BasicUndoLogBuilder: BasicUndoLogBuilder{},
+	}
+}
+
+func (u *MySQLMultiUpdateUndoLogBuilder) BeforeImage(ctx context.Context, execCtx *types.ExecContext) ([]*types.RecordImage, error) {
+	vals := execCtx.Values
+	if vals == nil {
+		for n, param := range execCtx.NamedValues {
+			vals[n] = param.Value
+		}
+	}
+
+	var updateStmts []*ast.UpdateStmt
+	for _, v := range execCtx.ParseContext.MultiStmt {
+		updateStmts = append(updateStmts, v.UpdateStmt)
+	}
+
+	// use
+	selectSQL, selectArgs, err := u.buildBeforeImageSQL(updateStmts, vals)
+	if err != nil {
+		return nil, err
+	}
+
+	stmt, err := execCtx.Conn.Prepare(selectSQL)
+	if err != nil {
+		log.Errorf("build prepare stmt: %+v", err)
+		return nil, err
+	}
+
+	rows, err := stmt.Query(selectArgs)
+	if err != nil {
+		log.Errorf("stmt query: %+v", err)
+		return nil, err
+	}
+
+	tableName := execCtx.ParseContext.UpdateStmt.TableRefs.TableRefs.Left.(*ast.TableSource).Source.(*ast.TableName).Name.O
+	metaData := execCtx.MetaDataMap[tableName]
+
+	image, err := u.buildRecordImages(rows, metaData)
+	if err != nil {
+		return nil, err
+	}
+
+	return []*types.RecordImage{image}, nil
+}
+
+func (u *MySQLMultiUpdateUndoLogBuilder) AfterImage(ctx context.Context, execCtx *types.ExecContext, beforeImages []*types.RecordImage) ([]*types.RecordImage, error) {
+	var beforeImage *types.RecordImage
+	if len(beforeImages) > 0 {
+		beforeImage = beforeImages[0]
+	}
+
+	tableName := execCtx.ParseContext.UpdateStmt.TableRefs.TableRefs.Left.(*ast.TableSource).Source.(*ast.TableName).Name.O
+	metaData := execCtx.MetaDataMap[tableName]
+	selectSQL, selectArgs := u.buildAfterImageSQL(beforeImage, metaData)
+
+	stmt, err := execCtx.Conn.Prepare(selectSQL)
+	if err != nil {
+		log.Errorf("build prepare stmt: %+v", err)
+		return nil, err
+	}
+
+	rows, err := stmt.Query(selectArgs)
+	if err != nil {
+		log.Errorf("stmt query: %+v", err)
+		return nil, err
+	}
+
+	image, err := u.buildRecordImages(rows, metaData)
+	if err != nil {
+		return nil, err
+	}
+
+	return []*types.RecordImage{image}, nil
+}
+
+func (u *MySQLMultiUpdateUndoLogBuilder) buildAfterImageSQL(beforeImage *types.RecordImage, meta types.TableMeta) (string, []driver.Value) {
+	sb := strings.Builder{}
+	// todo use ONLY_CARE_UPDATE_COLUMNS to judge select all columns or not
+	sb.WriteString("SELECT * FROM " + meta.Name + " ")
+	whereSQL := u.buildWhereConditionByPKs(meta.GetPrimaryKeyOnlyName(), len(beforeImage.Rows), "mysql", maxInSize)
+	sb.WriteString(" " + whereSQL + " ")
+	return sb.String(), u.buildPKParams(beforeImage.Rows, meta.GetPrimaryKeyOnlyName())
+}
+
+// buildSelectSQLByUpdate build select sql from update sql
+func (u *MySQLMultiUpdateUndoLogBuilder) buildBeforeImageSQL(updateStmts []*ast.UpdateStmt, args []driver.Value) (string, []driver.Value, error) {
+	if len(updateStmts) == 0 {
+		log.Errorf("invalid multi update stmt")
+		return "", nil, fmt.Errorf("invalid muliti update stmt")
+	}
+
+	var newArgs []driver.Value
+	var fields []*ast.SelectField
+	fieldsExits := make(map[string]struct{})
+	var whereCondition strings.Builder
+	for _, updateStmt := range updateStmts {
+		if updateStmt.Limit != nil {
+			return "", nil, errors.New("multi update SQL with limit condition is not support yet")
+		}
+		if updateStmt.Order != nil {
+			return "", nil, errors.New("multi update SQL with orderBy condition is not support yet")
+		}
+
+		// todo use ONLY_CARE_UPDATE_COLUMNS to judge select all columns or not
+		for _, column := range updateStmt.List {
+			if _, exist := fieldsExits[column.Column.String()]; exist {
+				continue
+			}
+			fieldsExits[column.Column.String()] = struct{}{}
+			fields = append(fields, &ast.SelectField{
+				Expr: &ast.ColumnNameExpr{
+					Name: column.Column,
+				},
+			})
+		}
+
+		tmpSelectStmt := ast.SelectStmt{
+			SelectStmtOpts: &ast.SelectStmtOpts{},
+			From:           updateStmt.TableRefs,
+			Where:          updateStmt.Where,
+			Fields:         &ast.FieldList{Fields: fields},
+			OrderBy:        updateStmt.Order,
+			Limit:          updateStmt.Limit,
+			TableHints:     updateStmt.TableHints,
+			LockInfo: &ast.SelectLockInfo{
+				LockType: ast.SelectLockForUpdate,
+			},
+		}
+		newArgs = append(newArgs, u.buildSelectArgs(&tmpSelectStmt, args)...)
+
+		in := bytes.NewByteBuffer([]byte{})
+		updateStmt.Where.Restore(format.NewRestoreCtx(format.RestoreKeyWordUppercase, in))
+		whereConditionStr := string(in.Bytes())
+
+		if whereCondition.Len() > 0 {
+			whereCondition.Write([]byte(" OR "))
+		}
+		whereCondition.Write([]byte(whereConditionStr))
+	}
+
+	// only just get the where condition
+	fakeSql := "select * from t where " + whereCondition.String()
+	fakeStmt, err := parser.New().ParseOneStmt(fakeSql, "", "")
+	if err != nil {
+		return "", nil, errors.Wrap(err, "multi update parse fake sql error")
+	}
+	fakeNode, ok := fakeStmt.Accept(&updateVisitor{})
+	if !ok {
+		return "", nil, errors.Wrap(err, "multi update accept update visitor error")
+	}
+	fakeSelectStmt, ok := fakeNode.(*ast.SelectStmt)
+	if !ok {
+		return "", nil, errors.New("multi update fake node is not select stmt")
+	}
+
+	selStmt := ast.SelectStmt{
+		SelectStmtOpts: &ast.SelectStmtOpts{},
+		From:           updateStmts[0].TableRefs,
+		Where:          fakeSelectStmt.Where,
+		Fields:         &ast.FieldList{Fields: fields},
+		TableHints:     updateStmts[0].TableHints,
+		LockInfo: &ast.SelectLockInfo{
+			LockType: ast.SelectLockForUpdate,
+		},
+	}
+
+	b := bytes.NewByteBuffer([]byte{})
+	selStmt.Restore(format.NewRestoreCtx(format.RestoreKeyWordUppercase, b))
+	sql := string(b.Bytes())
+	log.Infof("build select sql by update sourceQuery, sql {}", sql)
+
+	return sql, newArgs, nil
+}
+
+func (u *MySQLMultiUpdateUndoLogBuilder) GetExecutorType() types.ExecutorType {
+	return types.UpdateExecutor
+}

pkg/datasource/sql/undo/builder/mysql_multi_update_undo_log_builder_test.go

@@ -0,0 +1,88 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 builder
+
+import (
+	"database/sql/driver"
+	"testing"
+
+	"github.com/arana-db/parser/ast"
+	_ "github.com/arana-db/parser/test_driver"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/seata/seata-go/pkg/datasource/sql/parser"
+	_ "github.com/seata/seata-go/pkg/util/log"
+)
+
+func TestBuildSelectSQLByMultiUpdate(t *testing.T) {
+	var builder MySQLMultiUpdateUndoLogBuilder
+	tests := []struct {
+		name            string
+		sourceQuery     string
+		sourceQueryArgs []driver.Value
+		expectQuery     string
+		expectQueryArgs []driver.Value
+	}{
+		{
+			sourceQuery:     "update t_user set name = ?, age = ? where id = ?;update t_user set name = ?, age = ? where id = ?;",
+			sourceQueryArgs: []driver.Value{"Jack", 1, 100, "TOM", 2, 200},
+			expectQuery:     "SELECT SQL_NO_CACHE name,age FROM t_user WHERE id=? OR id=? FOR UPDATE",
+			expectQueryArgs: []driver.Value{100, 200},
+		},
+		{
+			sourceQuery:     "update t_user set name = ?, age = ? where id = ? and name = 'Jack' and age between ? and ?;update t_user set name = ?, age = ? where id = ? and name = 'Jack2' and age between ? and ?",
+			sourceQueryArgs: []driver.Value{"Jack", 1, 100, 18, 28, "Jack2", 2, 200, 28, 38},
+			expectQuery:     "SELECT SQL_NO_CACHE name,age FROM t_user WHERE id=? AND name=_UTF8MB4Jack AND age BETWEEN ? AND ? OR id=? AND name=_UTF8MB4Jack2 AND age BETWEEN ? AND ? FOR UPDATE",
+			expectQueryArgs: []driver.Value{100, 18, 28, 200, 28, 38},
+		},
+		{
+			sourceQuery:     "update t_user set name = ?, age = ? where id = ? and name = 'Jack' and age in (?,?);update t_user set name = ?, age = ? where id = ? and name = 'Jack2' and age in (?,?)",
+			sourceQueryArgs: []driver.Value{"Jack", 1, 100, 18, 28, "Jack2", 2, 200, 48, 58},
+			expectQuery:     "SELECT SQL_NO_CACHE name,age FROM t_user WHERE id=? AND name=_UTF8MB4Jack AND age IN (?,?) OR id=? AND name=_UTF8MB4Jack2 AND age IN (?,?) FOR UPDATE",
+			expectQueryArgs: []driver.Value{100, 18, 28, 200, 48, 58},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c, err := parser.DoParser(tt.sourceQuery)
+			assert.Nil(t, err)
+			var updateStmts []*ast.UpdateStmt
+			for _, v := range c.MultiStmt {
+				updateStmts = append(updateStmts, v.UpdateStmt)
+			}
+
+			query, args, err := builder.buildBeforeImageSQL(updateStmts, tt.sourceQueryArgs)
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expectQuery, query)
+			assert.Equal(t, tt.expectQueryArgs, args)
+		})
+	}
+
+	sourceQuery := "update t_user set name = ?, age = ? where kk between ? and ? and id = ? and addr in(?,?) and age > ? order by name desc;update t_user set name = ?, age = ? where kk between ? and ? and id = ? and addr in(?,?) and age > ? order by name"
+	sourceQueryArgs := []driver.Value{"Jack", 1, 10, 20, 17, "Beijing", "Guangzhou", 18, 2, "Jack2", 1, 10, 20, 17, "Beijing", "Guangzhou", 18, 2}
+	c, err := parser.DoParser(sourceQuery)
+	assert.NoError(t, err)
+	var updateStmts []*ast.UpdateStmt
+	for _, v := range c.MultiStmt {
+		updateStmts = append(updateStmts, v.UpdateStmt)
+	}
+	_, _, err = builder.buildBeforeImageSQL(updateStmts, sourceQueryArgs)
+	assert.Error(t, err)
+	assert.Equal(t, err.Error(), "multi update SQL with orderBy condition is not support yet")
+}

pkg/datasource/sql/undo/builder/mysql_update_undo_log_builder.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/arana-db/parser/ast"
 	"github.com/arana-db/parser/format"
+
 	"github.com/seata/seata-go/pkg/datasource/sql/types"
 	"github.com/seata/seata-go/pkg/datasource/sql/undo"
 	"github.com/seata/seata-go/pkg/util/bytes"