From a5eda3267ad2ba46a0c38058637b06e545c399eb Mon Sep 17 00:00:00 2001 From: Uzziah <120019273+uzziahlin@users.noreply.github.com> Date: Mon, 14 Aug 2023 14:25:26 +0800 Subject: [PATCH] fix: refactor UpdateBatch method (#5295) --- CHANGELOG.md | 1 + client/orm/db.go | 157 +++++++++++----- client/orm/db_test.go | 419 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 526 insertions(+), 51 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b46f2239..1747197a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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) diff --git a/client/orm/db.go b/client/orm/db.go index 8bb4bec4..bded065f 100644 --- a/client/orm/db.go +++ b/client/orm/db.go @@ -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 { diff --git a/client/orm/db_test.go b/client/orm/db_test.go index 43fa3798..6a61f380 100644 --- a/client/orm/db_test.go +++ b/client/orm/db_test.go @@ -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) + }) + } +}