sequel 4.41.0 → 4.42.0

data/spec/extensions/skip_create_refresh_spec.rb

@@ -4,11 +4,12 @@ describe "Sequel::Plugins::SkipCreateRefresh" do
   it "should skip the refresh after saving a new object" do
     c = Class.new(Sequel::Model(:a))
     c.columns :id, :x
+    c.dataset = c.dataset.with_autoid(2)
     c.db.reset
-    c.instance_dataset.meta_def(:insert){|*a| super(*a); 2}
     c.create(:x=>1)
     c.db.sqls.must_equal ['INSERT INTO a (x) VALUES (1)', 'SELECT * FROM a WHERE (id = 2) LIMIT 1']
 
+    c.dataset = c.dataset.with_autoid(2)
     c.plugin :skip_create_refresh
     c.db.reset
     c.create(:x=>3).values.must_equal(:id=>2, :x=>3)

data/spec/extensions/spec_helper.rb

@@ -45,7 +45,7 @@ class << Sequel::Model
   def columns(*cols)
     return super if cols.empty?
     define_method(:columns){cols}
-    @dataset.instance_variable_set(:@columns, cols) if @dataset
+    @dataset.send(:columns=, cols) if @dataset
     def_column_accessor(*cols)
     @columns = cols
     @db_schema = {}

data/spec/extensions/split_values_spec.rb

@@ -1,11 +1,11 @@
 require File.join(File.dirname(File.expand_path(__FILE__)), "spec_helper")
 
 describe "Sequel::Plugins::SplitValues" do
-  it "should skip the refresh after saving a new object" do
+  it "stores non-columns in a separate hash" do
     c = Class.new(Sequel::Model(:a))
     c.columns :id, :x
     c.plugin :split_values
-    c.dataset._fetch = {:id=>1, :x=>2, :y=>3}
+    c.dataset = c.dataset.with_fetch(:id=>1, :x=>2, :y=>3)
     o = c.first
     c.db.reset
 

data/spec/extensions/sql_comments_spec.rb

@@ -24,4 +24,10 @@ describe "sql_comments extension" do
     ds = @ds.comment("Some\nComment\r\n Here")
     ds.where(:id=>ds).select_sql.must_equal "SELECT * FROM t WHERE (id IN (SELECT * FROM t -- Some Comment Here\n)) -- Some Comment Here\n"
   end
+
+  it "should handle frozen SQL strings" do
+    @ds = Sequel.mock[:t].with_extend{def select_sql; super.freeze; end}.extension(:sql_comments)
+    ds = @ds.comment("Some\nComment\r\n Here")
+    ds.select_sql.must_equal "SELECT * FROM t -- Some Comment Here\n"
+  end
 end

data/spec/extensions/static_cache_spec.rb

@@ -25,10 +25,10 @@ describe "Sequel::Plugins::StaticCache" do
     it "should make .[] method with primary key use the cache" do
       @c[1].must_equal @c1
       @c[2].must_equal @c2
-      @c[3].must_equal nil
-      @c[[1, 2]].must_equal nil
-      @c[nil].must_equal nil
-      @c[].must_equal nil
+      @c[3].must_be_nil
+      @c[[1, 2]].must_be_nil
+      @c[nil].must_be_nil
+      @c[].must_be_nil
       @db.sqls.must_equal []
     end
 
@@ -41,7 +41,7 @@ describe "Sequel::Plugins::StaticCache" do
     it "should support cache_get_pk" do
       @c.cache_get_pk(1).must_equal @c1
       @c.cache_get_pk(2).must_equal @c2
-      @c.cache_get_pk(3).must_equal nil
+      @c.cache_get_pk(3).must_be_nil
       @db.sqls.must_equal []
     end
 
@@ -171,9 +171,7 @@ describe "Sequel::Plugins::StaticCache" do
     end
 
     it "set_dataset should work correctly" do
-      ds = @c.dataset.from(:t2)
-      ds.instance_variable_set(:@columns, [:id])
-      ds._fetch = {:id=>3}
+      ds = @c.dataset.from(:t2).columns(:id).with_fetch(:id=>3)
       @c.dataset = ds
       @c.all.must_equal [@c.load(:id=>3)]
       @c.to_hash.must_equal(3=>@c.load(:id=>3))
