From cc7bbe812a2e50a5b2293be13301ed3cd7318bd9 Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@jeremyevans.net>
Date: Fri, 13 Jun 2008 13:49:26 -0700
Subject: [PATCH] Add Model.sti_key, for easily setting up single table
 inheritance

Model.sti_key is a simple helper method.  It doesn't add any new
functionality, it just calls some existing methods to make it easy
to set up STI quickly and correctly.  It is only useful if you have
a single string column in your table that contains the name of the
class you would like to use.  Model.sti_key should only be called in
the parent class, not in the subclasses.

The two things it does are:

1) Call dataset.set_model with a hash that results in a valid ruby
class for any key.  If there are problems using the string (like an
empty string or NULL), it uses the class that called sti_key.

2) Add a before_create hook that sets the name of the sti_key column
to the class name, so that when it is loaded from the database, you
get the same class back.

Example:

  class Artist < Sequel::Model
    sti_key :kind
  end
---
 sequel/CHANGELOG                |  2 ++
 sequel/lib/sequel_model/base.rb | 13 +++++++++
 sequel/spec/model_spec.rb       | 47 ++++++++++++++++++++++++++++++---
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/sequel/CHANGELOG b/sequel/CHANGELOG
index 70295ce400..d770d30f8a 100644
--- a/sequel/CHANGELOG
+++ b/sequel/CHANGELOG
@@ -1,5 +1,7 @@
 === HEAD
 
+* Add Model.sti_key, for easily setting up single table inheritance (jeremyevans)
+
 * Make all associations support a :read_only option, which doesn't add methods that modify the database (jeremyevans)
 
 * Make *_to_many associations support a :limit option, for specifying a limit to the resulting records (and possibly an offset) (jeremyevans)
diff --git a/sequel/lib/sequel_model/base.rb b/sequel/lib/sequel_model/base.rb
index b4936cf7a0..fbb633db41 100644
--- a/sequel/lib/sequel_model/base.rb
+++ b/sequel/lib/sequel_model/base.rb
@@ -271,6 +271,19 @@ def self.set_primary_key(*key)
       @primary_key = (key.length == 1) ? key[0] : key.flatten
     end
 
+    # Makes this model a polymorphic model with the given key being a string
+    # field in the database holding the name of the class to use.  If the
+    # key given has a NULL value or there are any problems looking up the
+    # class, uses the current class.
+    #
+    # This should be used to set up single table inheritance for the model,
+    # and it only makes sense to use this in the parent class.
+    def self.sti_key(key)
+      m = self
+      dataset.set_model(key, Hash.new{|h,k| h[k] = (k.constantize rescue m)})
+      before_create(:set_sti_key){send("#{key}=", model.name)}
+    end
+
     # Returns the columns as a list of frozen strings instead
     # of a list of symbols.  This makes it possible to check
     # whether a column exists without creating a symbol, which
diff --git a/sequel/spec/model_spec.rb b/sequel/spec/model_spec.rb
index 89b251dc23..bf17fff85b 100644
--- a/sequel/spec/model_spec.rb
+++ b/sequel/spec/model_spec.rb
@@ -25,7 +25,6 @@
 end
 
 describe Sequel::Model, "dataset & schema" do
-
   before do
     @model = Class.new(Sequel::Model(:items))
   end
@@ -80,9 +79,51 @@
     @model.primary_key.should == :id
     @model.table_name.should == :foo
   end
+end
+
+describe Sequel::Model, "#sti_key" do
+  before do
+    class StiTest < Sequel::Model
+      def kind=(x); self[:kind] = x; end
+      def refresh; end
+      sti_key :kind
+    end
+    class StiTestSub1 < StiTest
+    end
+    class StiTestSub2 < StiTest
+    end
+    @ds = StiTest.dataset
+    MODEL_DB.reset
+  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
+    StiTest.all.collect{|x| x.class}.should == [StiTest, StiTestSub1, StiTestSub2]
+  end
 
-  it "puts the lotion in the basket or it gets the hose again" do
-    # just kidding!
+  it "should fallback to the main class if polymophic_key value is NULL" do
+    def @ds.fetch_rows(sql)
+      yield({:kind=>nil})
+    end
+    StiTest.all.collect{|x| x.class}.should == [StiTest]
+  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
+    StiTest.all.collect{|x| x.class}.should == [StiTest]
+  end
+  
+  it "should add a before_create hook that sets the model class name for the 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')"]
   end
 end