sequel 3.21.0 → 3.41.0

data/spec/extensions/sharding_spec.rb

@@ -2,271 +2,197 @@ require File.join(File.dirname(File.expand_path(__FILE__)), "spec_helper")
 
 describe "sharding plugin" do
   before do
-    @db = Sequel::Model.db.clone
+    @db = Sequel.mock(:numrows=>1, :autoid=>proc{1}, :servers=>{:s1=>{}, :s2=>{}, :s3=>{}, :s4=>{}})
     @Artist = Class.new(Sequel::Model(@db[:artists]))
     @Artist.class_eval do
+      instance_dataset._fetch = dataset._fetch = {:id=>2, :name=>'YJM'}
       columns :id, :name
-
-      def self.y
-        {:id=>2, :name=>'YJM'}
-      end
+      plugin :sharding
     end
     @Album = Class.new(Sequel::Model(@db[:albums]))
     @Album.class_eval do
+      instance_dataset._fetch = dataset._fetch = {:id=>1, :name=>'RF', :artist_id=>2}
       columns :id, :artist_id, :name
-
-      def self.ds_ext(m=nil)
-        @ds_ext = m if m
-        @ds_ext
-      end
-
-      def self.y
-        {:id=>1, :name=>'RF', :artist_id=>2}
-      end
-      
-      private
-
-      def _join_table_dataset(opts)
-        ds = super
-        m = model
-        ds.meta_def(:model){m}
-        ds.extend model.ds_ext
-        ds
-      end
+      plugin :sharding
     end
     @Tag = Class.new(Sequel::Model(@db[:tags]))
     @Tag.class_eval do
+      instance_dataset._fetch = dataset._fetch = {:id=>3, :name=>'M'}
       columns :id, :name
-
-      def self.y
-        {:id=>3, :name=>'M'}
-      end
+      plugin :sharding
     end
-    models = [@Artist, @Album, @Tag]
     @Artist.one_to_many :albums, :class=>@Album, :key=>:artist_id
     @Album.many_to_one :artist, :class=>@Artist
     @Album.many_to_many :tags, :class=>@Tag, :left_key=>:album_id, :right_key=>:tag_id, :join_table=>:albums_tags
-    m = Module.new do
-      def actions
-        @actions ||= []
-      end
-    end
-    models.each do |model|
-      model.extend m
-      model.plugin :sharding
-      model.dataset.extend(ds_ext = Module.new do
-        def insert(h={})
-          model.actions << [:insert, h.dup, opts[:server]]
-          1
-        end
-        def delete
-          model.actions << [:delete,(literal(opts[:where]) if opts[:where]), opts[:server]]
-          1
-        end
-        def update(h={})
-          model.actions << [:update, h.dup, (literal(opts[:where]) if opts[:where]), opts[:server]]
-          1
-        end
-        def fetch_rows(sql)
-          model.actions << [:fetch, (literal(opts[:where] || opts[:join]) if opts[:where] || opts[:join]), opts[:server]]
-          yield(model.y)
-        end
-      end)
-      @Album.ds_ext(ds_ext)
-    end
-    def @db.actions; @actions ||= []; end
-    def @db.transaction(opts)
-      actions << [:transaction, opts[:server]]
-      super
-    end
+    @db.sqls
   end 
 
   specify "should allow you to instantiate a new object for a specified shard" do
     @Album.new_using_server(:s1, :name=>'RF').save
-    @Album.actions.should == [[:insert, {:name=>"RF"}, :s1], [:fetch, "(id = 1)", :s1]]
+    @db.sqls.should == ["INSERT INTO albums (name) VALUES ('RF') -- s1", "SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s1"]
     
-    @Album.actions.clear
     @Album.new_using_server(:s2){|o| o.name = 'MO'}.save
-    @Album.actions.should == [[:insert, {:name=>"MO"}, :s2], [:fetch, "(id = 1)", :s2]]
+    @db.sqls.should == ["INSERT INTO albums (name) VALUES ('MO') -- s2", "SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s2"]
   end 
 
   specify "should allow you to create and save a new object for a specified shard" do
     @Album.create_using_server(:s1, :name=>'RF')
-    @Album.actions.should == [[:insert, {:name=>"RF"}, :s1], [:fetch, "(id = 1)", :s1]]
-    
-    @Album.actions.clear
+    @db.sqls.should == ["INSERT INTO albums (name) VALUES ('RF') -- s1", "SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s1"]
+
     @Album.create_using_server(:s2){|o| o.name = 'MO'}