@@ -341,7 +339,7 @@ describe "Sequel::Plugins::StaticCache" do
       @db.autoid = 3
       @db.fetch = [[{:id=>1}, {:id=>2}, {:id=>3}], [{:id=>3}]]
       o.save
-      @c[3].must_equal nil
+      @c[3].must_be_nil
     end
 
     it "should not automatically update the cache when updating model objects" do

data/spec/extensions/string_agg_spec.rb

@@ -30,39 +30,40 @@ describe "string_agg extension" do
   end
 
   it "should correctly literalize on Postgres" do
-    db = dbf.call(:postgres)
-    db.literal(@sa1).must_equal "string_agg(c, ',')"
-    db.literal(@sa2).must_equal "string_agg(c, '-')"
-    db.literal(@sa3).must_equal "string_agg(c, '-' ORDER BY o)"
-    db.literal(@sa4).must_equal "string_agg(DISTINCT c, ',' ORDER BY o)"
+    ds = dbf.call(:postgres).dataset.with_quote_identifiers(false)
+    ds.literal(@sa1).must_equal "string_agg(c, ',')"
+    ds.literal(@sa2).must_equal "string_agg(c, '-')"
+    ds.literal(@sa3).must_equal "string_agg(c, '-' ORDER BY o)"
+    ds.literal(@sa4).must_equal "string_agg(DISTINCT c, ',' ORDER BY o)"
   end
 
   it "should correctly literalize on SQLAnywhere" do
-    db = dbf.call(:sqlanywhere)
-    db.literal(@sa1).must_equal "list(c, ',')"
-    db.literal(@sa2).must_equal "list(c, '-')"
-    db.literal(@sa3).must_equal "list(c, '-' ORDER BY o)"
-    db.literal(@sa4).must_equal "list(DISTINCT c, ',' ORDER BY o)"
+    ds = dbf.call(:sqlanywhere).dataset.with_quote_identifiers(false).with_extend{def input_identifier(v) v.to_s end}
+    ds.literal(@sa1).must_equal "list(c, ',')"
+    ds.literal(@sa2).must_equal "list(c, '-')"
+    ds.literal(@sa3).must_equal "list(c, '-' ORDER BY o)"
+    ds.literal(@sa4).must_equal "list(DISTINCT c, ',' ORDER BY o)"
   end
 
   it "should correctly literalize on MySQL, H2, HSQLDB, CUBRID" do
     [:mysql, :h2, :hsqldb, :cubrid].each do |type|
       db = dbf.call(type)
       db.meta_def(:database_type){type}
-      db.literal(@sa1).must_equal "GROUP_CONCAT(c SEPARATOR ',')"
-      db.literal(@sa2).must_equal "GROUP_CONCAT(c SEPARATOR '-')"
-      db.literal(@sa3).must_equal "GROUP_CONCAT(c ORDER BY o SEPARATOR '-')"
-      db.literal(@sa4).must_equal "GROUP_CONCAT(DISTINCT c ORDER BY o SEPARATOR ',')"
+      ds = db.dataset.with_quote_identifiers(false).with_extend{def input_identifier(v) v.to_s end}
+      ds.literal(@sa1).upcase.must_equal "GROUP_CONCAT(C SEPARATOR ',')"
+      ds.literal(@sa2).upcase.must_equal "GROUP_CONCAT(C SEPARATOR '-')"
+      ds.literal(@sa3).upcase.must_equal "GROUP_CONCAT(C ORDER BY O SEPARATOR '-')"
+      ds.literal(@sa4).upcase.must_equal "GROUP_CONCAT(DISTINCT C ORDER BY O SEPARATOR ',')"
     end
   end
 
   it "should correctly literalize on Oracle and DB2" do
     [:oracle, :db2].each do |type|
