parse select sql from update sql

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
+}

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)
+		})
+	}
+}