-    @Album.actions.should == [[:insert, {:name=>"MO"}, :s2], [:fetch, "(id = 1)", :s2]]
+    @db.sqls.should == ["INSERT INTO albums (name) VALUES ('MO') -- s2", "SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s2"]
   end 
 
   specify "should have objects retrieved from a specific shard update that shard" do
     @Album.server(:s1).first.update(:name=>'MO')
-    @Album.actions.should == [[:fetch, nil, :s1], [:update, {:name=>"MO"}, "(id = 1)", :s1]]
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "UPDATE albums SET name = 'MO' WHERE (id = 1) -- s1"]
   end 
 
   specify "should have objects retrieved from a specific shard delete from that shard" do
     @Album.server(:s1).first.delete
-    @Album.actions.should == [[:fetch, nil, :s1], [:delete, "(id = 1)", :s1]]
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "DELETE FROM albums WHERE (id = 1) -- s1"]
   end 
 
   specify "should have objects retrieved from a specific shard reload from that shard" do
     @Album.server(:s1).first.reload
-    @Album.actions.should == [[:fetch, nil, :s1], [:fetch, "(id = 1)", :s1]]
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s1"]
   end 
 
   specify "should use current dataset's shard when eager loading if eagerly loaded dataset doesn't have its own shard" do
     albums = @Album.server(:s1).eager(:artist).all
-    @Album.actions.should == [[:fetch, nil, :s1]]
-    @Artist.actions.should == [[:fetch, "(artists.id IN (2))", :s1]]
-    @Artist.actions.clear
-    albums.length == 1
+    @db.sqls.should == ["SELECT * FROM albums -- s1", "SELECT * FROM artists WHERE (artists.id IN (2)) -- s1"]
+    albums.length.should == 1
     albums.first.artist.save
-    @Artist.actions.should == [[:update, {:name=>"YJM"}, "(id = 2)", :s1]]
+    @db.sqls.should == ["UPDATE artists SET name = 'YJM' WHERE (id = 2) -- s1"]
   end 
 
   specify "should not use current dataset's shard when eager loading if eagerly loaded dataset has its own shard" do
-    @Artist.dataset.opts[:server] = :s2
+    @Artist.instance_dataset.opts[:server] = @Artist.dataset.opts[:server] = :s2
     albums = @Album.server(:s1).eager(:artist).all
-    @Album.actions.should == [[:fetch, nil, :s1]]
-    @Artist.actions.should == [[:fetch, "(artists.id IN (2))", :s2]]
-    @Artist.actions.clear
-    albums.length == 1
+    @db.sqls.should == ["SELECT * FROM albums -- s1", "SELECT * FROM artists WHERE (artists.id IN (2)) -- s2"]
+    albums.length.should == 1
     albums.first.artist.save
-    @Artist.actions.should == [[:update, {:name=>"YJM"}, "(id = 2)", :s2]]
+    @db.sqls.should == ["UPDATE artists SET name = 'YJM' WHERE (id = 2) -- s2"]
   end 
 
   specify "should use current dataset's shard when eager graphing if eagerly graphed dataset doesn't have its own shard" do
     ds = @Album.server(:s1).eager_graph(:artist)
-    def ds.fetch_rows(sql)
-      super(sql)
-      yield({:id=>1, :artist_id=>2, :name=>'RF', :artist_id_0=>2, :artist_name=>'YJM'})
-    end
+    ds._fetch = {:id=>1, :artist_id=>2, :name=>'RF', :artist_id_0=>2, :artist_name=>'YJM'}
     albums = ds.all
-    @Album.actions.should == [[:fetch, "( LEFT OUTER JOIN artists AS artist ON (artist.id = albums.artist_id))", :s1]]
-    albums.length == 1
+    @db.sqls.should == ["SELECT albums.id, albums.artist_id, albums.name, artist.id AS artist_id_0, artist.name AS artist_name FROM albums LEFT OUTER JOIN artists AS artist ON (artist.id = albums.artist_id) -- s1"]
+    albums.length.should == 1
     albums.first.artist.save
-    @Artist.actions.should == [[:update, {:name=>"YJM"}, "(id = 2)", :s1]]
+    @db.sqls.should == ["UPDATE artists SET name = 'YJM' WHERE (id = 2) -- s1"]
   end 
 
   specify "should not use current dataset's shard when eager graphing if eagerly graphed dataset has its own shard" do