-      db = dbf.call(type)
-      db.literal(@sa1).must_equal "listagg(c, ',') WITHIN GROUP (ORDER BY 1)"
-      db.literal(@sa2).must_equal "listagg(c, '-') WITHIN GROUP (ORDER BY 1)"
-      db.literal(@sa3).must_equal "listagg(c, '-') WITHIN GROUP (ORDER BY o)"
-      proc{db.literal(@sa4)}.must_raise Sequel::Error
+      ds = dbf.call(type).dataset.with_quote_identifiers(false).with_extend{def input_identifier(v) v.to_s end}
+      ds.literal(@sa1).must_equal "listagg(c, ',') WITHIN GROUP (ORDER BY 1)"
+      ds.literal(@sa2).must_equal "listagg(c, '-') WITHIN GROUP (ORDER BY 1)"
+      ds.literal(@sa3).must_equal "listagg(c, '-') WITHIN GROUP (ORDER BY o)"
+      proc{ds.literal(@sa4)}.must_raise Sequel::Error
     end
   end
 
@@ -72,18 +73,18 @@ describe "string_agg extension" do
 
   it "should handle order without arguments" do
     db = dbf.call(:postgres)
-    db.literal(@sa1.order).must_equal "string_agg(c, ',')"
+    db.dataset.with_quote_identifiers(false).literal(@sa1.order).must_equal "string_agg(c, ',')"
   end
 
   it "should handle operations on object" do
-    db = dbf.call(:postgres)
-    db.literal(@sa1 + 'b').must_equal "(string_agg(c, ',') || 'b')"
-    db.literal(@sa1.like('b')).must_equal "(string_agg(c, ',') LIKE 'b' ESCAPE '\\')"
-    db.literal(@sa1 < 'b').must_equal "(string_agg(c, ',') < 'b')"
-    db.literal(@sa1.as(:b)).must_equal "string_agg(c, ',') AS b"
-    db.literal(@sa1.cast(:b)).must_equal "CAST(string_agg(c, ',') AS b)"
-    db.literal(@sa1.desc).must_equal "string_agg(c, ',') DESC"
-    db.literal(@sa1 =~ /a/).must_equal "(string_agg(c, ',') ~ 'a')"
-    db.literal(@sa1.sql_subscript(1)).must_equal "string_agg(c, ',')[1]"
+    ds = dbf.call(:postgres).dataset.with_quote_identifiers(false)
+    ds.literal(@sa1 + 'b').must_equal "(string_agg(c, ',') || 'b')"
+    ds.literal(@sa1.like('b')).must_equal "(string_agg(c, ',') LIKE 'b' ESCAPE '\\')"
+    ds.literal(@sa1 < 'b').must_equal "(string_agg(c, ',') < 'b')"
+    ds.literal(@sa1.as(:b)).must_equal "string_agg(c, ',') AS b"
+    ds.literal(@sa1.cast(:b)).must_equal "CAST(string_agg(c, ',') AS b)"
+    ds.literal(@sa1.desc).must_equal "string_agg(c, ',') DESC"
+    ds.literal(@sa1 =~ /a/).must_equal "(string_agg(c, ',') ~ 'a')"
+    ds.literal(@sa1.sql_subscript(1)).must_equal "string_agg(c, ',')[1]"
   end
 end

data/spec/extensions/tactical_eager_loading_spec.rb

@@ -21,7 +21,7 @@ describe "Sequel::Plugins::TacticalEagerLoading" do
       columns :id, :parent_id
       many_to_one :parent, :class=>self
       one_to_many :children, :class=>self, :key=>:parent_id
-      dataset._fetch = proc do |sql|
+      set_dataset dataset.with_fetch(proc do |sql|
         if sql !~ /WHERE/
           [{:id=>1, :parent_id=>101}, {:id=>2, :parent_id=>102}, {:id=>101, :parent_id=>nil}, {:id=>102, :parent_id=>nil}]
         elsif sql =~ /WHERE.*\bid = (\d+)/
@@ -31,7 +31,7 @@ describe "Sequel::Plugins::TacticalEagerLoading" do
         elsif sql =~ /WHERE.*\bparent_id IN \(([\d, ]*)\)/
           $1.split(', ').map{|x| {:id=>x.to_i - 100, :parent_id=>x.to_i} if x.to_i > 100}.compact
         end
-      end
+      end)
     end
     @c = ::TacticalEagerLoadingModel
     @ds = TacticalEagerLoadingModel.dataset
