diff --git a/pkg/datasource/sql/exec/hook/basic_undo_builder.go b/pkg/datasource/sql/exec/hook/basic_undo_builder.go index 3e39f860d..b767d60ef 100644 --- a/pkg/datasource/sql/exec/hook/basic_undo_builder.go +++ b/pkg/datasource/sql/exec/hook/basic_undo_builder.go @@ -19,8 +19,14 @@ package exec import ( "context" + "fmt" + "github.com/pingcap/tidb/parser/ast" + "github.com/pingcap/tidb/parser/format" + "github.com/seata/seata-go/pkg/common/bytes" + "github.com/seata/seata-go/pkg/common/log" "github.com/seata/seata-go/pkg/datasource/sql/exec" + "github.com/seata/seata-go/pkg/datasource/sql/parser" "github.com/seata/seata-go/pkg/datasource/sql/types" ) @@ -36,3 +42,42 @@ func (u *BasicUndoBuilder) buildRowImages(ctx context.Context, execCtx *exec.Exe func (u *BasicUndoBuilder) buildRecordImage(ctx context.Context, execCtx *exec.ExecContext) ([]*types.RecordImage, error) { panic("implement me") } + +func (u *BasicUndoBuilder) buildSelectSQLByUpdate(query string) (string, error) { + p, err := parser.DoParser(query) + if err != nil { + return "", err + } + + if p.UpdateStmt == nil { + return "", fmt.Errorf("invalid update stmt") + } + + updateColumns := p.UpdateStmt.List + fields := []*ast.SelectField{} + + for _, column := range updateColumns { + fields = append(fields, &ast.SelectField{ + Expr: &ast.ColumnNameExpr{ + Name: column.Column, + }, + }) + } + + selStmt := ast.SelectStmt{ + SelectStmtOpts: &ast.SelectStmtOpts{}, + From: p.UpdateStmt.TableRefs, + Where: p.UpdateStmt.Where, + Fields: &ast.FieldList{Fields: fields}, + OrderBy: p.UpdateStmt.Order, + Limit: p.UpdateStmt.Limit, + TableHints: p.UpdateStmt.TableHints, + } + + b := bytes.NewByteBuffer([]byte{}) + selStmt.Restore(format.NewRestoreCtx(format.RestoreKeyWordUppercase, b)) + sql := string(b.Bytes()) + log.Infof("build select sql by update query, sql {}", sql) + + return sql, nil +} diff --git a/pkg/datasource/sql/exec/hook/basic_undo_builder_test.go b/pkg/datasource/sql/exec/hook/basic_undo_builder_test.go new file mode 100644 index 000000000..1e49ddff0 --- /dev/null +++ b/pkg/datasource/sql/exec/hook/basic_undo_builder_test.go @@ -0,0 +1,54 @@ +/* + * 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 exec + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBuildSelectSQLByUpdate(t *testing.T) { + builder := BasicUndoBuilder{} + + tests := []struct { + name string + query string + expect string + }{ + { + query: "update t_user set name = ?, age = ? where id = ?", + expect: "SELECT SQL_NO_CACHE name,age FROM t_user WHERE id=?", + }, + { + query: "update t_user set name = ?, age = ? where id = ? and addr = ? order by id, name desc", + expect: "SELECT SQL_NO_CACHE name,age FROM t_user WHERE id=? AND addr=? ORDER BY id,name DESC", + }, + { + query: "update t_user set name = ?, age = ? where id = ? and addr = ? order by id, name desc limit 99", + expect: "SELECT SQL_NO_CACHE name,age FROM t_user WHERE id=? AND addr=? ORDER BY id,name DESC LIMIT 99", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sql, err := builder.buildSelectSQLByUpdate(tt.query) + assert.Nil(t, err) + assert.Equal(t, tt.expect, sql) + }) + } +}