-    @Artist.dataset.opts[:server] = :s2
+    @Artist.instance_dataset.opts[:server] = @Artist.dataset.opts[:server] = :s2
     ds = @Album.server(:s1).eager_graph(:artist)
-    def ds.fetch_rows(sql)
-      super(sql)
-      yield({:id=>1, :artist_id=>2, :name=>'RF', :artist_id_0=>2, :artist_name=>'YJM'})
-    end
+    ds._fetch = {:id=>1, :artist_id=>2, :name=>'RF', :artist_id_0=>2, :artist_name=>'YJM'}
     albums = ds.all
-    @Album.actions.should == [[:fetch, "( LEFT OUTER JOIN artists AS artist ON (artist.id = albums.artist_id))", :s1]]
-    albums.length == 1
+    @db.sqls.should == ["SELECT albums.id, albums.artist_id, albums.name, artist.id AS artist_id_0, artist.name AS artist_name FROM albums LEFT OUTER JOIN artists AS artist ON (artist.id = albums.artist_id) -- s1"]
+    albums.length.should == 1
     albums.first.artist.save
-    @Artist.actions.should == [[:update, {:name=>"YJM"}, "(id = 2)", :s2]]
+    @db.sqls.should == ["UPDATE artists SET name = 'YJM' WHERE (id = 2) -- s2"]
   end 
 
   specify "should use eagerly graphed dataset shard for eagerly graphed objects even if current dataset does not have a shard" do
-    @Artist.dataset.opts[:server] = :s2
+    @Artist.instance_dataset.opts[:server] = @Artist.dataset.opts[:server] = :s2
     ds = @Album.eager_graph(:artist)
-    def ds.fetch_rows(sql)
-      super(sql)
-      yield({:id=>1, :artist_id=>2, :name=>'RF', :artist_id_0=>2, :artist_name=>'YJM'})
-    end
+    ds._fetch = {:id=>1, :artist_id=>2, :name=>'RF', :artist_id_0=>2, :artist_name=>'YJM'}
     albums = ds.all
-    @Album.actions.should == [[:fetch, "( LEFT OUTER JOIN artists AS artist ON (artist.id = albums.artist_id))", nil]]
-    albums.length == 1
+    @db.sqls.should == ["SELECT albums.id, albums.artist_id, albums.name, artist.id AS artist_id_0, artist.name AS artist_name FROM albums LEFT OUTER JOIN artists AS artist ON (artist.id = albums.artist_id)"]
+    albums.length.should == 1
     albums.first.artist.save
-    @Artist.actions.should == [[:update, {:name=>"YJM"}, "(id = 2)", :s2]]
+    @db.sqls.should == ["UPDATE artists SET name = 'YJM' WHERE (id = 2) -- s2"]
   end 
 
   specify "should have objects retrieved from a specific shard use associated objects from that shard, with modifications to the associated objects using that shard" do
     album = @Album.server(:s1).first
-    @Album.actions.should == [[:fetch, nil, :s1]]
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1"]
     album.artist.update(:name=>'AS')
-    @Artist.actions.should == [[:fetch, "(artists.id = 2)", :s1], [:update, {:name=>"AS"}, "(id = 2)", :s1]]
+    @db.sqls.should == ["SELECT * FROM artists WHERE (artists.id = 2) LIMIT 1 -- s1", "UPDATE artists SET name = 'AS' WHERE (id = 2) -- s1"]
     album.tags.map{|a| a.update(:name=>'SR')}
-    @Tag.actions.should == [[:fetch, "( INNER JOIN albums_tags ON ((albums_tags.tag_id = tags.id) AND (albums_tags.album_id = 1)))", :s1], [:update, {:name=>"SR"}, "(id = 3)", :s1]]
-    
-    @Album.actions.clear
-    @Artist.actions.clear
+    @db.sqls.should == ["SELECT tags.* FROM tags INNER JOIN albums_tags ON ((albums_tags.tag_id = tags.id) AND (albums_tags.album_id = 1)) -- s1", "UPDATE tags SET name = 'SR' WHERE (id = 3) -- s1"]
     @Artist.server(:s2).first.albums.map{|a| a.update(:name=>'MO')}