@@ -63,7 +63,7 @@ describe "Sequel::Plugins::TacticalEagerLoading" do
   it "association getter methods should not eagerly load the association if the association is cached" do
     ts.map{|x| x.parent}.must_equal [ts[2], ts[3], nil, nil]
     sql_match(/\ASELECT \* FROM t WHERE \(t\.id IN \(10[12], 10[12]\)\)\z/)
-    def @ds.eager_load(*) raise end
+    @c.dataset = @c.dataset.with_extend{def eager_load(*) raise end}
     ts.map{|x| x.parent}.must_equal [ts[2], ts[3], nil, nil]
   end
 
@@ -113,8 +113,7 @@ describe "Sequel::Plugins::TacticalEagerLoading" do
   it "should handle case where an association is valid on an instance, but not on all instances" do
     c = Class.new(@c)
     c.many_to_one :parent2, :class=>@c, :key=>:parent_id
-    @c.dataset.row_proc = proc{|r| (r[:parent_id] == 101 ? c : @c).call(r)}
-    @c.all{|x| x.parent2 if x.is_a?(c)}
+    @c.dataset.with_row_proc(proc{|r| (r[:parent_id] == 101 ? c : @c).call(r)}).all{|x| x.parent2 if x.is_a?(c)}
     sql_match('SELECT * FROM t', 'SELECT * FROM t WHERE id = 101')
   end
 

data/spec/extensions/thread_local_timezones_spec.rb

@@ -22,7 +22,7 @@ describe "Sequel thread_local_timezones extension" do
     Sequel.thread_database_timezone = :utc
     Sequel.database_timezone.must_equal :utc
     Sequel.thread_typecast_timezone = nil
-    Sequel.typecast_timezone.must_equal nil
+    Sequel.typecast_timezone.must_be_nil
   end
   
   it "should fallback to default timezone if no thread_local timezone" do
@@ -37,7 +37,7 @@ describe "Sequel thread_local_timezones extension" do
     Sequel.thread_typecast_timezone = nil
     Sequel.typecast_timezone.must_equal :utc
     Sequel.thread_typecast_timezone = :nil
-    Sequel.typecast_timezone.must_equal nil
+    Sequel.typecast_timezone.must_be_nil
   end
     
   it "should be thread safe" do

data/spec/extensions/timestamps_spec.rb

@@ -16,7 +16,6 @@ describe "Sequel::Plugins::Timestamps" do
       def _save_refresh(*) end
       db.reset
     end
-    @c.dataset.autoid = nil
   end 
   after do
     Sequel.datetime_class = Time
@@ -71,7 +70,7 @@ describe "Sequel::Plugins::Timestamps" do
   end
 
   it "should not update the update timestamp on creation" do
-    @c.create.updated_at.must_equal nil
+    @c.create.updated_at.must_be_nil
   end
 
   it "should use the same value for the creation and update timestamps when creating if the :update_on_create option is given" do
@@ -133,6 +132,32 @@ describe "Sequel::Plugins::Timestamps" do
     o.created_at.must_equal '2009-08-01'
   end
 
+  it "should set update timestamp to same timestamp as create timestamp when setting creating timestamp" do
+    i = 1
+    (class << (Sequel.datetime_class); self end).send(:define_method, :now){"2009-08-0#{i+=1}"}
+    @c.plugin :timestamps, :update_on_create=>true
+    o = @c.create
+    sqls = @c.db.sqls
+    sqls.length.must_equal 1
+    ["INSERT INTO t (created_at, updated_at) VALUES ('2009-08-02', '2009-08-02')",
+     "INSERT INTO t (updated_at, created_at) VALUES ('2009-08-02', '2009-08-02')"].must_include sqls.first
+    o.created_at.must_equal '2009-08-02'
+    o.updated_at.must_equal '2009-08-02'
+  end
+
+  it "should set update timestamp when using not overriding create timestamp" do
+    i = 1
+    (class << (Sequel.datetime_class); self end).send(:define_method, :now){"2009-08-0#{i+=1}"}
+    @c.plugin :timestamps, :update_on_create=>true
+    o = @c.create(:created_at=>'2009-08-10')
+    sqls = @c.db.sqls
+    sqls.length.must_equal 1
+    ["INSERT INTO t (created_at, updated_at) VALUES ('2009-08-10', '2009-08-02')",
+     "INSERT INTO t (updated_at, created_at) VALUES ('2009-08-02', '2009-08-10')"].must_include sqls.first
+    o.created_at.must_equal '2009-08-10'
+    o.updated_at.must_equal '2009-08-02'
+  end
+
   it "should have create_timestamp_field give the create timestamp field" do
     @c.create_timestamp_field.must_equal :created_at
     @c.plugin :timestamps, :create=>:c
