Add order by on update statement

Author: jeffgbutler

src/main/java/org/mybatis/dynamic/sql/update/UpdateDSL.java

@@ -16,6 +16,8 @@
 package org.mybatis.dynamic.sql.update;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
@@ -24,8 +26,10 @@
 
 import org.jetbrains.annotations.NotNull;
 import org.mybatis.dynamic.sql.BasicColumn;
+import org.mybatis.dynamic.sql.SortSpecification;
 import org.mybatis.dynamic.sql.SqlColumn;
 import org.mybatis.dynamic.sql.SqlTable;
+import org.mybatis.dynamic.sql.common.OrderByModel;
 import org.mybatis.dynamic.sql.configuration.StatementConfiguration;
 import org.mybatis.dynamic.sql.select.SelectModel;
 import org.mybatis.dynamic.sql.util.AbstractColumnMapping;
@@ -52,6 +56,7 @@ public class UpdateDSL<R> extends AbstractWhereSupport<UpdateDSL<R>.UpdateWhereB
     private UpdateWhereBuilder whereBuilder;
     private final StatementConfiguration statementConfiguration = new StatementConfiguration();
     private Long limit;
+    private OrderByModel orderByModel;
 
     private UpdateDSL(SqlTable table, String tableAlias, Function<UpdateModel, R> adapterFunction) {
         this.table = Objects.requireNonNull(table);
@@ -77,6 +82,15 @@ public UpdateDSL<R> limit(long limit) {
         return this;
     }
 
+    public UpdateDSL<R> orderBy(SortSpecification... columns) {
+        return orderBy(Arrays.asList(columns));
+    }
+
+    public UpdateDSL<R> orderBy(Collection<SortSpecification> columns) {
+        orderByModel = OrderByModel.of(columns);
+        return this;
+    }
+
     /**
      * WARNING! Calling this method could result in an update statement that updates
      * all rows in a table.
@@ -89,7 +103,8 @@ public R build() {
         UpdateModel.Builder updateModelBuilder = UpdateModel.withTable(table)
                 .withTableAlias(tableAlias)
                 .withColumnMappings(columnMappings)
-                .withLimit(limit);
+                .withLimit(limit)
+                .withOrderByModel(orderByModel);
 
         if (whereBuilder != null) {
             updateModelBuilder.withWhereModel(whereBuilder.buildWhereModel());
@@ -187,6 +202,15 @@ public UpdateDSL<R> limit(long limit) {
             return UpdateDSL.this.limit(limit);
         }
 
+        public UpdateDSL<R> orderBy(SortSpecification... columns) {
+            return orderBy(Arrays.asList(columns));
+        }
+
+        public UpdateDSL<R> orderBy(Collection<SortSpecification> columns) {
+            orderByModel = OrderByModel.of(columns);
+            return UpdateDSL.this;
+        }
+
         @NotNull
         @Override
         public R build() {

src/main/java/org/mybatis/dynamic/sql/update/UpdateModel.java

@@ -24,6 +24,7 @@
 
 import org.jetbrains.annotations.NotNull;
 import org.mybatis.dynamic.sql.SqlTable;
+import org.mybatis.dynamic.sql.common.OrderByModel;
 import org.mybatis.dynamic.sql.exception.InvalidSqlException;
 import org.mybatis.dynamic.sql.render.RenderingStrategy;
 import org.mybatis.dynamic.sql.update.render.UpdateRenderer;
@@ -38,13 +39,15 @@ public class UpdateModel {
     private final WhereModel whereModel;
     private final List<AbstractColumnMapping> columnMappings;
     private final Long limit;
+    private final OrderByModel orderByModel;
 
     private UpdateModel(Builder builder) {
         table = Objects.requireNonNull(builder.table);
         whereModel = builder.whereModel;
         columnMappings = Objects.requireNonNull(builder.columnMappings);
         tableAlias = builder.tableAlias;
         limit = builder.limit;
+        orderByModel = builder.orderByModel;
 
         if (columnMappings.isEmpty()) {
             throw new InvalidSqlException(Messages.getString("ERROR.17")); //$NON-NLS-1$
@@ -71,6 +74,10 @@ public Optional<Long> limit() {
         return Optional.ofNullable(limit);
     }
 
+    public Optional<OrderByModel> orderByModel() {
+        return Optional.ofNullable(orderByModel);
+    }
+
     @NotNull
     public UpdateStatementProvider render(RenderingStrategy renderingStrategy) {
         return UpdateRenderer.withUpdateModel(this)
@@ -89,6 +96,7 @@ public static class Builder {
         private WhereModel whereModel;
         private final List<AbstractColumnMapping> columnMappings = new ArrayList<>();
         private Long limit;
+        private OrderByModel orderByModel;
 
         public Builder withTable(SqlTable table) {
             this.table = table;
@@ -115,6 +123,11 @@ public Builder withLimit(Long limit) {
             return this;
         }
 
+        public Builder withOrderByModel(OrderByModel orderByModel) {
+            this.orderByModel = orderByModel;
+            return this;
+        }
+
         public UpdateModel build() {
             return new UpdateModel(this);
         }

src/main/java/org/mybatis/dynamic/sql/update/render/UpdateRenderer.java

@@ -22,6 +22,8 @@
 import java.util.stream.Collectors;
 
 import org.mybatis.dynamic.sql.SqlTable;
+import org.mybatis.dynamic.sql.common.OrderByModel;
+import org.mybatis.dynamic.sql.common.OrderByRenderer;
 import org.mybatis.dynamic.sql.exception.InvalidSqlException;
 import org.mybatis.dynamic.sql.render.ExplicitTableAliasCalculator;
 import org.mybatis.dynamic.sql.render.RenderingStrategy;
@@ -53,6 +55,7 @@ public UpdateStatementProvider render() {
         fragmentCollector.add(calculateUpdateStatementStart());
         fragmentCollector.add(calculateSetPhrase());
         calculateWhereClause().ifPresent(fragmentCollector::add);
+        calculateOrderByClause().ifPresent(fragmentCollector::add);
         calculateLimitClause().ifPresent(fragmentCollector::add);
 
         return toUpdateStatementProvider(fragmentCollector);
@@ -128,6 +131,14 @@ private FragmentAndParameters renderLimitClause(Long limit) {
                 .build();
     }
 
+    private Optional<FragmentAndParameters> calculateOrderByClause() {
+        return updateModel.orderByModel().map(this::renderOrderByClause);
+    }
+
+    private FragmentAndParameters renderOrderByClause(OrderByModel orderByModel) {
+        return new OrderByRenderer().render(orderByModel);
+    }
+
     public static Builder withUpdateModel(UpdateModel updateModel) {
         return new Builder().withUpdateModel(updateModel);
     }

src/test/java/examples/mariadb/MariaDBTest.java

@@ -231,6 +231,47 @@ void testUpdateWithLimit() {
         }
     }
 
+    @Test
+    void testUpdateWithOrderBy() {
+        try (SqlSession session = sqlSessionFactory.openSession()) {
+            CommonUpdateMapper mapper = session.getMapper(CommonUpdateMapper.class);
+
+            UpdateStatementProvider updateStatement = update(items)
+                    .set(amount).equalTo(add(amount, constant("100")))
+                    .orderBy(amount.descending())
+                    .build()
+                    .render(RenderingStrategies.MYBATIS3);
+
+            assertThat(updateStatement.getUpdateStatement())
+                    .isEqualTo("update items set amount = (amount + 100) order by amount DESC");
+            assertThat(updateStatement.getParameters()).isEmpty();
+
+            int rows = mapper.update(updateStatement);
+            assertThat(rows).isEqualTo(20);
+        }
+    }
+
+    @Test
+    void testUpdateWithOrderByAndLimit() {
+        try (SqlSession session = sqlSessionFactory.openSession()) {
+            CommonUpdateMapper mapper = session.getMapper(CommonUpdateMapper.class);
+
+            UpdateStatementProvider updateStatement = update(items)
+                    .set(amount).equalTo(add(amount, constant("100")))
+                    .orderBy(amount.descending())
+                    .limit(5)
+                    .build()
+                    .render(RenderingStrategies.MYBATIS3);
+
+            String expected = "update items set amount = (amount + 100) order by amount DESC limit #{parameters.p1}";
+            assertThat(updateStatement.getUpdateStatement()).isEqualTo(expected);
+            assertThat(updateStatement.getParameters()).containsOnly(entry("p1", 5L));
+
+            int rows = mapper.update(updateStatement);
+            assertThat(rows).isEqualTo(5);
+        }
+    }
+
     @Test
     void testUpdateWithWhereAndLimit() {
         try (SqlSession session = sqlSessionFactory.openSession()) {
@@ -252,4 +293,49 @@ void testUpdateWithWhereAndLimit() {
             assertThat(rows).isEqualTo(5);
         }
     }
+
+    @Test
+    void testUpdateWithWhereAndOrderBy() {
+        try (SqlSession session = sqlSessionFactory.openSession()) {
+            CommonUpdateMapper mapper = session.getMapper(CommonUpdateMapper.class);
+
+            UpdateStatementProvider updateStatement = update(items)
+                    .set(amount).equalTo(add(amount, constant("100")))
+                    .where(id, isLessThan(10))
+                    .orderBy(amount.descending())
+                    .build()
+                    .render(RenderingStrategies.MYBATIS3);
+
+            String expected = "update items set amount = (amount + 100) where id < #{parameters.p1,jdbcType=INTEGER} "
+                    + "order by amount DESC";
+            assertThat(updateStatement.getUpdateStatement()).isEqualTo(expected);
+            assertThat(updateStatement.getParameters()).containsOnly(entry("p1", 10));
+
+            int rows = mapper.update(updateStatement);
+            assertThat(rows).isEqualTo(9);
+        }
+    }
+
+    @Test
+    void testUpdateWithWhereAndOrderByAndLimit() {
+        try (SqlSession session = sqlSessionFactory.openSession()) {
+            CommonUpdateMapper mapper = session.getMapper(CommonUpdateMapper.class);
+
+            UpdateStatementProvider updateStatement = update(items)
+                    .set(amount).equalTo(add(amount, constant("100")))
+                    .where(id, isLessThan(10))
+                    .orderBy(amount)
+                    .limit(5)
+                    .build()
+                    .render(RenderingStrategies.MYBATIS3);
+
+            String expected = "update items set amount = (amount + 100) where id < #{parameters.p1,jdbcType=INTEGER} "
+                    + "order by amount limit #{parameters.p2}";
+            assertThat(updateStatement.getUpdateStatement()).isEqualTo(expected);
+            assertThat(updateStatement.getParameters()).containsOnly(entry("p1", 10), entry("p2", 5L));
+
+            int rows = mapper.update(updateStatement);
+            assertThat(rows).isEqualTo(5);
+        }
+    }
 }