-    @Artist.actions.should == [[:fetch, nil, :s2]]
-    @Album.actions.should == [[:fetch, "(albums.artist_id = 2)", :s2], [:update, {:name=>"MO"}, "(id = 1)", :s2]]
+    @db.sqls.should == ["SELECT * FROM artists LIMIT 1 -- s2", "SELECT * FROM albums WHERE (albums.artist_id = 2) -- s2", "UPDATE albums SET name = 'MO' WHERE (id = 1) -- s2"]
   end 
 
   specify "should have objects retrieved from a specific shard add associated objects to that shard" do
     album = @Album.server(:s1).first
     artist = @Artist.server(:s2).first
-    @Album.actions.clear
-    @Artist.actions.clear
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "SELECT * FROM artists LIMIT 1 -- s2"]
 
     artist.add_album(:name=>'MO')
-    @Album.actions.should == [[:insert, {:name=>"MO", :artist_id=>2}, :s2], [:fetch, "(id = 1)", :s2]]
-    @Album.actions.clear
+    sqls = @db.sqls
+    ["INSERT INTO albums (artist_id, name) VALUES (2, 'MO') -- s2", "INSERT INTO albums (name, artist_id) VALUES ('MO', 2) -- s2"].should include(sqls.shift)
+    sqls.should == ["SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s2"]
     
     album.add_tag(:name=>'SR')
-    @Tag.actions.should == [[:insert, {:name=>"SR"}, :s1], [:fetch, "(id = 1)", :s1]]
-    @Album.actions.should == [[:insert, {:album_id=>1, :tag_id=>3}, :s1]]
+    sqls = @db.sqls
+    ["INSERT INTO albums_tags (album_id, tag_id) VALUES (1, 3) -- s1", "INSERT INTO albums_tags (tag_id, album_id) VALUES (3, 1) -- s1"].should include(sqls.pop)
+    sqls.should == ["INSERT INTO tags (name) VALUES ('SR') -- s1", "SELECT * FROM tags WHERE (id = 1) LIMIT 1 -- s1", ]
   end 
 
   specify "should have objects retrieved from a specific shard remove associated objects from that shard" do
     album = @Album.server(:s1).first
     artist = @Artist.server(:s2).first
-    @Album.actions.clear
-    @Artist.actions.clear
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "SELECT * FROM artists LIMIT 1 -- s2"]
 
     artist.remove_album(1)
-    @Album.actions.should == [[:fetch, "((albums.artist_id = 2) AND (albums.id = 1))", :s2], [:update, {:name=>"RF", :artist_id=>nil}, "(id = 1)", :s2]]
-    @Album.actions.clear
+    sqls = @db.sqls
+    ["UPDATE albums SET artist_id = NULL, name = 'RF' WHERE (id = 1) -- s2", "UPDATE albums SET name = 'RF', artist_id = NULL WHERE (id = 1) -- s2"].should include(sqls.pop)
+    sqls.should == ["SELECT * FROM albums WHERE ((albums.artist_id = 2) AND (albums.id = 1)) LIMIT 1 -- s2"]
     
     album.remove_tag(3)
-    @Tag.actions.should == [[:fetch, "(tags.id = 3)", :s1]]
-    @Album.actions.should == [[:delete, "((album_id = 1) AND (tag_id = 3))", :s1]]
+    @db.sqls.should == ["SELECT tags.* FROM tags INNER JOIN albums_tags ON ((albums_tags.tag_id = tags.id) AND (albums_tags.album_id = 1)) WHERE (tags.id = 3) LIMIT 1 -- s1", "DELETE FROM albums_tags WHERE ((album_id = 1) AND (tag_id = 3)) -- s1"]
   end 
 
   specify "should have objects retrieved from a specific shard remove all associated objects from that shard" do
     album = @Album.server(:s1).first
     artist = @Artist.server(:s2).first
-    @Album.actions.clear
-    @Artist.actions.clear
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "SELECT * FROM artists LIMIT 1 -- s2"]
 
     artist.remove_all_albums
-    @Album.actions.should == [[:update, {:artist_id=>nil}, "(artist_id = 2)", :s2]]
-    @Album.actions.clear
+    @db.sqls.should == ["UPDATE albums SET artist_id = NULL WHERE (artist_id = 2) -- s2"]
     
     album.remove_all_tags
-    @Album.actions.should == [[:delete, "(album_id = 1)", :s1]]
+    @db.sqls.should == ["DELETE FROM albums_tags WHERE (album_id = 1) -- s1"]
   end 
 
   specify "should not override a server already set on an associated object" do
     album = @Album.server(:s1).first
     artist = @Artist.server(:s2).first
