feat(mysql): support for UPDATE ... ORDER BY ... LIMIT ...

uptrace · Nov 4, 2024 · cfed0ca · cfed0ca
1 parent c1541c2
commit cfed0ca
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 0 deletions.
diff --git a/internal/dbtest/query_test.go b/internal/dbtest/query_test.go
@@ -1578,6 +1578,13 @@ func TestQuery(t *testing.T) {
 					Where("src.id = story.id").Order("src.id").Limit(1)
 			},
 		},
+		{
+			id: 171,
+			query: func(db *bun.DB) schema.QueryAppender {
+				// UPDATE ... SET ... ORDER BY ... LIMIT ... (MySQL, MariaDB)
+				return db.NewUpdate().Model(new(Story)).Set("name = ?", "new-name").WherePK().Order("id").Limit(1)
+			},
+		},
 	}
 
 	timeRE := regexp.MustCompile(`'2\d{3}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}(\.\d+)?(\+\d{2}:\d{2})?'`)

diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mariadb-171 b/internal/dbtest/testdata/snapshots/TestQuery-mariadb-171
@@ -0,0 +1 @@
+UPDATE `stories` AS `story` SET name = 'new-name' WHERE (`story`.`id` = NULL) ORDER BY `id` LIMIT 1
diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-171 b/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-171
@@ -0,0 +1 @@
+bun: limit is not supported for current dialect
diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mysql8-171 b/internal/dbtest/testdata/snapshots/TestQuery-mysql8-171
@@ -0,0 +1 @@
+UPDATE `stories` AS `story` SET name = 'new-name' WHERE (`story`.`id` = NULL) ORDER BY `id` LIMIT 1
diff --git a/internal/dbtest/testdata/snapshots/TestQuery-pg-171 b/internal/dbtest/testdata/snapshots/TestQuery-pg-171
@@ -0,0 +1 @@
+bun: limit is not supported for current dialect
diff --git a/internal/dbtest/testdata/snapshots/TestQuery-pgx-171 b/internal/dbtest/testdata/snapshots/TestQuery-pgx-171
@@ -0,0 +1 @@
+bun: limit is not supported for current dialect
diff --git a/internal/dbtest/testdata/snapshots/TestQuery-sqlite-171 b/internal/dbtest/testdata/snapshots/TestQuery-sqlite-171
@@ -0,0 +1 @@
+bun: limit is not supported for current dialect
diff --git a/query_update.go b/query_update.go
@@ -15,6 +15,7 @@ import (
 
 type UpdateQuery struct {
 	whereBaseQuery
+	orderLimitOffsetQuery
 	returningQuery
 	customValueQuery
 	setQuery
@@ -200,6 +201,31 @@ func (q *UpdateQuery) WhereAllWithDeleted() *UpdateQuery {
 	return q
 }
 
+// ------------------------------------------------------------------------------
+func (q *UpdateQuery) Order(orders ...string) *UpdateQuery {
+	if q.Dialect().Name() != dialect.MySQL {
+		q.err = errors.New("bun: order is not supported for current dialect")
+	}
+	q.addOrder(orders...)
+	return q
+}
+
+func (q *UpdateQuery) OrderExpr(query string, args ...interface{}) *UpdateQuery {
+	if q.Dialect().Name() != dialect.MySQL {
+		q.err = errors.New("bun: order is not supported for current dialect")
+	}
+	q.addOrderExpr(query, args...)
+	return q
+}
+
+func (q *UpdateQuery) Limit(n int) *UpdateQuery {
+	if q.Dialect().Name() != dialect.MySQL {
+		q.err = errors.New("bun: limit is not supported for current dialect")
+	}
+	q.addLimit(n)
+	return q
+}
+
 //------------------------------------------------------------------------------
 
 // Returning adds a RETURNING clause to the query.
@@ -278,6 +304,16 @@ func (q *UpdateQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, e
 		return nil, err
 	}
 
+	b, err = q.appendOrder(fmter, b)
+	if err != nil {
+		return nil, err
+	}
+
+	b, err = q.appendLimitOffset(fmter, b)
+	if err != nil {
+		return nil, err
+	}
+
 	if q.hasFeature(feature.Returning) && q.hasReturning() {
 		b = append(b, " RETURNING "...)
 		b, err = q.appendReturning(fmter, b)