Add support for WITH and UNION

PostgreSQL WITH RECURSIVE support

Make WITH be a unary node

8 files changed, 149 insertions, 6 deletions

diff --git a/lib/arel/nodes.rb b/lib/arel/nodes.rb
index 64f2cc13fd..cbd10c31e0 100644
--- a/lib/arel/nodes.rb
+++ b/lib/arel/nodes.rb
@@ -9,6 +9,7 @@ require 'arel/nodes/update_statement'
 # unary
 require 'arel/nodes/unary'
 require 'arel/nodes/unqualified_column'
+require 'arel/nodes/with'
 
 # binary
 require 'arel/nodes/binary'
diff --git a/lib/arel/nodes/binary.rb b/lib/arel/nodes/binary.rb
index 0d02554199..e72228f11c 100644
--- a/lib/arel/nodes/binary.rb
+++ b/lib/arel/nodes/binary.rb
@@ -29,6 +29,8 @@ module Arel
       NotEqual
       NotIn
       Or
+      Union
+      UnionAll
     }.each do |name|
       const_set name, Class.new(Binary)
     end
diff --git a/lib/arel/nodes/select_statement.rb b/lib/arel/nodes/select_statement.rb
index c9a0cde4e0..cd74dd8e28 100644
--- a/lib/arel/nodes/select_statement.rb
+++ b/lib/arel/nodes/select_statement.rb
@@ -2,15 +2,17 @@ module Arel
   module Nodes
     class SelectStatement < Arel::Nodes::Node
       attr_reader :cores
-      attr_accessor :limit, :orders, :lock, :offset
+      attr_accessor :limit, :orders, :lock, :offset, :with, :with_recursive
 
       def initialize cores = [SelectCore.new]
         #puts caller
-        @cores  = cores
-        @orders = []
-        @limit  = nil
-        @lock   = nil
-        @offset = nil
+        @cores          = cores
+        @orders         = []
+        @limit          = nil
+        @lock           = nil
+        @offset         = nil
+        @with           = nil
+        @with_recursive = nil
       end
 
       def initialize_copy other
diff --git a/lib/arel/nodes/with.rb b/lib/arel/nodes/with.rb
new file mode 100644
index 0000000000..ce1644b60e
--- /dev/null
+++ b/lib/arel/nodes/with.rb
@@ -0,0 +1,17 @@
+module Arel
+	module Nodes
+		class With < Arel::Nodes::Unary
+			attr_reader :children
+			alias value children
+			alias expr  children
+
+			def initialize *children
+				@children = children
+			end
+
+		end
+
+		class WithRecursive < With; end
+	end
+end
+
diff --git a/lib/arel/select_manager.rb b/lib/arel/select_manager.rb
index b4103144fe..4f510d096b 100644
--- a/lib/arel/select_manager.rb
+++ b/lib/arel/select_manager.rb
@@ -134,6 +134,35 @@ module Arel
       Nodes::SqlLiteral.new viz.accept @ctx
     end
 
+    def union operation, other = nil
+    	if operation.is_a? Symbol
+    		if operation === :all
+					node_class = Nodes::UnionAll
+				else
+					raise "Only supported UNION operation is :all"
+				end
+			else
+				other = operation
+				node_class = Nodes::Union
+			end
+
+    	node_class.new self.ast, other.ast
+		end
+
+    def with *subqueries
+    	if subqueries.first.is_a? Symbol
+				if subqueries.shift == :recursive
+					node_class = Nodes::WithRecursive
+				else
+					raise "Only supported WITH modifier is :recursive"
+				end
+			else
+				node_class = Nodes::With
+			end
+			@ast.with = node_class.new(*subqueries)
+		end
+
+
     def take limit
       @ast.limit = Nodes::Limit.new(limit)
       @ctx.top   = Nodes::Top.new(limit)
diff --git a/lib/arel/visitors/to_sql.rb b/lib/arel/visitors/to_sql.rb
index d5534384f8..f3c8710068 100644
--- a/lib/arel/visitors/to_sql.rb
+++ b/lib/arel/visitors/to_sql.rb
@@ -129,6 +129,8 @@ eowarn
 
       def visit_Arel_Nodes_SelectStatement o
         [
+        	(visit(o.with) if o.with),
+        	(visit(o.with_recursive) if o.with_recursive),
           o.cores.map { |x| visit_Arel_Nodes_SelectCore x }.join,
           ("ORDER BY #{o.orders.map { |x| visit x }.join(', ')}" unless o.orders.empty?),
           (visit(o.limit) if o.limit),
@@ -149,6 +151,22 @@ eowarn
         ].compact.join ' '
       end
 