-    @Album.actions.clear
-    @Artist.actions.clear
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "SELECT * FROM artists LIMIT 1 -- s2"]
 
     artist.add_album(@Album.load(:id=>4, :name=>'MO').set_server(:s3))
-    @Album.actions.should == [[:update, {:name=>"MO", :artist_id=>2}, "(id = 4)", :s3]]
-    @Album.actions.clear
+    ["UPDATE albums SET artist_id = 2, name = 'MO' WHERE (id = 4) -- s3", "UPDATE albums SET name = 'MO', artist_id = 2 WHERE (id = 4) -- s3"].should include(@db.sqls.pop)
 
     artist.remove_album(@Album.load(:id=>5, :name=>'T', :artist_id=>2).set_server(:s4))
     # Should select from current object's shard to check existing association, but update associated object's shard
-    @Album.actions.should == [[:fetch, "((albums.artist_id = 2) AND (id = 5))", :s2], [:update, {:name=>"T", :artist_id=>nil}, "(id = 5)", :s4]]
-    @Album.actions.clear
+    sqls = @db.sqls
+    ["UPDATE albums SET artist_id = NULL, name = 'T' WHERE (id = 5) -- s4", "UPDATE albums SET name = 'T', artist_id = NULL WHERE (id = 5) -- s4"].should include(sqls.pop)
+    sqls.should == ["SELECT 1 AS one FROM albums WHERE ((albums.artist_id = 2) AND (id = 5)) LIMIT 1 -- s2"]
   end 
 
   specify "should be able to set a shard to use for any object using set_server" do
     @Album.server(:s1).first.set_server(:s2).reload
-    @Album.actions.should == [[:fetch, nil, :s1], [:fetch, "(id = 1)", :s2]]
+    @db.sqls.should == ["SELECT * FROM albums LIMIT 1 -- s1", "SELECT * FROM albums WHERE (id = 1) LIMIT 1 -- s2"]
   end 
 
   specify "should use transactions on the correct shard" do
     @Album.use_transactions = true
     @Album.server(:s2).first.save
-    @Album.actions.should == [[:fetch, nil, :s2], [:update, {:name=>"RF", :artist_id=>2}, "(id = 1)", :s2]]
-    @db.actions.should == [[:transaction, :s2]]
+    sqls = @db.sqls
+    ["UPDATE albums SET artist_id = 2, name = 'RF' WHERE (id = 1) -- s2", "UPDATE albums SET name = 'RF', artist_id = 2 WHERE (id = 1) -- s2"].should include(sqls.slice!(2))
+    sqls.should == ["SELECT * FROM albums LIMIT 1 -- s2", "BEGIN -- s2", "COMMIT -- s2"]
   end 
 
-  specify "should use not override shard given when saving" do
+  specify "should use override current shard when saving with given :server option" do
     @Album.use_transactions = true
     @Album.server(:s2).first.save(:server=>:s1)
-    @Album.actions.should == [[:fetch, nil, :s2], [:update, {:name=>"RF", :artist_id=>2}, "(id = 1)", :s2]]
-    @db.actions.should == [[:transaction, :s1]]
+    sqls = @db.sqls
+    ["UPDATE albums SET artist_id = 2, name = 'RF' WHERE (id = 1) -- s1", "UPDATE albums SET name = 'RF', artist_id = 2 WHERE (id = 1) -- s1"].should include(sqls.slice!(2))
+    sqls.should == ["SELECT * FROM albums LIMIT 1 -- s2", "BEGIN -- s1", "COMMIT -- s1"]
   end 
 end

data/spec/extensions/single_table_inheritance_spec.rb

@@ -3,9 +3,7 @@ require File.join(File.dirname(File.expand_path(__FILE__)), "spec_helper")
 describe Sequel::Model, "#sti_key" do
   before do
     class ::StiTest < Sequel::Model
-      def kind; self[:kind]; end 
-      def kind=(x); self[:kind] = x; end 
-      def _refresh(x); end 
+      columns :id, :kind, :blah
       plugin :single_table_inheritance, :kind
     end 
     class ::StiTestSub1 < StiTest
@@ -22,7 +20,7 @@ describe Sequel::Model, "#sti_key" do
   end
 
   specify "should have simple_table = nil" do
-    StiTest.simple_table.should == nil
+    StiTest.simple_table.should == "sti_tests"
     StiTestSub1.simple_table.should == nil
   end
   
@@ -30,15 +28,9 @@ describe Sequel::Model, "#sti_key" do
     StiTest.plugin :single_table_inheritance, :blah
     Object.send(:remove_const, :StiTestSub1)
     Object.send(:remove_const, :StiTestSub2)
