Merge pull request #432 from tradegecko/pg-grouping-elements-support

PG: Add Support for GroupBy's Cube, Rollup and Grouping Set

4 files changed, 50 insertions, 2 deletions

diff --git a/lib/arel/nodes/unary.rb b/lib/arel/nodes/unary.rb
index a0062ff5be..50946980b4 100644
--- a/lib/arel/nodes/unary.rb
+++ b/lib/arel/nodes/unary.rb
@@ -22,15 +22,19 @@ module Arel
 
     %w{
       Bin
+      Cube
+      DistinctOn
       Group
+      GroupingElement
+      GroupingSet
       Limit
+      Lock
       Not
       Offset
       On
       Ordering
+      RollUp
       Top
-      Lock
-      DistinctOn
     }.each do |name|
       const_set(name, Class.new(Unary))
     end
diff --git a/lib/arel/visitors/depth_first.rb b/lib/arel/visitors/depth_first.rb
index fb21fb1e70..d38795e640 100644
--- a/lib/arel/visitors/depth_first.rb
+++ b/lib/arel/visitors/depth_first.rb
@@ -18,6 +18,10 @@ module Arel
       end
       alias :visit_Arel_Nodes_Else              :unary
       alias :visit_Arel_Nodes_Group             :unary
+      alias :visit_Arel_Nodes_Cube              :unary
+      alias :visit_Arel_Nodes_RollUp            :unary
+      alias :visit_Arel_Nodes_GroupingSet       :unary
+      alias :visit_Arel_Nodes_GroupingElement   :unary
       alias :visit_Arel_Nodes_Grouping          :unary
       alias :visit_Arel_Nodes_Having            :unary
       alias :visit_Arel_Nodes_Limit             :unary
diff --git a/lib/arel/visitors/dot.rb b/lib/arel/visitors/dot.rb
index da75423523..95da652278 100644
--- a/lib/arel/visitors/dot.rb
+++ b/lib/arel/visitors/dot.rb
@@ -69,6 +69,10 @@ module Arel
         visit_edge o, "expr"
       end
       alias :visit_Arel_Nodes_Group             :unary
+      alias :visit_Arel_Nodes_Cube              :unary
+      alias :visit_Arel_Nodes_RollUp            :unary
+      alias :visit_Arel_Nodes_GroupingSet       :unary
+      alias :visit_Arel_Nodes_GroupingElement   :unary
       alias :visit_Arel_Nodes_Grouping          :unary
       alias :visit_Arel_Nodes_Having            :unary
       alias :visit_Arel_Nodes_Limit             :unary
diff --git a/lib/arel/visitors/postgresql.rb b/lib/arel/visitors/postgresql.rb
index 1ef0261bdd..ef0f0ea2ef 100644
--- a/lib/arel/visitors/postgresql.rb
+++ b/lib/arel/visitors/postgresql.rb
@@ -1,6 +1,10 @@
 module Arel
   module Visitors
     class PostgreSQL < Arel::Visitors::ToSql
+      CUBE = 'CUBE'
+      ROLLUP = 'ROLLUP'
+      GROUPING_SET = 'GROUPING SET'
+
       private
 
       def visit_Arel_Nodes_Matches o, collector
@@ -43,6 +47,38 @@ module Arel
       def visit_Arel_Nodes_BindParam o, collector
         collector.add_bind(o) { |i| "$#{i}" }
       end
+
+      def visit_Arel_Nodes_GroupingElement o, collector
+        collector << "( "
+        visit(o.expr, collector) << " )"
+      end
+
+      def visit_Arel_Nodes_Cube o, collector
+        collector << CUBE
+        grouping_array_or_grouping_element o, collector
+      end
+
+      def visit_Arel_Nodes_RollUp o, collector
+        collector << ROLLUP
+        grouping_array_or_grouping_element o, collector
+      end
+
+      def visit_Arel_Nodes_GroupingSet o, collector
+        collector << GROUPING_SET
+        grouping_array_or_grouping_element o, collector
+      end
+
+      # Utilized by GroupingSet, Cube & RollUp visitors to
+      # handle grouping aggregation semantics
+      def grouping_array_or_grouping_element o, collector
+        if o.expr.is_a? Array
+          collector << "( "
+          visit o.expr, collector
+          collector << " )"
+        else
+          visit o.expr, collector
+        end
+      end
     end
   end
 end