+      def visit_Arel_Nodes_With o
+				"WITH #{visit o.children}"
+			end
+
+			def visit_Arel_Nodes_WithRecursive o
+				"WITH RECURSIVE #{visit o.children}"
+			end
+
+			def visit_Arel_Nodes_Union o
+				"( #{visit o.left} UNION #{visit o.right} )"
+			end
+
+			def visit_Arel_Nodes_UnionAll o
+				"( #{visit o.left} UNION ALL #{visit o.right} )"
+			end
+
       def visit_Arel_Nodes_Having o
         "HAVING #{visit o.expr}"
       end
diff --git a/test/test_select_manager.rb b/test/test_select_manager.rb
index 271a8ae0b0..f2873b2a65 100644
--- a/test/test_select_manager.rb
+++ b/test/test_select_manager.rb
@@ -178,6 +178,79 @@ module Arel
       end
     end
 
+    describe 'union' do
+    	before do
+				table = Table.new :users
+				@m1 = Arel::SelectManager.new Table.engine, table
+				@m1.project Arel.star
+				@m1.where(table[:age].lt(18))
+
+				@m2 = Arel::SelectManager.new Table.engine, table
+				@m2.project Arel.star
+				@m2.where(table[:age].gt(99))
+
+
+			end
+
+    	it 'should union two managers' do
+				# FIXME should this union "managers" or "statements" ?
+				# FIXME this probably shouldn't return a node
+				node = @m1.union @m2
+
+				# maybe FIXME: decide when wrapper parens are needed
+				node.to_sql.must_be_like %{
+					( SELECT * FROM "users"  WHERE "users"."age" < 18 UNION SELECT * FROM "users"  WHERE "users"."age" > 99 )
+				}
+			end
+
+			it 'should union all' do
+				node = @m1.union :all, @m2
+
+				node.to_sql.must_be_like %{
+					( SELECT * FROM "users"  WHERE "users"."age" < 18 UNION ALL SELECT * FROM "users"  WHERE "users"."age" > 99 )
+				}
+			end
+
+		end
+
+		describe 'with' do
+
+      it "should support WITH RECURSIVE" do
+        comments           = Table.new(:comments)
+        comments_id        = comments[:id]
+        comments_parent_id = comments[:parent_id]
+
+        replies            = Table.new(:replies)
+        replies_id         = replies[:id]
+
+        recursive_term = Arel::SelectManager.new Table.engine
+        recursive_term.from(comments).project(comments_id, comments_parent_id).where(comments_id.eq 42)
+
+        non_recursive_term = Arel::SelectManager.new Table.engine
+        non_recursive_term.from(comments).project(comments_id, comments_parent_id).join(replies).on(comments_parent_id.eq replies_id)
+
+        union = recursive_term.union(non_recursive_term)
+
+        as_statement = Arel::Nodes::As.new replies, union
+
+        manager = Arel::SelectManager.new Table.engine
+        manager.from replies
+        manager.with :recursive, as_statement
+        manager.project Arel.star
+
+        sql = manager.to_sql
+        sql.must_be_like %{
+          WITH RECURSIVE "replies" AS (
+              SELECT "comments"."id", "comments"."parent_id" FROM "comments" WHERE "comments"."id" = 42
+            UNION
+              SELECT "comments"."id", "comments"."parent_id" FROM "comments" INNER JOIN "replies" ON "comments"."parent_id" = "replies"."id"
+          )
+          SELECT * FROM "replies"
+        }
+      end
+
+		end
+
     describe 'ast' do
       it 'should return the ast' do
         table   = Table.new :users
diff --git a/test/visitors/test_postgres.rb b/test/visitors/test_postgres.rb
index 6e8f399b6b..8d3f19aa6e 100644
--- a/test/visitors/test_postgres.rb
+++ b/test/visitors/test_postgres.rb
@@ -22,6 +22,7 @@ module Arel
         assert_match(/LIMIT 'omg'/, sql)
         assert_equal 1, sql.scan(/LIMIT/).length, 'should have one limit'
       end
+
     end
   end
 end