-    class ::StiTestSub1 < StiTest
-    end 
-    class ::StiTestSub2 < StiTest
-    end 
-    def @ds.fetch_rows(sql)
-      yield({:blah=>'StiTest'})
-      yield({:blah=>'StiTestSub1'})
-      yield({:blah=>'StiTestSub2'})
-    end 
+    class ::StiTestSub1 < StiTest; end 
+    class ::StiTestSub2 < StiTest; end 
+    @ds._fetch = [{:blah=>'StiTest'}, {:blah=>'StiTestSub1'}, {:blah=>'StiTestSub2'}]
     StiTest.all.collect{|x| x.class}.should == [StiTest, StiTestSub1, StiTestSub2]
     StiTest.dataset.sql.should == "SELECT * FROM sti_tests"
     StiTestSub1.dataset.sql.should == "SELECT * FROM sti_tests WHERE (sti_tests.blah IN ('StiTestSub1'))"
@@ -46,18 +38,23 @@ describe Sequel::Model, "#sti_key" do
   end 
   
   it "should return rows with the correct class based on the polymorphic_key value" do
-    def @ds.fetch_rows(sql)
-      yield({:kind=>'StiTest'})
-      yield({:kind=>'StiTestSub1'})
-      yield({:kind=>'StiTestSub2'})
-    end 
+    @ds._fetch = [{:kind=>'StiTest'}, {:kind=>'StiTestSub1'}, {:kind=>'StiTestSub2'}]
     StiTest.all.collect{|x| x.class}.should == [StiTest, StiTestSub1, StiTestSub2]
   end 
 
+  it "should return rows with the correct class based on the polymorphic_key value when retreiving by primary key" do
+    @ds._fetch = [{:kind=>'StiTestSub1'}]
+    StiTest[1].class.should == StiTestSub1
+  end 
+
+  it "should return rows with the correct class for subclasses based on the polymorphic_key value" do
+    class ::StiTestSub1Sub < StiTestSub1; end 
+    StiTestSub1.dataset._fetch = [{:kind=>'StiTestSub1'}, {:kind=>'StiTestSub1Sub'}]
+    StiTestSub1.all.collect{|x| x.class}.should == [StiTestSub1, StiTestSub1Sub]
+  end 
+
   it "should fallback to the main class if the given class does not exist" do
-    def @ds.fetch_rows(sql)
-      yield({:kind=>'StiTestSub3'})
-    end
+    @ds._fetch = {:kind=>'StiTestSub3'}
     StiTest.all.collect{|x| x.class}.should == [StiTest]
   end
 
@@ -68,10 +65,7 @@ describe Sequel::Model, "#sti_key" do
       Object
     end
     StiTest.plugin :single_table_inheritance, :kind
-    def @ds.fetch_rows(sql)
-      yield({:kind=>''})
-      yield({:kind=>nil})
-    end
+    @ds._fetch = [{:kind=>''}, {:kind=>nil}]
     StiTest.all.collect{|x| x.class}.should == [StiTest, StiTest]
     called.should == false
   end
@@ -80,19 +74,19 @@ describe Sequel::Model, "#sti_key" do
     StiTest.new.save
     StiTestSub1.new.save
     StiTestSub2.new.save
-    MODEL_DB.sqls.should == ["INSERT INTO sti_tests (kind) VALUES ('StiTest')", "INSERT INTO sti_tests (kind) VALUES ('StiTestSub1')", "INSERT INTO sti_tests (kind) VALUES ('StiTestSub2')"]
+    MODEL_DB.sqls.should == ["INSERT INTO sti_tests (kind) VALUES ('StiTest')", "SELECT * FROM sti_tests WHERE (id = 10) LIMIT 1", "INSERT INTO sti_tests (kind) VALUES ('StiTestSub1')", "SELECT * FROM sti_tests WHERE ((sti_tests.kind IN ('StiTestSub1')) AND (id = 10)) LIMIT 1", "INSERT INTO sti_tests (kind) VALUES ('StiTestSub2')", "SELECT * FROM sti_tests WHERE ((sti_tests.kind IN ('StiTestSub2')) AND (id = 10)) LIMIT 1"]
   end
 
   it "should have the before_create hook not override an existing value" do
     StiTest.create(:kind=>'StiTestSub1')