@@ -161,7 +186,7 @@ describe "Sequel::Plugins::Timestamps" do
     c = Class.new(@c)
     o = c.create
     o.created_at.must_equal '2009-08-01'
-    o.updated_at.must_equal nil
+    o.updated_at.must_be_nil
     o = c.load(:id=>1).save
     o.updated_at.must_equal '2009-08-01'
     c.db.sqls.must_equal ["INSERT INTO t (created_at) VALUES ('2009-08-01')", "UPDATE t SET updated_at = '2009-08-01' WHERE (id = 1)"]

data/spec/extensions/to_dot_spec.rb

@@ -25,8 +25,7 @@ END
   end
 
   it "should handle WITH" do
-    def @ds.supports_cte?(*) true end
-    a = dot(@ds.with(:a, @ds))
+    a = dot(@ds.with_extend{def supports_cte?(*) true end}.with(:a, @ds))
     a[0..3].must_equal ["1 -> 2 [label=\"with\"];", "2 [label=\"Array\"];", "2 -> 3 [label=\"0\"];", "3 [label=\"Hash\"];"]
     [["3 -> 4 [label=\"dataset\"];", "4 [label=\"Dataset\"];", "3 -> 5 [label=\"name\"];", "5 [label=\":a\"];"],
      ["3 -> 4 [label=\"name\"];", "4 [label=\":a\"];", "3 -> 5 [label=\"dataset\"];", "5 [label=\"Dataset\"];"]].must_include(a[4..-1])

data/spec/extensions/tree_spec.rb

@@ -48,7 +48,7 @@ describe Sequel::Model, "tree plugin" do
   end
 
   it "should have tree_order give the order of the association" do
-    @c.tree_order.must_equal nil
+    @c.tree_order.must_be_nil
     klass(:order=>:name).tree_order.must_equal :name
     klass(:order=>[:parent_id, :name]).tree_order.must_equal [:parent_id, :name]
   end
@@ -60,7 +60,7 @@ describe Sequel::Model, "tree plugin" do
   end
 
   it "should have roots return an array of the tree's roots" do
-    @ds._fetch = [{:id=>1, :parent_id=>nil, :name=>'r'}]
+    @c.dataset = @c.dataset.with_fetch([{:id=>1, :parent_id=>nil, :name=>'r'}])
     @c.roots.must_equal [@c.load(:id=>1, :parent_id=>nil, :name=>'r')]
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE (parent_id IS NULL)"]
   end
@@ -70,14 +70,14 @@ describe Sequel::Model, "tree plugin" do
   end
 
   it "should have ancestors return the ancestors of the current node" do
-    @ds._fetch = [[{:id=>1, :parent_id=>5, :name=>'r'}], [{:id=>5, :parent_id=>nil, :name=>'r2'}]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :parent_id=>5, :name=>'r'}], [{:id=>5, :parent_id=>nil, :name=>'r2'}]])
     @o.ancestors.must_equal [@c.load(:id=>1, :parent_id=>5, :name=>'r'), @c.load(:id=>5, :parent_id=>nil, :name=>'r2')]
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE id = 1",
       "SELECT * FROM nodes WHERE id = 5"]
   end
 
   it "should have descendants return the descendants of the current node" do
