Support the ORDER BY clause for the UNION set operation

Author: nakamura-to

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/context/OrderByItem.java

@@ -0,0 +1,42 @@
+package org.seasar.doma.jdbc.criteria.context;
+
+import java.util.Objects;
+import org.seasar.doma.jdbc.criteria.def.PropertyDef;
+
+public interface OrderByItem {
+
+  void accept(Visitor visitor);
+
+  class Name implements OrderByItem {
+    public final PropertyDef<?> value;
+
+    public Name(PropertyDef<?> value) {
+      Objects.requireNonNull(value);
+      this.value = value;
+    }
+
+    @Override
+    public void accept(Visitor visitor) {
+      visitor.visit(this);
+    }
+  }
+
+  class Index implements OrderByItem {
+    public final int value;
+
+    public Index(int value) {
+      this.value = value;
+    }
+
+    @Override
+    public void accept(Visitor visitor) {
+      visitor.visit(this);
+    }
+  }
+
+  interface Visitor {
+    void visit(Name name);
+
+    void visit(Index index);
+  }
+}

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/context/SelectContext.java

@@ -21,7 +21,7 @@ public class SelectContext implements Context {
   public List<Criterion> where = new ArrayList<>();
   public final List<PropertyDef<?>> groupBy = new ArrayList<>();
   public List<Criterion> having = new ArrayList<>();
-  public final List<Pair<PropertyDef<?>, String>> orderBy = new ArrayList<>();
+  public final List<Pair<OrderByItem, String>> orderBy = new ArrayList<>();
   public Integer limit;
   public Integer offset;
   public ForUpdate forUpdate;

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/context/SetOperationContext.java

@@ -1,6 +1,9 @@
 package org.seasar.doma.jdbc.criteria.context;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
+import org.seasar.doma.internal.util.Pair;
 
 public interface SetOperationContext<ELEMENT> {
 
@@ -23,6 +26,7 @@ public <R> R accept(Visitor<R> visitor) {
   class Union<ELEMENT> implements SetOperationContext<ELEMENT> {
     public final SetOperationContext<ELEMENT> left;
     public final SetOperationContext<ELEMENT> right;
+    public final List<Pair<OrderByItem.Index, String>> orderBy = new ArrayList<>();
 
     public Union(SetOperationContext<ELEMENT> left, SetOperationContext<ELEMENT> right) {
       Objects.requireNonNull(left);
@@ -40,6 +44,7 @@ public <R> R accept(Visitor<R> visitor) {
   class UnionAll<ELEMENT> implements SetOperationContext<ELEMENT> {
     public final SetOperationContext<ELEMENT> left;
     public final SetOperationContext<ELEMENT> right;
+    public final List<Pair<OrderByItem.Index, String>> orderBy = new ArrayList<>();
 
     public UnionAll(SetOperationContext<ELEMENT> left, SetOperationContext<ELEMENT> right) {
       Objects.requireNonNull(left);

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/declaration/OrderByIndexDeclaration.java

@@ -0,0 +1,46 @@
+package org.seasar.doma.jdbc.criteria.declaration;
+
+import java.util.Objects;
+import org.seasar.doma.internal.util.Pair;
+import org.seasar.doma.jdbc.criteria.context.OrderByItem;
+import org.seasar.doma.jdbc.criteria.context.SetOperationContext;
+
+public class OrderByIndexDeclaration {
+
+  private final SetOperationContext<?> context;
+
+  public OrderByIndexDeclaration(SetOperationContext<?> context) {
+    Objects.requireNonNull(context);
+    this.context = context;
+  }
+
+  public void asc(int index) {
+    add(new Pair<>(new OrderByItem.Index(index), "asc"));
+  }
+
+  public void desc(int index) {
+    add(new Pair<>(new OrderByItem.Index(index), "desc"));
+  }
+
+  private void add(Pair<OrderByItem.Index, String> pair) {
+    context.accept(
+        new SetOperationContext.Visitor<Void>() {
+          @Override
+          public Void visit(SetOperationContext.Select<?> select) {
+            return null;
+          }
+
+          @Override
+          public Void visit(SetOperationContext.Union<?> union) {
+            union.orderBy.add(pair);
+            return null;
+          }
+
+          @Override
+          public Void visit(SetOperationContext.UnionAll<?> unionAll) {
+            unionAll.orderBy.add(pair);
+            return null;
+          }
+        });
+  }
+}

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/declaration/OrderByDeclaration.java renamed to doma-core/src/main/java/org/seasar/doma/jdbc/criteria/declaration/OrderByNameDeclaration.java

@@ -2,27 +2,28 @@
 
 import java.util.Objects;
 import org.seasar.doma.internal.util.Pair;
+import org.seasar.doma.jdbc.criteria.context.OrderByItem;
 import org.seasar.doma.jdbc.criteria.context.SelectContext;
 import org.seasar.doma.jdbc.criteria.def.PropertyDef;
 
-public class OrderByDeclaration {
+public class OrderByNameDeclaration {
 
   private final SelectContext context;
 
-  public OrderByDeclaration(SelectContext context) {
+  public OrderByNameDeclaration(SelectContext context) {
     Objects.requireNonNull(context);
     this.context = context;
   }
 
   public void asc(PropertyDef<?> propertyDef) {
     Objects.requireNonNull(propertyDef);
-    Pair<PropertyDef<?>, String> pair = new Pair<>(propertyDef, "asc");
+    Pair<OrderByItem, String> pair = new Pair<>(new OrderByItem.Name(propertyDef), "asc");
     context.orderBy.add(pair);
   }
 
   public void desc(PropertyDef<?> propertyDef) {
     Objects.requireNonNull(propertyDef);
-    Pair<PropertyDef<?>, String> pair = new Pair<>(propertyDef, "desc");
+    Pair<OrderByItem, String> pair = new Pair<>(new OrderByItem.Name(propertyDef), "desc");
     context.orderBy.add(pair);
   }
 }

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/declaration/SelectFromDeclaration.java

@@ -70,9 +70,9 @@ public void having(Consumer<HavingDeclaration> block) {
     block.accept(declaration);
   }
 
-  public void orderBy(Consumer<OrderByDeclaration> block) {
+  public void orderBy(Consumer<OrderByNameDeclaration> block) {
     Objects.requireNonNull(block);
-    OrderByDeclaration declaration = new OrderByDeclaration(context);
+    OrderByNameDeclaration declaration = new OrderByNameDeclaration(context);
     block.accept(declaration);
   }
 

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/declaration/SubSelectFromDeclaration.java

@@ -36,7 +36,7 @@ public SubSelectFromDeclaration where(Consumer<WhereDeclaration> block) {
     return this;
   }
 
-  public SubSelectFromDeclaration orderBy(Consumer<OrderByDeclaration> block) {
+  public SubSelectFromDeclaration orderBy(Consumer<OrderByNameDeclaration> block) {
     declaration.orderBy(block);
     return this;
   }

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/query/SelectBuilder.java

@@ -15,6 +15,7 @@
 import org.seasar.doma.jdbc.criteria.context.Criterion;
 import org.seasar.doma.jdbc.criteria.context.Join;
 import org.seasar.doma.jdbc.criteria.context.JoinKind;
+import org.seasar.doma.jdbc.criteria.context.OrderByItem;
 import org.seasar.doma.jdbc.criteria.context.SelectContext;
 import org.seasar.doma.jdbc.criteria.declaration.AggregateFunction;
 import org.seasar.doma.jdbc.criteria.def.EntityDef;
@@ -166,8 +167,20 @@ private void having() {
   private void orderBy() {
     if (!context.orderBy.isEmpty()) {
       buf.appendSql(" order by ");
-      for (Pair<PropertyDef<?>, String> pair : context.orderBy) {
-        column(pair.fst);
+      for (Pair<OrderByItem, String> pair : context.orderBy) {
+        pair.fst.accept(
+            new OrderByItem.Visitor() {
+
+              @Override
+              public void visit(OrderByItem.Name name) {
+                column(name.value);
+              }
+
+              @Override
+              public void visit(OrderByItem.Index index) {
+                buf.appendSql(String.valueOf(index.value));
+              }
+            });
         buf.appendSql(" " + pair.snd + ", ");
       }
       buf.cutBackSql(2);

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/query/SetOperationBuilder.java

@@ -1,12 +1,15 @@
 package org.seasar.doma.jdbc.criteria.query;
 
+import java.util.List;
 import java.util.Objects;
 import java.util.function.Function;
 import org.seasar.doma.internal.jdbc.sql.PreparedSqlBuilder;
+import org.seasar.doma.internal.util.Pair;
 import org.seasar.doma.jdbc.Config;
 import org.seasar.doma.jdbc.PreparedSql;
 import org.seasar.doma.jdbc.SqlKind;
 import org.seasar.doma.jdbc.SqlLogType;
+import org.seasar.doma.jdbc.criteria.context.OrderByItem;
 import org.seasar.doma.jdbc.criteria.context.SelectContext;
 import org.seasar.doma.jdbc.criteria.context.SetOperationContext;
 
@@ -22,13 +25,10 @@ public SetOperationBuilder(
       SetOperationContext<?> context,
       Function<String, String> commenter,
       SqlLogType sqlLogType) {
-    Objects.requireNonNull(config);
-    Objects.requireNonNull(context);
-    Objects.requireNonNull(commenter);
+    this.config = Objects.requireNonNull(config);
+    this.context = Objects.requireNonNull(context);
+    this.commenter = Objects.requireNonNull(commenter);
     Objects.requireNonNull(sqlLogType);
-    this.config = config;
-    this.context = context;
-    this.commenter = commenter;
     this.buf = new PreparedSqlBuilder(config, SqlKind.SELECT, sqlLogType);
   }
 
@@ -46,19 +46,49 @@ public Void visit(SetOperationContext.Select<?> select) {
 
           @Override
           public Void visit(SetOperationContext.Union<?> union) {
-            union.left.accept(this);
-            buf.appendSql(" union ");
-            union.right.accept(this);
+            union("union", union.left, union.right, union.orderBy);
             return null;
           }
 
           @Override
           public Void visit(SetOperationContext.UnionAll<?> unionAll) {
-            unionAll.left.accept(this);
-            buf.appendSql(" union all ");
-            unionAll.right.accept(this);
+            union("union all", unionAll.left, unionAll.right, unionAll.orderBy);
             return null;
           }
+
+          private void union(
+              String op,
+              SetOperationContext<?> left,
+              SetOperationContext<?> right,
+              List<Pair<OrderByItem.Index, String>> orderBy) {
+            String open;
+            String close;
+            if (orderBy.isEmpty()) {
+              open = "";
+              close = "";
+            } else {
+              open = "(";
+              close = ")";
+            }
+            buf.appendSql(open);
+            left.accept(this);
+            buf.appendSql(close + " " + op + " " + open);
+            right.accept(this);
+            buf.appendSql(close);
+            orderBy(orderBy);
+          }
+
+          private void orderBy(List<Pair<OrderByItem.Index, String>> orderBy) {
+            if (!orderBy.isEmpty()) {
+              buf.appendSql(" order by ");
+              for (Pair<OrderByItem.Index, String> pair : orderBy) {
+                OrderByItem.Index index = pair.fst;
+                buf.appendSql(String.valueOf(index.value));
+                buf.appendSql(" " + pair.snd + ", ");
+              }
+              buf.cutBackSql(2);
+            }
+          }
         });
     return buf.build(commenter);
   }

doma-core/src/main/java/org/seasar/doma/jdbc/criteria/statement/AbstractSetOperand.java

@@ -20,15 +20,15 @@ protected AbstractSetOperand(
   }
 
   @Override
-  public SetOperand<ELEMENT> union(SetOperand<ELEMENT> other) {
+  public SetOperator<ELEMENT> union(SetOperand<ELEMENT> other) {
     Objects.requireNonNull(other);
     SetOperationContext<ELEMENT> newContext =
         new SetOperationContext.Union<>(getContext(), other.getContext());
     return new NativeSqlSetStarting<>(config, newContext, objectProviderFactory);
   }
 
   @Override
-  public SetOperand<ELEMENT> unionAll(SetOperand<ELEMENT> other) {
+  public SetOperator<ELEMENT> unionAll(SetOperand<ELEMENT> other) {
     Objects.requireNonNull(other);
     SetOperationContext<ELEMENT> newContext =
         new SetOperationContext.UnionAll<>(getContext(), other.getContext());