-    MODEL_DB.sqls.should == ["INSERT INTO sti_tests (kind) VALUES ('StiTestSub1')"]
+    MODEL_DB.sqls.should == ["INSERT INTO sti_tests (kind) VALUES ('StiTestSub1')", "SELECT * FROM sti_tests WHERE (id = 10) LIMIT 1"]
   end
 
   it "should have the before_create hook handle columns with the same name as existing method names" do
     StiTest.plugin :single_table_inheritance, :type
     StiTest.columns :id, :type
     StiTest.create
-    MODEL_DB.sqls.should == ["INSERT INTO sti_tests (type) VALUES ('StiTest')"]
+    MODEL_DB.sqls.should == ["INSERT INTO sti_tests (type) VALUES ('StiTest')", "SELECT * FROM sti_tests WHERE (id = 10) LIMIT 1"]
   end
 
   it "should add a filter to model datasets inside subclasses hook to only retreive objects with the matching key" do
@@ -118,7 +112,7 @@ describe Sequel::Model, "#sti_key" do
     before do
       class ::StiTest2 < Sequel::Model
         columns :id, :kind
-        def _refresh(x); end
+        def _save_refresh; end
       end
     end
     after do
@@ -187,5 +181,47 @@ describe Sequel::Model, "#sti_key" do
       proc{StiTest2.dataset.row_proc.call(:kind=>1)}.should raise_error(Sequel::Error)
       proc{StiTest2.dataset.row_proc.call(:kind=>2)}.should raise_error(Sequel::Error)
     end
+
+    it "should work with non-bijective mappings" do
+      StiTest2.plugin :single_table_inheritance, :kind, :model_map=>{0=>'StiTest3', 1=>'StiTest3', 2=>'StiTest4'}
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+      StiTest2.dataset.row_proc.call(:kind=>0).should be_a_instance_of(StiTest3)
+      StiTest2.dataset.row_proc.call(:kind=>1).should be_a_instance_of(StiTest3)
+      StiTest2.dataset.row_proc.call(:kind=>2).should be_a_instance_of(StiTest4)
+
+      StiTest3.create.kind.should == 1
+      StiTest4.create.kind.should == 2
+    end
+
+    it "should work with non-bijective mappings and key map procs" do
+      StiTest2.plugin :single_table_inheritance, :kind,
+        :key_map=>proc{|model| model.to_s == 'StiTest4' ? 2 : [0,1] }
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+
+      StiTest2.dataset.sql.should == "SELECT * FROM sti_test2s"
+      StiTest3.dataset.sql.should == "SELECT * FROM sti_test2s WHERE (sti_test2s.kind IN (0, 1))"
+      StiTest4.dataset.sql.should == "SELECT * FROM sti_test2s WHERE (sti_test2s.kind IN (2))"
+    end
+
+    it "should create correct sql with non-bijective mappings" do
+      StiTest2.plugin :single_table_inheritance, :kind, :model_map=>{0=>'StiTest3', 1=>'StiTest3', 2=>'StiTest4'}
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+    
+      StiTest2.dataset.sql.should == "SELECT * FROM sti_test2s"
+      StiTest3.dataset.sql.should == "SELECT * FROM sti_test2s WHERE (sti_test2s.kind IN (0, 1))"
+    end
+
+    it "should honor a :key_chooser" do
+      StiTest2.plugin :single_table_inheritance, :kind, :key_chooser => proc{|inst| inst.model.to_s.downcase }
+      class ::StiTest3 < ::StiTest2; end
+      class ::StiTest4 < ::StiTest2; end
+
+      StiTest3.create.kind.should == 'stitest3'
+      StiTest4.create.kind.should == 'stitest4'
+    end
+
   end
 end

data/spec/extensions/skip_create_refresh_spec.rb

@@ -5,7 +5,7 @@ describe "Sequel::Plugins::SkipCreateRefresh" do
     c = Class.new(Sequel::Model(:a))
     c.columns :id, :x
     c.db.reset
-    c.dataset.meta_def(:insert){|*a| super(*a); 2}
+    c.instance_dataset.meta_def(:insert){|*a| super(*a); 2}
     c.create(:x=>1)
     c.db.sqls.should == ['INSERT INTO a (x) VALUES (1)', 'SELECT * FROM a WHERE (id = 2) LIMIT 1']
 

data/spec/extensions/spec_helper.rb

@@ -1,80 +1,45 @@
 require 'rubygems'
