Update rcte plugin to allow the user to specify union_all in rcte

CHANGELOG

@@ -1,5 +1,7 @@
 === master
 
+* Support generating rcte queries using UNION or UNION ALL in the rcte plugin (jonathanfrias)
+
 * Make Database#table_exists? on PostgreSQL handle lock or statement timeout errors as evidence the table exists (jeremyevans) (#2106)
 
 * Work around DateTime.jd fractional second bug on JRuby in named_timezones extension (jeremyevans)

lib/sequel/plugins/rcte_tree.rb

@@ -71,6 +71,8 @@ module Plugins
     #              (default: :t)
     # :level_alias :: The symbol identifier to use when eagerly loading descendants
     #                 up to a given level (default: :x_level_x)
+    # :union_all :: Whether to use UNION ALL or UNION with the recursive
+    #               common table expression (default: true)
     module RcteTree
       # Create the appropriate parent, children, ancestors, and descendants
       # associations for the model.
@@ -80,6 +82,7 @@ def self.apply(model, opts=OPTS)
         opts = opts.dup
         opts[:class] = model
         opts[:methods_module] = Module.new
+        opts[:union_all] = opts[:union_all].nil? ? true : opts[:union_all]
         model.send(:include, opts[:methods_module])
 
         key = opts[:key] ||= :parent_id
@@ -142,7 +145,7 @@ def self.apply(model, opts=OPTS)
           model.from(SQL::AliasedExpression.new(t, table_alias)).
            with_recursive(t, col_aliases ? base_ds.select(*col_aliases) : base_ds.select_all,
             recursive_ds.select(*c_all),
-            :args=>col_aliases)
+            :args=>col_aliases, union_all: opts[:union_all])
         end
         aal = Array(a[:after_load])
         aal << proc do |m, ancs|
@@ -191,7 +194,7 @@ def self.apply(model, opts=OPTS)
           table_alias = model.dataset.schema_and_table(model.table_name)[1].to_sym
           ds = model.from(SQL::AliasedExpression.new(t, table_alias)).
             with_recursive(t, base_case, recursive_case,
-             :args=>((key_aliases + col_aliases) if col_aliases))
+             :args=>((key_aliases + col_aliases) if col_aliases), union_all: opts[:union_all])
           ds = r.apply_eager_dataset_changes(ds)
           ds = ds.select_append(ka) unless ds.opts[:select] == nil
           model.eager_load_results(r, eo.merge(:loader=>false, :initialize_rows=>false, :dataset=>ds, :id_map=>nil)) do |obj|
@@ -240,7 +243,7 @@ def self.apply(model, opts=OPTS)
           model.from(SQL::AliasedExpression.new(t, table_alias)).
            with_recursive(t, col_aliases ? base_ds.select(*col_aliases) : base_ds.select_all,
             recursive_ds.select(*c_all),
-            :args=>col_aliases)
+            :args=>col_aliases, union_all: opts[:union_all])
           end
         dal = Array(d[:after_load])
         dal << proc do |m, descs|
@@ -299,7 +302,7 @@ def self.apply(model, opts=OPTS)
           table_alias = model.dataset.schema_and_table(model.table_name)[1].to_sym
           ds = model.from(SQL::AliasedExpression.new(t, table_alias)).
             with_recursive(t, base_case, recursive_case,
-              :args=>((key_aliases + col_aliases + (level ? [la] : [])) if col_aliases))
+              :args=>((key_aliases + col_aliases + (level ? [la] : [])) if col_aliases), union_all: opts[:union_all])
           ds = r.apply_eager_dataset_changes(ds)
           ds = ds.select_append(ka) unless ds.opts[:select] == nil
           model.eager_load_results(r, eo.merge(:loader=>false, :initialize_rows=>false, :dataset=>ds, :id_map=>nil, :associations=>OPTS)) do |obj|

spec/extensions/rcte_tree_spec.rb

@@ -54,7 +54,18 @@ def self.name; 'Node'; end
     @c.eager(:descendants).all
     @db.sqls.must_equal ["SELECT * FROM nodes", "WITH t(x_root_x, id, name, parent_id, i, pi) AS (SELECT parent_id AS x_root_x, nodes.id, nodes.name, nodes.parent_id, nodes.i, nodes.pi FROM nodes WHERE (parent_id IN (1)) UNION ALL SELECT t.x_root_x, nodes.id, nodes.name, nodes.parent_id, nodes.i, nodes.pi FROM nodes INNER JOIN t ON (t.id = nodes.parent_id)) SELECT * FROM t AS nodes"]
   end
-
+
+  it "should use the correct SQL for eager loading when recursive CTEs set union_all" do
+    @c.plugin :rcte_tree, :union_all => false
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :name=>'A', :parent_id=>3}]])
+    @c.eager(:ancestors).all
+    @db.sqls.must_equal ["SELECT * FROM nodes", "WITH t AS (SELECT id AS x_root_x, nodes.* FROM nodes WHERE (id IN (3)) UNION SELECT t.x_root_x, nodes.* FROM nodes INNER JOIN t ON (t.parent_id = nodes.id)) SELECT * FROM t AS nodes"]
+
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :name=>'A', :parent_id=>3}]])
+    @c.eager(:descendants).all
+    @db.sqls.must_equal ["SELECT * FROM nodes", "WITH t AS (SELECT parent_id AS x_root_x, nodes.* FROM nodes WHERE (parent_id IN (1)) UNION SELECT t.x_root_x, nodes.* FROM nodes INNER JOIN t ON (t.id = nodes.parent_id)) SELECT * FROM t AS nodes"]
+  end
+
   it "should use the correct SQL for lazy associations when giving options" do
     @c.plugin :rcte_tree, :primary_key=>:i, :key=>:pi, :cte_name=>:cte, :order=>:name, :ancestors=>{:name=>:as}, :children=>{:name=>:cs}, :descendants=>{:name=>:ds}, :parent=>{:name=>:p}
     @o.p_dataset.sql.must_equal 'SELECT * FROM nodes WHERE (nodes.i = 4) ORDER BY name LIMIT 1'