-    @ds._fetch = [[{:id=>3, :parent_id=>2, :name=>'r'}, {:id=>4, :parent_id=>2, :name=>'r2'}], [{:id=>5, :parent_id=>4, :name=>'r3'}], []]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>3, :parent_id=>2, :name=>'r'}, {:id=>4, :parent_id=>2, :name=>'r2'}], [{:id=>5, :parent_id=>4, :name=>'r3'}], []])
     @o.descendants.must_equal [@c.load(:id=>3, :parent_id=>2, :name=>'r'), @c.load(:id=>4, :parent_id=>2, :name=>'r2'), @c.load(:id=>5, :parent_id=>4, :name=>'r3')] 
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE (nodes.parent_id = 2)",
       "SELECT * FROM nodes WHERE (nodes.parent_id = 3)",
@@ -86,7 +86,7 @@ describe Sequel::Model, "tree plugin" do
   end
 
   it "should have root return the root of the current node" do
-    @ds._fetch = [[{:id=>1, :parent_id=>5, :name=>'r'}], [{:id=>5, :parent_id=>nil, :name=>'r2'}]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :parent_id=>5, :name=>'r'}], [{:id=>5, :parent_id=>nil, :name=>'r2'}]])
     @o.root.must_equal @c.load(:id=>5, :parent_id=>nil, :name=>'r2')
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE id = 1",
       "SELECT * FROM nodes WHERE id = 5"]
@@ -102,14 +102,14 @@ describe Sequel::Model, "tree plugin" do
   end
 
   it "should have self_and_siblings return the children of the current node's parent" do
-    @ds._fetch = [[{:id=>1, :parent_id=>3, :name=>'r'}], [{:id=>7, :parent_id=>1, :name=>'r2'}, @o.values.dup]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :parent_id=>3, :name=>'r'}], [{:id=>7, :parent_id=>1, :name=>'r2'}, @o.values.dup]])
     @o.self_and_siblings.must_equal [@c.load(:id=>7, :parent_id=>1, :name=>'r2'), @o] 
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE id = 1",
       "SELECT * FROM nodes WHERE (nodes.parent_id = 1)"]
   end
 
   it "should have siblings return the children of the current node's parent, except for the current node" do
-    @ds._fetch = [[{:id=>1, :parent_id=>3, :name=>'r'}], [{:id=>7, :parent_id=>1, :name=>'r2'}, @o.values.dup]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :parent_id=>3, :name=>'r'}], [{:id=>7, :parent_id=>1, :name=>'r2'}, @o.values.dup]])
     @o.siblings.must_equal [@c.load(:id=>7, :parent_id=>1, :name=>'r2')] 
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE id = 1",
       "SELECT * FROM nodes WHERE (nodes.parent_id = 1)"]
@@ -121,23 +121,23 @@ describe Sequel::Model, "tree plugin" do
     end
 
     it "should have root class method return the root" do
-      @c.dataset._fetch = [{:id=>1, :parent_id=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :parent_id=>nil, :name=>'r'}])
       @c.root.must_equal @c.load(:id=>1, :parent_id=>nil, :name=>'r')
     end
 
     it "prevents creating a second root" do
-      @c.dataset._fetch = [{:id=>1, :parent_id=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :parent_id=>nil, :name=>'r'}])
       lambda { @c.create }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
     end
 
     it "errors when promoting an existing record to a second root" do
-      @c.dataset._fetch = [{:id=>1, :parent_id=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :parent_id=>nil, :name=>'r'}])
       n = @c.load(:id => 2, :parent_id => 1)
       lambda { n.update(:parent_id => nil) }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
     end
 
     it "allows updating existing root" do
-      @c.dataset._fetch = [{:id=>1, :parent_id=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :parent_id=>nil, :name=>'r'}])
        @c.root.update(:name => 'fdsa') 
     end
   end
@@ -152,6 +152,10 @@ describe Sequel::Model, "tree plugin with composite keys" do
       columns :id, :id2, :name, :parent_id, :parent_id2, :i, :pi
       set_primary_key [:id, :id2]
       plugin :tree, opts.merge(:key=>[:parent_id, :parent_id2])
+      def self.set_dataset(ds)
+        super
+        set_primary_key [:id, :id2]
+      end
     end
     c
   end
@@ -174,7 +178,7 @@ describe Sequel::Model, "tree plugin with composite keys" do
   end
 
   it "should have roots return an array of the tree's roots" do