-unless Object.const_defined?('Sequel')
-  $:.unshift(File.join(File.dirname(File.expand_path(__FILE__)), "../../lib/"))
-  require 'sequel/core'
-end
-unless Sequel.const_defined?('Model')
-  $:.unshift(File.join(File.dirname(File.expand_path(__FILE__)), "../../lib/"))
-  require 'sequel/model'
-end
-
-Sequel.extension(*%w'string_date_time inflector pagination query pretty_table blank migration schema_dumper looser_typecasting sql_expr thread_local_timezones to_dot')
-{:hook_class_methods=>[], :schema=>[], :validation_class_methods=>[]}.each{|p, opts| Sequel::Model.plugin(p, *opts)}
-
-def skip_warn(s)
-  warn "Skipping test of #{s}" if ENV["SKIPPED_TEST_WARN"]
-end
 
-class MockDataset < Sequel::Dataset
-  def insert(*args)
-    @db.execute insert_sql(*args)
-  end
-  
-  def update(*args)
-    @db.execute update_sql(*args)
-    1
-  end
-  
-  def delete(*args)
-    @db.execute delete_sql(*args)
-    1
-  end
-  
-  def fetch_rows(sql)
-    return if sql =~ /information_schema/
-    @db.execute(sql)
-    yield({:id => 1, :x => 1})
+if defined?(RSpec)
+  begin
+    require 'rspec/expectations'
+  rescue LoadError
+    nil
   end
+end
 
-  def quoted_identifier(c)
-    "\"#{c}\""
-  end
+unless Object.const_defined?('Sequel') && Sequel.const_defined?('Model')
+  $:.unshift(File.join(File.dirname(File.expand_path(__FILE__)), "../../lib/"))
+  require 'sequel/no_core_ext'
 end
 
-class MockDatabase < Sequel::Database
-  @@quote_identifiers = false
-  self.identifier_input_method = nil
-  self.identifier_output_method = nil
-  attr_reader :sqls
-  
-  def execute(sql, opts={})
-    @sqls ||= []
-    @sqls << sql
-  end
-  
-  def new_sqls
-    s = sqls
-    reset
-    s
-  end
+begin
+  # Attempt to load ActiveSupport inflector first, so Sequel inflector
+  # can override it.
+  require 'active_support/inflector'
+  require 'active_support/core_ext/string/inflections'
+rescue LoadError
+  nil
+end
 
-  def reset
-    @sqls = []
-  end
+# Load most extensions by default, so that any conflicts are easily detectable.
+Sequel.extension(*%w'string_date_time inflector pagination query pretty_table blank migration schema_dumper looser_typecasting sql_expr thread_local_timezones to_dot columns_introspection server_block arbitrary_servers pg_auto_parameterize pg_statement_cache pg_array pg_array_ops pg_hstore pg_hstore_ops pg_range pg_range_ops pg_json pg_inet pg_row pg_row_ops schema_caching null_dataset select_remove query_literals eval_inspect')
 
-  def schema(table_name, opts)
-    [[:id, {:primary_key=>true}]]
-  end
+Sequel::Dataset.introspect_all_columns if ENV['SEQUEL_COLUMNS_INTROSPECTION']
 
-  def transaction(opts={}); yield; end
-  
-  def dataset(opts=nil); MockDataset.new(self, opts); end
+def skip_warn(s)
+  warn "Skipping test of #{s}" if ENV["SKIPPED_TEST_WARN"]
 end
 
+Sequel.quote_identifiers = false
+Sequel.identifier_input_method = nil
+Sequel.identifier_output_method = nil
+
 class << Sequel::Model
+  attr_writer :db_schema
   alias orig_columns columns
   def columns(*cols)
-    return @columns if cols.empty?
+    return super if cols.empty?
     define_method(:columns){cols}
     @dataset.instance_variable_set(:@columns, cols) if @dataset
     def_column_accessor(*cols)
@@ -82,10 +47,12 @@ class << Sequel::Model
     @db_schema = {}
     cols.each{|c| @db_schema[c] = {}}
   end
-  def simple_table
-    nil
-  end
 end
 
-Sequel::Model.db = MODEL_DB = MockDatabase.new
 Sequel::Model.use_transactions = false
+Sequel::Model.cache_anonymous_models = false
+
+db = Sequel.mock(:fetch=>{:id => 1, :x => 1}, :numrows=>1, :autoid=>proc{|sql| 10})
+def db.schema(*) [[:id, {:primary_key=>true}]] end
+def db.reset() sqls end
+Sequel::Model.db = MODEL_DB = db