@@ -86,6 +97,12 @@ def self.name; 'Node'; end
     @o.ancestors_dataset.sql.must_equal 'WITH t AS (SELECT * FROM nodes WHERE ((id = 1) AND (i = 1)) UNION ALL SELECT nodes.* FROM nodes INNER JOIN t ON (t.parent_id = nodes.id) WHERE (i = 1)) SELECT * FROM t AS nodes WHERE (i = 1)'
     @o.descendants_dataset.sql.must_equal 'WITH t AS (SELECT * FROM nodes WHERE ((parent_id = 2) AND (i = 1)) UNION ALL SELECT nodes.* FROM nodes INNER JOIN t ON (t.id = nodes.parent_id) WHERE (i = 1)) SELECT * FROM t AS nodes WHERE (i = 1)'
   end
+
+  it "should use the correct SQL for UNION queries when using :union_all option" do
+    @c.plugin :rcte_tree, :union_all => false
+    @o.ancestors_dataset.sql.must_equal 'WITH t AS (SELECT * FROM nodes WHERE (id = 1) UNION SELECT nodes.* FROM nodes INNER JOIN t ON (t.parent_id = nodes.id)) SELECT * FROM t AS nodes'
+    @o.descendants_dataset.sql.must_equal 'WITH t AS (SELECT * FROM nodes WHERE (parent_id = 2) UNION SELECT nodes.* FROM nodes INNER JOIN t ON (t.id = nodes.parent_id)) SELECT * FROM t AS nodes'
+  end
 
   it "should add all parent associations when lazily loading ancestors" do
     @c.plugin :rcte_tree
@@ -171,6 +188,27 @@ def self.name; 'Node'; end
     os.map{|o| o.parent.parent.parent.parent if o.parent and o.parent.parent and o.parent.parent.parent}.must_equal [nil, nil, nil, nil]
     @db.sqls.must_equal []
   end
+
+  it "should eagerly load ancestors with a union_all configured" do
+    @c.plugin :rcte_tree, :union_all => false
+    @ds = @c.dataset = @c.dataset.with_fetch([[{:id=>2, :parent_id=>1, :name=>'AA'}, {:id=>6, :parent_id=>2, :name=>'C'}, {:id=>7, :parent_id=>1, :name=>'D'}, {:id=>9, :parent_id=>nil, :name=>'E'}],
+      [{:id=>2, :name=>'AA', :parent_id=>1, :x_root_x=>2},
+       {:id=>1, :name=>'00', :parent_id=>8, :x_root_x=>1}, {:id=>1, :name=>'00', :parent_id=>8, :x_root_x=>2},
+       {:id=>8, :name=>'?', :parent_id=>nil, :x_root_x=>2}, {:id=>8, :name=>'?', :parent_id=>nil, :x_root_x=>1}]])
+    os = @ds.eager(:ancestors).all
+    @db.sqls.must_equal ["SELECT * FROM nodes",
+      'WITH t AS (SELECT id AS x_root_x, nodes.* FROM nodes WHERE (id IN (1, 2)) UNION SELECT t.x_root_x, nodes.* FROM nodes INNER JOIN t ON (t.parent_id = nodes.id)) SELECT * FROM t AS nodes']
+    os.must_equal [@c.load(:id=>2, :parent_id=>1, :name=>'AA'), @c.load(:id=>6, :parent_id=>2, :name=>'C'), @c.load(:id=>7, :parent_id=>1, :name=>'D'), @c.load(:id=>9, :parent_id=>nil, :name=>'E')]
+    os.map{|o| o.ancestors}.must_equal [[@c.load(:id=>1, :name=>'00', :parent_id=>8), @c.load(:id=>8, :name=>'?', :parent_id=>nil)],
+      [@c.load(:id=>2, :name=>'AA', :parent_id=>1), @c.load(:id=>1, :name=>'00', :parent_id=>8), @c.load(:id=>8, :name=>'?', :parent_id=>nil)],
+      [@c.load(:id=>1, :name=>'00', :parent_id=>8), @c.load(:id=>8, :name=>'?', :parent_id=>nil)],
+      []]
+    os.map{|o| o.parent}.must_equal [@c.load(:id=>1, :name=>'00', :parent_id=>8), @c.load(:id=>2, :name=>'AA', :parent_id=>1), @c.load(:id=>1, :name=>'00', :parent_id=>8), nil]
+    os.map{|o| o.parent.parent if o.parent}.must_equal [@c.load(:id=>8, :name=>'?', :parent_id=>nil), @c.load(:id=>1, :name=>'00', :parent_id=>8), @c.load(:id=>8, :name=>'?', :parent_id=>nil), nil]
+    os.map{|o| o.parent.parent.parent if o.parent and o.parent.parent}.must_equal [nil, @c.load(:id=>8, :name=>'?', :parent_id=>nil), nil, nil]
+    os.map{|o| o.parent.parent.parent.parent if o.parent and o.parent.parent and o.parent.parent.parent}.must_equal [nil, nil, nil, nil]
+    @db.sqls.must_equal []
+  end
 
   it "should eagerly load ancestors on oracle when root column is a BigDecimal value" do
     def (@c.dataset.db).database_type; :oracle; end