-    @ds._fetch = [{:id=>1, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}]
+    @c.dataset = @c.dataset.with_fetch([{:id=>1, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}])
     @c.roots.must_equal [@c.load(:id=>1, :parent_id=>nil, :parent_id2=>nil, :name=>'r')]
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE ((parent_id IS NULL) OR (parent_id2 IS NULL))"]
   end
@@ -184,7 +188,7 @@ describe Sequel::Model, "tree plugin with composite keys" do
   end
 
   it "should have ancestors return the ancestors of the current node" do
-    @ds._fetch = [[{:id=>1, :id2=>6, :parent_id=>5, :parent_id2=>7, :name=>'r'}], [{:id=>5, :id2=>7, :parent_id=>nil, :parent_id2=>nil, :name=>'r2'}]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :id2=>6, :parent_id=>5, :parent_id2=>7, :name=>'r'}], [{:id=>5, :id2=>7, :parent_id=>nil, :parent_id2=>nil, :name=>'r2'}]])
     @o.ancestors.must_equal [@c.load(:id=>1, :id2=>6, :parent_id=>5, :parent_id2=>7, :name=>'r'), @c.load(:id=>5, :id2=>7, :parent_id=>nil, :parent_id2=>nil, :name=>'r2')]
     sqls = @db.sqls
     sqls.length.must_equal 2
@@ -193,7 +197,7 @@ describe Sequel::Model, "tree plugin with composite keys" do
   end
 
   it "should have descendants return the descendants of the current node" do
-    @ds._fetch = [[{:id=>3, :id2=>7, :parent_id=>2, :parent_id2=>5, :name=>'r'}, {:id=>4, :id2=>8, :parent_id=>2, :parent_id2=>5, :name=>'r2'}], [{:id=>5, :id2=>9, :parent_id=>4, :parent_id2=>8, :name=>'r3'}], []]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>3, :id2=>7, :parent_id=>2, :parent_id2=>5, :name=>'r'}, {:id=>4, :id2=>8, :parent_id=>2, :parent_id2=>5, :name=>'r2'}], [{:id=>5, :id2=>9, :parent_id=>4, :parent_id2=>8, :name=>'r3'}], []])
     @o.descendants.must_equal [@c.load(:id=>3, :id2=>7, :parent_id=>2, :parent_id2=>5, :name=>'r'), @c.load(:id=>4, :id2=>8, :parent_id=>2, :parent_id2=>5, :name=>'r2'), @c.load(:id=>5, :id2=>9, :parent_id=>4, :parent_id2=>8, :name=>'r3')] 
     @db.sqls.must_equal ["SELECT * FROM nodes WHERE ((nodes.parent_id = 2) AND (nodes.parent_id2 = 5))",
       "SELECT * FROM nodes WHERE ((nodes.parent_id = 3) AND (nodes.parent_id2 = 7))",
@@ -202,7 +206,7 @@ describe Sequel::Model, "tree plugin with composite keys" do
   end
 
   it "should have root return the root of the current node" do
-    @ds._fetch = [[{:id=>1, :id2=>6, :parent_id=>5, :parent_id2=>7, :name=>'r'}], [{:id=>5, :id2=>7, :parent_id=>nil, :parent_id2=>nil, :name=>'r2'}]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :id2=>6, :parent_id=>5, :parent_id2=>7, :name=>'r'}], [{:id=>5, :id2=>7, :parent_id=>nil, :parent_id2=>nil, :name=>'r2'}]])
     @o.root.must_equal @c.load(:id=>5, :id2=>7, :parent_id=>nil, :parent_id2=>nil, :name=>'r2')
     sqls = @db.sqls
     sqls.length.must_equal 2
@@ -222,7 +226,7 @@ describe Sequel::Model, "tree plugin with composite keys" do
   end
 
   it "should have self_and_siblings return the children of the current node's parent" do
