fix: refactor UpdateBatch method (#5295)

This commit is contained in:
Uzziah 2023-08-14 14:25:26 +08:00 committed by GitHub
parent 590fcd9b8f
commit a5eda3267a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 526 additions and 51 deletions

View File

@ -14,6 +14,7 @@
- [refactor cache/redis: Use redisConfig to receive incoming JSON (previously using a map)](https://github.com/beego/beego/pull/5268)
- [fix: refactor DeleteSQL method](https://github.com/beego/beego/pull/5271)
- [fix: refactor UpdateSQL method](https://github.com/beego/beego/pull/5274)
- [fix: refactor UpdateBatch method](https://github.com/beego/beego/pull/5295)
## ORM refactoring
- [introducing internal/models pkg](https://github.com/beego/beego/pull/5238)

View File

@ -819,58 +819,8 @@ func (d *dbBase) UpdateBatch(ctx context.Context, q dbQuerier, qs *querySet, mi
join := tables.getJoinSQL()
var query, T string
query := d.UpdateBatchSQL(mi, columns, values, specifyIndexes, join, where)
Q := d.ins.TableQuote()
if d.ins.SupportUpdateJoin() {
T = "T0."
}
cols := make([]string, 0, len(columns))
for i, v := range columns {
col := fmt.Sprintf("%s%s%s%s", T, Q, v, Q)
if c, ok := values[i].(colValue); ok {
switch c.opt {
case ColAdd:
cols = append(cols, col+" = "+col+" + ?")
case ColMinus:
cols = append(cols, col+" = "+col+" - ?")
case ColMultiply:
cols = append(cols, col+" = "+col+" * ?")
case ColExcept:
cols = append(cols, col+" = "+col+" / ?")
case ColBitAnd:
cols = append(cols, col+" = "+col+" & ?")
case ColBitRShift:
cols = append(cols, col+" = "+col+" >> ?")
case ColBitLShift:
cols = append(cols, col+" = "+col+" << ?")
case ColBitXOR:
cols = append(cols, col+" = "+col+" ^ ?")
case ColBitOr:
cols = append(cols, col+" = "+col+" | ?")
}
values[i] = c.value
} else {
cols = append(cols, col+" = ?")
}
}
sets := strings.Join(cols, ", ") + " "
if d.ins.SupportUpdateJoin() {
query = fmt.Sprintf("UPDATE %s%s%s T0 %s%sSET %s%s", Q, mi.Table, Q, specifyIndexes, join, sets, where)
} else {
supQuery := fmt.Sprintf("SELECT T0.%s%s%s FROM %s%s%s T0 %s%s%s",
Q, mi.Fields.Pk.Column, Q,
Q, mi.Table, Q,
specifyIndexes, join, where)
query = fmt.Sprintf("UPDATE %s%s%s SET %sWHERE %s%s%s IN ( %s )", Q, mi.Table, Q, sets, Q, mi.Fields.Pk.Column, Q, supQuery)
}
d.ins.ReplaceMarks(&query)
res, err := q.ExecContext(ctx, query, values...)
if err == nil {
return res.RowsAffected()
@ -878,6 +828,111 @@ func (d *dbBase) UpdateBatch(ctx context.Context, q dbQuerier, qs *querySet, mi
return 0, err
}
func (d *dbBase) UpdateBatchSQL(mi *models.ModelInfo, cols []string, values []interface{}, specifyIndexes, join, where string) string {
quote := d.ins.TableQuote()
buf := buffers.Get()
defer buffers.Put(buf)
_, _ = buf.WriteString("UPDATE ")
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(mi.Table)
_, _ = buf.WriteString(quote)
if d.ins.SupportUpdateJoin() {
_, _ = buf.WriteString(" T0 ")
_, _ = buf.WriteString(specifyIndexes)
_, _ = buf.WriteString(join)
d.buildSetSQL(buf, cols, values)
_, _ = buf.WriteString(" ")
_, _ = buf.WriteString(where)
} else {
_, _ = buf.WriteString(" ")
d.buildSetSQL(buf, cols, values)
_, _ = buf.WriteString(" WHERE ")
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(mi.Fields.Pk.Column)
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(" IN ( ")
_, _ = buf.WriteString("SELECT T0.")
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(mi.Fields.Pk.Column)
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(" FROM ")
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(mi.Table)
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(" T0 ")
_, _ = buf.WriteString(specifyIndexes)
_, _ = buf.WriteString(join)
_, _ = buf.WriteString(where)
_, _ = buf.WriteString(" )")
}
query := buf.String()
d.ins.ReplaceMarks(&query)
return query
}
func (d *dbBase) buildSetSQL(buf buffers.Buffer, cols []string, values []interface{}) {
var owner string
quote := d.ins.TableQuote()
if d.ins.SupportUpdateJoin() {
owner = "T0."
}
_, _ = buf.WriteString("SET ")
for i, v := range cols {
if i > 0 {
_, _ = buf.WriteString(", ")
}
_, _ = buf.WriteString(owner)
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(v)
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(" = ")
if c, ok := values[i].(colValue); ok {
_, _ = buf.WriteString(owner)
_, _ = buf.WriteString(quote)
_, _ = buf.WriteString(v)
_, _ = buf.WriteString(quote)
switch c.opt {
case ColAdd:
_, _ = buf.WriteString(" + ?")
case ColMinus:
_, _ = buf.WriteString(" - ?")
case ColMultiply:
_, _ = buf.WriteString(" * ?")
case ColExcept:
_, _ = buf.WriteString(" / ?")
case ColBitAnd:
_, _ = buf.WriteString(" & ?")
case ColBitRShift:
_, _ = buf.WriteString(" >> ?")
case ColBitLShift:
_, _ = buf.WriteString(" << ?")
case ColBitXOR:
_, _ = buf.WriteString(" ^ ?")
case ColBitOr:
_, _ = buf.WriteString(" | ?")
}
values[i] = c.value
} else {
_, _ = buf.WriteString("?")
}
}
}
// delete related records.
// do UpdateBanch or DeleteBanch by condition of tables' relationship.
func (d *dbBase) deleteRels(ctx context.Context, q dbQuerier, mi *models.ModelInfo, args []interface{}, tz *time.Location) error {

View File

@ -15,6 +15,7 @@
package orm
import (
"github.com/beego/beego/v2/client/orm/internal/buffers"
"testing"
"github.com/stretchr/testify/assert"
@ -229,3 +230,421 @@ func TestDbBase_DeleteSQL(t *testing.T) {
})
}
}
func TestDbBase_buildSetSQL(t *testing.T) {
testCases := []struct {
name string
db *dbBase
columns []string
values []interface{}
wantRes string
wantValues []interface{}
}{
{
name: "set add/mul operator by dbBase",
db: &dbBase{
ins: &dbBase{},
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColAdd,
value: 12,
},
colValue{
opt: ColMultiply,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET T0.`name` = ?, T0.`age` = T0.`age` + ?, T0.`score` = T0.`score` * ?",
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set min/except operator by dbBase",
db: &dbBase{
ins: &dbBase{},
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColMinus,
value: 12,
},
colValue{
opt: ColExcept,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET T0.`name` = ?, T0.`age` = T0.`age` - ?, T0.`score` = T0.`score` / ?",
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set bitRShift/bitLShift operator by dbBase",
db: &dbBase{
ins: &dbBase{},
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColBitRShift,
value: 12,
},
colValue{
opt: ColBitLShift,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET T0.`name` = ?, T0.`age` = T0.`age` >> ?, T0.`score` = T0.`score` << ?",
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set bitAnd/bitOr/bitXOR operator by dbBase",
db: &dbBase{
ins: &dbBase{},
},
columns: []string{"count", "age", "score"},
values: []interface{}{
colValue{
opt: ColBitAnd,
value: 28,
},
colValue{
opt: ColBitOr,
value: 12,
},
colValue{
opt: ColBitXOR,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET T0.`count` = T0.`count` & ?, T0.`age` = T0.`age` | ?, T0.`score` = T0.`score` ^ ?",
wantValues: []interface{}{int64(28), int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set add/mul operator by dbBasePostgres",
db: &dbBase{
ins: newdbBasePostgres(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColAdd,
value: 12,
},
colValue{
opt: ColMultiply,
value: 2,
},
"test_origin_name",
18,
},
wantRes: `SET "name" = ?, "age" = "age" + ?, "score" = "score" * ?`,
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set min/except operator by dbBasePostgres",
db: &dbBase{
ins: newdbBasePostgres(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColMinus,
value: 12,
},
colValue{
opt: ColExcept,
value: 2,
},
"test_origin_name",
18,
},
wantRes: `SET "name" = ?, "age" = "age" - ?, "score" = "score" / ?`,
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set bitRShift/bitLShift operator by dbBasePostgres",
db: &dbBase{
ins: newdbBasePostgres(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColBitRShift,
value: 12,
},
colValue{
opt: ColBitLShift,
value: 2,
},
"test_origin_name",
18,
},
wantRes: `SET "name" = ?, "age" = "age" >> ?, "score" = "score" << ?`,
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set bitAnd/bitOr/bitXOR operator by dbBasePostgres",
db: &dbBase{
ins: newdbBasePostgres(),
},
columns: []string{"count", "age", "score"},
values: []interface{}{
colValue{
opt: ColBitAnd,
value: 28,
},
colValue{
opt: ColBitOr,
value: 12,
},
colValue{
opt: ColBitXOR,
value: 2,
},
"test_origin_name",
18,
},
wantRes: `SET "count" = "count" & ?, "age" = "age" | ?, "score" = "score" ^ ?`,
wantValues: []interface{}{int64(28), int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set add/mul operator by dbBaseSqlite",
db: &dbBase{
ins: newdbBaseSqlite(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColAdd,
value: 12,
},
colValue{
opt: ColMultiply,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET `name` = ?, `age` = `age` + ?, `score` = `score` * ?",
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set min/except operator by dbBaseSqlite",
db: &dbBase{
ins: newdbBaseSqlite(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColMinus,
value: 12,
},
colValue{
opt: ColExcept,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET `name` = ?, `age` = `age` - ?, `score` = `score` / ?",
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set bitRShift/bitLShift operator by dbBaseSqlite",
db: &dbBase{
ins: newdbBaseSqlite(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColBitRShift,
value: 12,
},
colValue{
opt: ColBitLShift,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET `name` = ?, `age` = `age` >> ?, `score` = `score` << ?",
wantValues: []interface{}{"test_name", int64(12), int64(2), "test_origin_name", 18},
},
{
name: "set bitAnd/bitOr/bitXOR operator by dbBaseSqlite",
db: &dbBase{
ins: newdbBaseSqlite(),
},
columns: []string{"count", "age", "score"},
values: []interface{}{
colValue{
opt: ColBitAnd,
value: 28,
},
colValue{
opt: ColBitOr,
value: 12,
},
colValue{
opt: ColBitXOR,
value: 2,
},
"test_origin_name",
18,
},
wantRes: "SET `count` = `count` & ?, `age` = `age` | ?, `score` = `score` ^ ?",
wantValues: []interface{}{int64(28), int64(12), int64(2), "test_origin_name", 18},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
buf := buffers.Get()
defer buffers.Put(buf)
tc.db.buildSetSQL(buf, tc.columns, tc.values)
assert.Equal(t, tc.wantRes, buf.String())
assert.Equal(t, tc.wantValues, tc.values)
})
}
}
func TestDbBase_UpdateBatchSQL(t *testing.T) {
mi := &models.ModelInfo{
Table: "test_tab",
Fields: &models.Fields{
Pk: &models.FieldInfo{
Column: "test_id",
},
},
}
testCases := []struct {
name string
db *dbBase
columns []string
values []interface{}
specifyIndexes string
join string
where string
wantRes string
}{
{
name: "update batch by dbBase",
db: &dbBase{
ins: &dbBase{},
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColAdd,
value: 12,
},
colValue{
opt: ColMultiply,
value: 2,
},
"test_origin_name",
18,
},
specifyIndexes: " USE INDEX(`name`) ",
join: "LEFT OUTER JOIN `test_tab_2` T1 ON T1.`id` = T0.`test_id` ",
where: "WHERE T0.`name` = ? AND T1.`age` = ?",
wantRes: "UPDATE `test_tab` T0 USE INDEX(`name`) LEFT OUTER JOIN `test_tab_2` T1 ON T1.`id` = T0.`test_id` SET T0.`name` = ?, T0.`age` = T0.`age` + ?, T0.`score` = T0.`score` * ? WHERE T0.`name` = ? AND T1.`age` = ?",
},
{
name: "update batch by dbBasePostgres",
db: &dbBase{
ins: newdbBasePostgres(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColAdd,
value: 12,
},
colValue{
opt: ColMultiply,
value: 2,
},
"test_origin_name",
18,
},
specifyIndexes: ` USE INDEX("name") `,
join: `LEFT OUTER JOIN "test_tab_2" T1 ON T1."id" = T0."test_id" `,
where: `WHERE T0."name" = ? AND T1."age" = ?`,
wantRes: `UPDATE "test_tab" SET "name" = $1, "age" = "age" + $2, "score" = "score" * $3 WHERE "test_id" IN ( SELECT T0."test_id" FROM "test_tab" T0 USE INDEX("name") LEFT OUTER JOIN "test_tab_2" T1 ON T1."id" = T0."test_id" WHERE T0."name" = $4 AND T1."age" = $5 )`,
},
{
name: "update batch by dbBaseSqlite",
db: &dbBase{
ins: newdbBaseSqlite(),
},
columns: []string{"name", "age", "score"},
values: []interface{}{
"test_name",
colValue{
opt: ColAdd,
value: 12,
},
colValue{
opt: ColMultiply,
value: 2,
},
"test_origin_name",
18,
},
specifyIndexes: " USE INDEX(`name`) ",
join: "LEFT OUTER JOIN `test_tab_2` T1 ON T1.`id` = T0.`test_id` ",
where: "WHERE T0.`name` = ? AND T1.`age` = ?",
wantRes: "UPDATE `test_tab` SET `name` = ?, `age` = `age` + ?, `score` = `score` * ? WHERE `test_id` IN ( SELECT T0.`test_id` FROM `test_tab` T0 USE INDEX(`name`) LEFT OUTER JOIN `test_tab_2` T1 ON T1.`id` = T0.`test_id` WHERE T0.`name` = ? AND T1.`age` = ? )",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
res := tc.db.UpdateBatchSQL(mi, tc.columns, tc.values, tc.specifyIndexes, tc.join, tc.where)
assert.Equal(t, tc.wantRes, res)
})
}
}