do not use single quotes to wrap pk data

actiontech · May 31, 2024 · e6ea655 · e6ea655
1 parent f66cb1c
commit e6ea655
Showing 1 changed file with 15 additions and 2 deletions.
diff --git a/sqle/driver/mysql/rollback.go b/sqle/driver/mysql/rollback.go
@@ -11,6 +11,7 @@ import (
 	"github.com/actiontech/sqle/sqle/errors"
 
 	"github.com/pingcap/parser/ast"
+	"github.com/pingcap/parser/format"
 	_model "github.com/pingcap/parser/model"
 	parserMysql "github.com/pingcap/parser/mysql"
 )
@@ -626,7 +627,7 @@ func (i *MysqlDriverImpl) generateUpdateRollbackSql(stmt *ast.UpdateStmt) (strin
 					colChanged = true
 					if isPk {
 						isPkChanged = true
-						pkValue = util.ExprFormat(l.Expr)
+						pkValue = restore(l.Expr)
 					}
 				}
 			}
@@ -647,7 +648,7 @@ func (i *MysqlDriverImpl) generateUpdateRollbackSql(stmt *ast.UpdateStmt) (strin
 			}
 			if isPk {
 				if isPkChanged {
-					where = append(where, fmt.Sprintf("%s = '%s'", name, pkValue))
+					where = append(where, fmt.Sprintf("%s = %s", name, pkValue))
 				} else {
 					where = append(where, fmt.Sprintf("%s = %s", name, v))
 
@@ -660,6 +661,18 @@ func (i *MysqlDriverImpl) generateUpdateRollbackSql(stmt *ast.UpdateStmt) (strin
 	return rollbackSql, "", nil
 }
 
+// 还原抽象语法树节点至SQL
+func restore(node ast.Node) (sql string) {
+	var buf strings.Builder
+	rc := format.NewRestoreCtx(format.DefaultRestoreFlags, &buf)
+
+	if err := node.Restore(rc); err != nil {
+		return
+	}
+	sql = buf.String()
+	return
+}
+
 // getRecords select all data which will be update or delete.
 func (i *MysqlDriverImpl) getRecords(tableName *ast.TableName, tableAlias string, where ast.ExprNode,
 	order *ast.OrderByClause, limit int64) ([]map[string]sql.NullString, error) {