-    @ds._fetch = [[{:id=>1, :id2=>6, :parent_id=>3, :parent_id2=>7, :name=>'r'}], [{:id=>7, :id2=>9, :parent_id=>1, :parent_id2=>6, :name=>'r2'}, @o.values.dup]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :id2=>6, :parent_id=>3, :parent_id2=>7, :name=>'r'}], [{:id=>7, :id2=>9, :parent_id=>1, :parent_id2=>6, :name=>'r2'}, @o.values.dup]])
     @o.self_and_siblings.must_equal [@c.load(:id=>7, :id2=>9, :parent_id=>1, :parent_id2=>6, :name=>'r2'), @o] 
     sqls = @db.sqls
     sqls.length.must_equal 2
@@ -231,7 +235,7 @@ describe Sequel::Model, "tree plugin with composite keys" do
   end
 
   it "should have siblings return the children of the current node's parent, except for the current node" do
-    @ds._fetch = [[{:id=>1, :id2=>6, :parent_id=>3, :parent_id2=>7, :name=>'r'}], [{:id=>7, :id2=>9, :parent_id=>1, :parent_id2=>6, :name=>'r2'}, @o.values.dup]]
+    @c.dataset = @c.dataset.with_fetch([[{:id=>1, :id2=>6, :parent_id=>3, :parent_id2=>7, :name=>'r'}], [{:id=>7, :id2=>9, :parent_id=>1, :parent_id2=>6, :name=>'r2'}, @o.values.dup]])
     @o.siblings.must_equal [@c.load(:id=>7, :id2=>9, :parent_id=>1, :parent_id2=>6, :name=>'r2')] 
     sqls = @db.sqls
     sqls.length.must_equal 2
@@ -245,30 +249,30 @@ describe Sequel::Model, "tree plugin with composite keys" do
     end
 
     it "prevents creating a second root" do
-      @c.dataset._fetch = [{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}])
       lambda { @c.create }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
-      @c.dataset._fetch = [{:id=>1, :id2=>6, :parent_id=>1, :parent_id2=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :id2=>6, :parent_id=>1, :parent_id2=>nil, :name=>'r'}])
       lambda { @c.create(:parent_id2=>1) }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
-      @c.dataset._fetch = [{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>2, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>2, :name=>'r'}])
       lambda { @c.create(:parent_id=>2) }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
     end
 
     it "errors when promoting an existing record to a second root" do
-      @c.dataset._fetch = [{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}])
       lambda { @c.load(:id => 2, :id2=>7, :parent_id => 1, :parent_id2=>2).update(:parent_id => nil, :parent_id2=>nil) }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
-      @c.dataset._fetch = [{:id=>1, :id2=>6, :parent_id=>1, :parent_id2=>nil, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :id2=>6, :parent_id=>1, :parent_id2=>nil, :name=>'r'}])
       lambda { @c.load(:id => 2, :id2=>7, :parent_id => 1, :parent_id2=>2).update(:parent_id => nil) }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
-      @c.dataset._fetch = [{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>2, :name=>'r'}]
+      @c.dataset = @c.dataset.with_fetch([{:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>2, :name=>'r'}])
       lambda { @c.load(:id => 2, :id2=>7, :parent_id => 1, :parent_id2=>2).update(:parent_id2 => nil) }.must_raise(Sequel::Plugins::Tree::TreeMultipleRootError)
     end
 
     it "allows updating existing root" do
-      @c.dataset._fetch = {:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>nil, :name=>'r'}
-       @c.root.update(:name => 'fdsa') 
-      @c.dataset._fetch = {:id=>1, :id2=>6, :parent_id=>1, :parent_id2=>nil, :name=>'r'}
-       @c.root.update(:name => 'fdsa') 
-      @c.dataset._fetch = {:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>2, :name=>'r'}
-       @c.root.update(:name => 'fdsa') 
+      @c.dataset = @c.dataset.with_fetch(:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>nil, :name=>'r')
+      @c.root.update(:name => 'fdsa') 
+      @c.dataset = @c.dataset.with_fetch(:id=>1, :id2=>6, :parent_id=>1, :parent_id2=>nil, :name=>'r')
+      @c.root.update(:name => 'fdsa') 
+      @c.dataset = @c.dataset.with_fetch(:id=>1, :id2=>6, :parent_id=>nil, :parent_id2=>2, :name=>'r')
+      @c.root.update(:name => 'fdsa') 
     end
   end
 end